""" Token 预算管理器 — 追踪每次 LLM 调用的 token 消耗，提供预警和限额控制。 参考 Claude Code: - src/utils/tokenBudget.ts — 预算追踪与自动续行 - src/utils/tokenUsageTracker.ts — 累计用量追踪 - UI StatusLine 的 token 用量条 核心概念: - context_window: 模型上下文窗口大小（如 128K） - output_reserve: 留给模型输出的空间（默认 8K），只有 (window - reserve) 可被输入使用 - warning/critical/exhausted 三级预警 - 用户可设置 target budget（如 +500k），达到后自动继续 """ from __future__ import annotations import logging from dataclasses import dataclass, field from typing import Any, Dict, Optional from app.core.token_counter import TokenCounter, get_model_context_window logger = logging.getLogger(__name__) # ────────────── 配置 ────────────── @dataclass class TokenBudgetConfig: """Token 预算配置。""" # ── 总开关 ── enabled: bool = True # ── 窗口配置 ── context_window: int = 128_000 # 模型上下文窗口（token），0=自动检测 output_reserve: int = 8_192 # 留给模型输出的空间 # ── 预警阈值（占有效窗口的百分比） ── warning_threshold_pct: float = 0.75 # 75% → 开始预警 compact_threshold_pct: float = 0.85 # 85% → 触发自动压缩 hard_limit_pct: float = 0.95 # 95% → 下次调用前必须压缩 # ── 用户预算目标 ── user_budget: Optional[int] = None # 用户累计 token 目标（如 500_000） auto_continue: bool = False # 达到用户预算后是否自动继续 # ── 压缩协调 ── compaction_after_warning: bool = True # 预警后是否自动触发压缩 max_compaction_attempts: int = 3 # 单轮最多压缩尝试次数 @property def effective_window(self) -> int: """有效输入窗口 = 上下文窗口 - 输出预留。""" return max(0, self.context_window - self.output_reserve) @property def warning_at(self) -> int: """预警 token 数。""" return int(self.effective_window * self.warning_threshold_pct) @property def compact_at(self) -> int: """自动压缩触发 token 数。""" return int(self.effective_window * self.compact_threshold_pct) @property def hard_limit_at(self) -> int: """硬限制 token 数（超过则拒绝调用 LLM）。""" return int(self.effective_window * self.hard_limit_pct) # ────────────── 快照 ────────────── @dataclass class TokenSnapshot: """单次 LLM 调用的 token 快照。""" prompt_tokens: int = 0 completion_tokens: int = 0 total_tokens: int = 0 iteration: int = 0 step_type: str = "" # think / final model: str = "" # ────────────── 预算追踪器 ────────────── class TokenBudget: """会话级 token 预算追踪器。 追踪: - 当前消息列表的 token 数（输入侧） - 累计 LLM 消耗（输入 + 输出） - 用户预算目标的进度 - 预警/压缩/限额状态 用法:: budget = TokenBudget(TokenBudgetConfig(context_window=128000)) budget.update_input_estimate(counter.count_messages(messages)) if budget.needs_compaction: # trigger compaction budget.record_llm_call(prompt_tokens=5000, completion_tokens=800) print(budget.status_line) # "12.5k/128k (10%) | ⚠ near limit" """ def __init__( self, config: Optional[TokenBudgetConfig] = None, model: str = "deepseek-v4-flash", token_counter: Optional[TokenCounter] = None, ): self.config = config or TokenBudgetConfig() self.model = model self.counter = token_counter or TokenCounter(model=model) # 自动检测上下文窗口 if self.config.context_window <= 0: self.config.context_window = get_model_context_window(model) # ── 计数器 ── self._input_tokens_estimate: int = 0 # 当前输入消息列表的 token 估计 self._cumulative_prompt_tokens: int = 0 # 累计 prompt token（含重试） self._cumulative_completion_tokens: int = 0 # 累计 completion token self._llm_call_count: int = 0 self._compaction_attempts_this_turn: int = 0 # ── 历史快照（最近 20 次调用） ── self._snapshots: list[TokenSnapshot] = [] # ──────── 属性 ──────── @property def input_tokens(self) -> int: """当前输入消息列表的预估 token 数。""" return self._input_tokens_estimate @property def cumulative_total(self) -> int: """累计消耗 token（prompt + completion）。""" return self._cumulative_prompt_tokens + self._cumulative_completion_tokens @property def cumulative_prompt(self) -> int: return self._cumulative_prompt_tokens @property def cumulative_completion(self) -> int: return self._cumulative_completion_tokens @property def llm_call_count(self) -> int: return self._llm_call_count @property def input_usage_pct(self) -> float: """输入占用窗口的百分比。""" ew = self.config.effective_window return self._input_tokens_estimate / ew if ew > 0 else 0.0 @property def input_remaining(self) -> int: """输入侧剩余 token 空间。""" return max(0, self.config.effective_window - self._input_tokens_estimate) @property def user_budget_used(self) -> int: """用户预算消耗量。""" return self.cumulative_total @property def user_budget_remaining(self) -> Optional[int]: """用户预算剩余量（未设置则 None）。""" if self.config.user_budget is None: return None return max(0, self.config.user_budget - self.cumulative_total) @property def user_budget_pct(self) -> Optional[float]: """用户预算消耗百分比。""" if self.config.user_budget is None or self.config.user_budget <= 0: return None return self.cumulative_total / self.config.user_budget # ──────── 状态判断 ──────── @property def is_warning(self) -> bool: """是否达到预警线。""" return self._input_tokens_estimate >= self.config.warning_at @property def is_critical(self) -> bool: """是否达到紧急线（需要立即压缩）。""" return self._input_tokens_estimate >= self.config.compact_at @property def is_exhausted(self) -> bool: """是否达到硬限制（调用 LLM 前必须处理）。""" return self._input_tokens_estimate >= self.config.hard_limit_at @property def needs_compaction(self) -> bool: """是否需要触发压缩。""" if not self.config.compaction_after_warning: return False if self._compaction_attempts_this_turn >= self.config.max_compaction_attempts: return False # 熔断 return self.is_critical @property def compaction_attempts(self) -> int: return self._compaction_attempts_this_turn @property def is_user_budget_exhausted(self) -> bool: """用户预算是否用尽。""" rem = self.user_budget_remaining return rem is not None and rem <= 0 # ──────── 更新方法 ──────── def update_input_estimate(self, tokens: int) -> None: """更新当前输入消息列表的 token 估计值（每次消息列表变更后调用）。""" self._input_tokens_estimate = tokens logger.debug( "TokenBudget: input=%d tokens (%.1f%% of %d, compact_at=%d)", tokens, self.input_usage_pct * 100, self.config.effective_window, self.config.compact_at, ) def update_from_counter(self, messages: list) -> int: """从消息列表计算并更新输入 token 估计。返回估计值。""" tokens = self.counter.count_messages(messages) self.update_input_estimate(tokens) return tokens def record_llm_call( self, prompt_tokens: int = 0, completion_tokens: int = 0, iteration: int = 0, step_type: str = "think", ) -> TokenSnapshot: """记录一次 LLM 调用。 注意：prompt_tokens 应优先使用 API 返回的实际值； 若不可用则传入 0，由 update_input_estimate 的估算值代替。 """ if prompt_tokens <= 0: prompt_tokens = self._input_tokens_estimate self._cumulative_prompt_tokens += prompt_tokens self._cumulative_completion_tokens += completion_tokens self._llm_call_count += 1 snap = TokenSnapshot( prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=prompt_tokens + completion_tokens, iteration=iteration, step_type=step_type, model=self.model, ) self._snapshots.append(snap) # 只保留最近 50 次快照 if len(self._snapshots) > 50: self._snapshots = self._snapshots[-50:] if logger.isEnabledFor(logging.DEBUG): logger.debug( "TokenBudget: call #%d prompt=%d comp=%d total=%d cumulative=%d (%.1f%%)", self._llm_call_count, prompt_tokens, completion_tokens, prompt_tokens + completion_tokens, self.cumulative_total, self.input_usage_pct * 100, ) return snap def record_compaction_attempt(self) -> None: """记录一次压缩尝试（用于熔断计数）。""" self._compaction_attempts_this_turn += 1 def reset_compaction_attempts(self) -> None: """重置压缩尝试计数（新轮次开始时调用）。""" self._compaction_attempts_this_turn = 0 # ──────── 摘要/展示 ──────── @property def status_line(self) -> str: """单行状态摘要（用于日志/UI）。""" pct = self.input_usage_pct * 100 parts = [f"{self._input_tokens_estimate/1000:.1f}k/{self.config.effective_window/1000:.0f}k ({pct:.0f}%)"] if self.is_exhausted: parts.append("[EXHAUSTED]") elif self.is_critical: parts.append("[CRITICAL]") elif self.is_warning: parts.append("[WARNING]") if self.config.user_budget: parts.append(f"| budget: {self.cumulative_total/1000:.1f}k/{self.config.user_budget/1000:.0f}k") return " ".join(parts) def summary(self) -> Dict[str, Any]: """返回可供 API 响应的 token 预算摘要。""" result: Dict[str, Any] = { "input_tokens": self._input_tokens_estimate, "input_remaining": self.input_remaining, "input_usage_pct": round(self.input_usage_pct, 4), "effective_window": self.config.effective_window, "context_window": self.config.context_window, "cumulative_total": self.cumulative_total, "cumulative_prompt": self._cumulative_prompt_tokens, "cumulative_completion": self._cumulative_completion_tokens, "llm_call_count": self._llm_call_count, "is_warning": self.is_warning, "is_critical": self.is_critical, "is_exhausted": self.is_exhausted, "compaction_attempts": self._compaction_attempts_this_turn, } if self.config.user_budget is not None: result["user_budget"] = self.config.user_budget result["user_budget_used"] = self.user_budget_used result["user_budget_remaining"] = self.user_budget_remaining result["user_budget_pct"] = round(self.user_budget_pct, 4) if self.user_budget_pct else None return result def needs_user_budget_continue(self) -> bool: """用户预算用尽且配置了自动继续。""" return self.is_user_budget_exhausted and self.config.auto_continue # ────────────── 便捷工厂 ────────────── def create_token_budget( model: str = "deepseek-v4-flash", context_window: int = 0, user_budget: Optional[int] = None, enabled: bool = True, ) -> TokenBudget: """创建预配置的 TokenBudget（适合大多数场景）。""" config = TokenBudgetConfig( enabled=enabled, context_window=context_window or get_model_context_window(model), user_budget=user_budget, ) return TokenBudget(config=config, model=model)