35 lines
1.1 KiB
Python
35 lines
1.1 KiB
Python
"""Token 预算守卫：累计消耗超阈值时优雅收尾。"""
|
|
|
|
from app.application.agent.extensions import (
|
|
AgentRunContext,
|
|
ExtensionRegistry,
|
|
TurnSnapshot,
|
|
)
|
|
from app.application.agent.python_skills import PythonSkillDefinition
|
|
|
|
# 累计 total_tokens 软上限。达到后 TurnPolicy 要求停止，让 Agent 用当前内容自然
|
|
# 收尾——区别于 max_iterations 的硬截断（truncated=True，回复可能只说了一半）。
|
|
_BUDGET_TOTAL_TOKENS = 20000
|
|
|
|
|
|
async def _stop_when_budget_exhausted(
|
|
turn: TurnSnapshot, _run: AgentRunContext
|
|
) -> bool:
|
|
return turn.usage.total_tokens >= _BUDGET_TOTAL_TOKENS
|
|
|
|
|
|
def _setup(registry: ExtensionRegistry) -> None:
|
|
registry.add_turn_policy(_stop_when_budget_exhausted)
|
|
|
|
|
|
SKILL = PythonSkillDefinition(
|
|
name="budget_guard",
|
|
description="累计 token 消耗超预算时优雅停止 Agent 循环，避免成本失控。",
|
|
triggers=("预算", "budget", "长任务"),
|
|
instructions=(
|
|
"本轮受 token 预算约束。若已接近上限，请用当前已有信息给出最终回复，"
|
|
"不要发起新的工具调用或长篇推理。"
|
|
),
|
|
setup=_setup,
|
|
)
|