Neo
6f8f12314f
包含多个会话的累积代码变更： - backend: AI 聊天服务、触发器调度、认证增强、WebSocket、调度器最小间隔 - admin-web: ETL 状态页、任务管理、调度配置、登录优化 - miniprogram: 看板页面、聊天集成、UI 组件、导航更新 - etl: DWS 新任务（finance_area_daily/board_cache）、连接器增强 - tenant-admin: 项目初始化 - db: 19 个迁移脚本（etl_feiqiu 11 + zqyy_app 8） - packages/shared: 枚举和工具函数更新 - tools: 数据库工具、报表生成、健康检查 - docs: PRD/架构/部署/合约文档更新 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
102 lines
3.1 KiB
Python
102 lines
3.1 KiB
Python
"""Token 预算追踪器 — 从 ai_run_logs 聚合日/月 token 消耗。
|
||
|
||
每次 AI 调用前检查预算，超限时拒绝请求。
|
||
日预算默认 100,000 tokens，月预算默认 2,000,000 tokens。
|
||
|
||
聚合数据通过构造函数注入的 callable 获取（解耦 AIRunLogService），
|
||
callable 签名：() -> int，分别返回当日/当月已消耗 token 数。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from typing import Callable, Protocol
|
||
|
||
|
||
class UsageProvider(Protocol):
|
||
"""Token 用量数据提供者协议。"""
|
||
|
||
def get_daily_usage(self) -> int:
|
||
"""返回当日已消耗 token 数。"""
|
||
...
|
||
|
||
def get_monthly_usage(self) -> int:
|
||
"""返回当月已消耗 token 数。"""
|
||
...
|
||
|
||
|
||
@dataclass
|
||
class BudgetStatus:
|
||
"""预算检查结果。"""
|
||
|
||
allowed: bool
|
||
daily_used: int
|
||
monthly_used: int
|
||
reason: str | None = None # "daily_exceeded" / "monthly_exceeded" / None
|
||
|
||
|
||
class BudgetTracker:
|
||
"""Token 预算追踪器，从 ai_run_logs 聚合。
|
||
|
||
支持两种注入方式：
|
||
1. 传入 UsageProvider 实例（如 AIRunLogService）
|
||
2. 传入两个 callable：get_daily_usage / get_monthly_usage
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
daily_limit: int = 100_000,
|
||
monthly_limit: int = 2_000_000,
|
||
*,
|
||
get_daily_usage: Callable[[], int] | None = None,
|
||
get_monthly_usage: Callable[[], int] | None = None,
|
||
usage_provider: UsageProvider | None = None,
|
||
) -> None:
|
||
self.daily_limit = daily_limit
|
||
self.monthly_limit = monthly_limit
|
||
|
||
# 优先使用 usage_provider，其次使用独立 callable
|
||
if usage_provider is not None:
|
||
self._get_daily_usage = usage_provider.get_daily_usage
|
||
self._get_monthly_usage = usage_provider.get_monthly_usage
|
||
elif get_daily_usage is not None and get_monthly_usage is not None:
|
||
self._get_daily_usage = get_daily_usage
|
||
self._get_monthly_usage = get_monthly_usage
|
||
else:
|
||
raise ValueError(
|
||
"必须提供 usage_provider 或同时提供 "
|
||
"get_daily_usage 和 get_monthly_usage callable"
|
||
)
|
||
|
||
def check_budget(self) -> BudgetStatus:
|
||
"""检查当前预算状态。
|
||
|
||
先检查日预算，再检查月预算。
|
||
任一超限即返回 allowed=False 并附带原因。
|
||
"""
|
||
daily_used = self._get_daily_usage()
|
||
monthly_used = self._get_monthly_usage()
|
||
|
||
if daily_used >= self.daily_limit:
|
||
return BudgetStatus(
|
||
allowed=False,
|
||
daily_used=daily_used,
|
||
monthly_used=monthly_used,
|
||
reason="daily_exceeded",
|
||
)
|
||
|
||
if monthly_used >= self.monthly_limit:
|
||
return BudgetStatus(
|
||
allowed=False,
|
||
daily_used=daily_used,
|
||
monthly_used=monthly_used,
|
||
reason="monthly_exceeded",
|
||
)
|
||
|
||
return BudgetStatus(
|
||
allowed=True,
|
||
daily_used=daily_used,
|
||
monthly_used=monthly_used,
|
||
reason=None,
|
||
)
|