mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-07-21 21:48:04 +00:00
feat(middleware): add TokenBudgetMiddleware for per-run token budget e… (#3412)
* eat(middleware): add TokenBudgetMiddleware for per-run token budget enforcement * address copilot comments * resolve feedback --------- Co-authored-by: Willem Jiang <willem.jiang@gmail.com>
This commit is contained in:
Chetan Sharma
GitHub
parent
c177d0e542
commit
78fff5a5e2
@ -282,7 +282,17 @@ def _assemble_from_features(
|
||||
|
||||
chain.append(LoopDetectionMiddleware.from_config(LoopDetectionConfig()))
|
||||
|
||||
# --- [13] Clarification (always last among built-ins) ---
|
||||
# --- [13] TokenBudget ---
|
||||
if feat.token_budget is not False:
|
||||
if isinstance(feat.token_budget, AgentMiddleware):
|
||||
chain.append(feat.token_budget)
|
||||
else:
|
||||
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
|
||||
from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
|
||||
chain.append(TokenBudgetMiddleware.from_config(TokenBudgetConfig()))
|
||||
|
||||
# --- [14] Clarification (always last among built-ins) ---
|
||||
chain.append(ClarificationMiddleware())
|
||||
extra_tools.append(ask_clarification_tool)
|
||||
|
||||
|
||||
@ -32,6 +32,7 @@ class RuntimeFeatures:
|
||||
auto_title: bool | AgentMiddleware = False
|
||||
guardrail: Literal[False] | AgentMiddleware = False
|
||||
loop_detection: bool | AgentMiddleware = True
|
||||
token_budget: bool | AgentMiddleware = False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@ -359,6 +359,13 @@ def build_middlewares(
|
||||
if loop_detection_config.enabled:
|
||||
middlewares.append(LoopDetectionMiddleware.from_config(loop_detection_config))
|
||||
|
||||
# TokenBudgetMiddleware - enforce per-run token limits
|
||||
token_budget_config = resolved_app_config.token_budget
|
||||
if token_budget_config.enabled:
|
||||
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
|
||||
|
||||
middlewares.append(TokenBudgetMiddleware.from_config(token_budget_config))
|
||||
|
||||
# Inject custom middlewares before ClarificationMiddleware
|
||||
if custom_middlewares:
|
||||
middlewares.extend(custom_middlewares)
|
||||
|
||||
@ -0,0 +1,286 @@
|
||||
"""Middleware to enforce per-run token budget limits.
|
||||
Tracks cumulative token usage (input, output, total) across model calls within
|
||||
a single agent run and enforces configurable soft-warning and hard-stop
|
||||
thresholds.
|
||||
Detection strategy:
|
||||
1. After each model response, sum the `usage_metadata` of all `AIMessage`s
|
||||
in the current thread history. This automatically captures tokens from
|
||||
subagents because `TokenUsageMiddleware` retroactively adds them to the
|
||||
history.
|
||||
2. If the highest fraction (input, output, or total) >= warn_threshold,
|
||||
queue a warning.
|
||||
3. If the highest fraction >= hard_stop_threshold, strip tool_calls.
|
||||
Warning injection uses the deferred pattern:
|
||||
- after_model queues the warning (does NOT mutate state).
|
||||
- wrap_model_call injects it as a HumanMessage at the next model call.
|
||||
This preserves AIMessage(tool_calls) → ToolMessage pairing.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
from collections.abc import Awaitable, Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, override
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.agents.middleware import AgentMiddleware
|
||||
from langchain.agents.middleware.types import ModelCallResult, ModelRequest, ModelResponse
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langgraph.runtime import Runtime
|
||||
|
||||
from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_BUDGET_WARNING_MSG = (
|
||||
"[TOKEN BUDGET WARNING] You have used {used:,} of your {budget:,} {reason} token budget ({percent:.0f}%). Wrap up your current work and produce a final answer. Avoid starting new tool calls unless absolutely necessary."
|
||||
)
|
||||
_BUDGET_EXCEEDED_MSG = "[TOKEN BUDGET EXCEEDED] The {reason} token usage ({used:,}) has exceeded the safety limit ({budget:,}). Producing final answer with results collected so far."
|
||||
|
||||
|
||||
@dataclass
|
||||
class TokenUsage:
|
||||
input: int = 0
|
||||
output: int = 0
|
||||
total: int = 0
|
||||
|
||||
|
||||
class BoundedDict(OrderedDict):
|
||||
"""A bounded dictionary to prevent unbounded state growth on abandoned runs."""
|
||||
|
||||
def __init__(self, maxsize=1000, *args, **kwds):
|
||||
self.maxsize = maxsize
|
||||
super().__init__(*args, **kwds)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
if key not in self:
|
||||
if len(self) >= self.maxsize:
|
||||
self.popitem(last=False)
|
||||
super().__setitem__(key, value)
|
||||
|
||||
|
||||
class TokenBudgetMiddleware(AgentMiddleware[AgentState]):
|
||||
"""Enforce per-run token budget limits."""
|
||||
|
||||
def __init__(self, config: TokenBudgetConfig) -> None:
|
||||
super().__init__()
|
||||
self._config = config
|
||||
self._lock = threading.Lock()
|
||||
|
||||
# Keyed strictly by run_id (clobber-safe) and bounded (leak-safe)
|
||||
self._warned: BoundedDict[str, bool] = BoundedDict(1000)
|
||||
self._pending_warnings: BoundedDict[str, list[str]] = BoundedDict(1000)
|
||||
self._seen_messages: BoundedDict[str, dict[str, tuple[int, int]]] = BoundedDict(1000)
|
||||
self._cumulative_usage: BoundedDict[str, TokenUsage] = BoundedDict(1000)
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: TokenBudgetConfig) -> TokenBudgetMiddleware:
|
||||
return cls(config=config)
|
||||
|
||||
def reset(self) -> None:
|
||||
with self._lock:
|
||||
self._warned.clear()
|
||||
self._pending_warnings.clear()
|
||||
self._seen_messages.clear()
|
||||
self._cumulative_usage.clear()
|
||||
|
||||
@staticmethod
|
||||
def _get_run_id(runtime: Runtime) -> str:
|
||||
ctx = getattr(runtime, "context", None)
|
||||
if isinstance(ctx, dict) and "run_id" in ctx:
|
||||
return ctx["run_id"]
|
||||
# Fallback to runtime object ID to prevent collisions across embedded client runs
|
||||
return str(id(runtime))
|
||||
|
||||
def _clear_run_state(self, run_id: str) -> None:
|
||||
self._warned.pop(run_id, None)
|
||||
self._pending_warnings.pop(run_id, None)
|
||||
self._seen_messages.pop(run_id, None)
|
||||
self._cumulative_usage.pop(run_id, None)
|
||||
|
||||
@override
|
||||
def before_agent(self, state: AgentState, runtime: Runtime) -> None:
|
||||
if not self._config.enabled:
|
||||
return
|
||||
|
||||
# Mark all old messages from previous runss as 'seen' so they don't count toward THIS run's budget
|
||||
messages = state.get("messages", [])
|
||||
if not messages:
|
||||
return
|
||||
|
||||
run_id = self._get_run_id(runtime)
|
||||
seen = self._seen_messages.setdefault(run_id, {})
|
||||
self._cumulative_usage.setdefault(run_id, TokenUsage())
|
||||
|
||||
for msg in messages:
|
||||
if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
|
||||
usage = msg.usage_metadata or {}
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
seen[msg.id] = (input_tokens, output_tokens)
|
||||
|
||||
@override
|
||||
async def abefore_agent(self, state: AgentState, runtime: Runtime) -> None:
|
||||
self.before_agent(state, runtime)
|
||||
|
||||
@override
|
||||
def after_agent(self, state: AgentState, runtime: Runtime) -> None:
|
||||
if not self._config.enabled:
|
||||
return
|
||||
self._clear_run_state(self._get_run_id(runtime))
|
||||
|
||||
@override
|
||||
async def aafter_agent(self, state: AgentState, runtime: Runtime) -> None:
|
||||
self.after_agent(state, runtime)
|
||||
|
||||
@staticmethod
|
||||
def _append_text(content: str | list[dict | None] | None, stop_msg: str) -> str | list[dict | str]:
|
||||
"""Append a stop message to an AIMessage.content field."""
|
||||
if content is None:
|
||||
return stop_msg
|
||||
if isinstance(content, str):
|
||||
if content:
|
||||
return f"{content}\n\n{stop_msg}"
|
||||
return f"\n\n{stop_msg}"
|
||||
if isinstance(content, list):
|
||||
new_content = list(content)
|
||||
new_content.append({"type": "text", "text": f"\n\n{stop_msg}"})
|
||||
return new_content
|
||||
return f"{content}\n\n{stop_msg}"
|
||||
|
||||
def _build_hard_stop_update(self, msg: AIMessage, stop_msg: str) -> dict[str, Any]:
|
||||
"""Build the state update dictionary for a hard stop."""
|
||||
updated_content = self._append_text(msg.content, stop_msg)
|
||||
kwargs = dict(msg.additional_kwargs) if msg.additional_kwargs else {}
|
||||
if "tool_calls" in kwargs:
|
||||
del kwargs["tool_calls"]
|
||||
if "function_call" in kwargs:
|
||||
del kwargs["function_call"]
|
||||
|
||||
response_metadata = dict(getattr(msg, "response_metadata", {}) or {})
|
||||
|
||||
if response_metadata.get("finish_reason") == "tool_calls":
|
||||
response_metadata["finish_reason"] = "stop"
|
||||
|
||||
stopped_msg = msg.model_copy(update={"content": updated_content, "tool_calls": [], "additional_kwargs": kwargs, "response_metadata": response_metadata})
|
||||
return {"messages": [stopped_msg]}
|
||||
|
||||
def _apply(self, state: AgentState, runtime: Runtime) -> dict | None:
|
||||
if not self._config.enabled:
|
||||
return None
|
||||
|
||||
messages = state.get("messages", [])
|
||||
if not messages:
|
||||
return None
|
||||
|
||||
last_msg = messages[-1]
|
||||
if not isinstance(last_msg, AIMessage):
|
||||
return None
|
||||
|
||||
run_id = self._get_run_id(runtime)
|
||||
|
||||
seen = self._seen_messages.setdefault(run_id, {})
|
||||
usage_accum = self._cumulative_usage.setdefault(run_id, TokenUsage())
|
||||
|
||||
for msg in messages:
|
||||
if isinstance(msg, AIMessage) and msg.id and hasattr(msg, "usage_metadata"):
|
||||
usage = msg.usage_metadata or {}
|
||||
|
||||
input_tokens = usage.get("input_tokens", 0)
|
||||
output_tokens = usage.get("output_tokens", 0)
|
||||
|
||||
# Check what previously recorded for this exact message
|
||||
prev_input, prev_output = seen.get(msg.id, (0, 0))
|
||||
|
||||
# Calculate if any new tokens were added (handles retroactive subagent tokens)
|
||||
diff_input = max(0, input_tokens - prev_input)
|
||||
diff_output = max(0, output_tokens - prev_output)
|
||||
|
||||
if diff_input > 0 or diff_output > 0:
|
||||
usage_accum.input += diff_input
|
||||
usage_accum.output += diff_output
|
||||
usage_accum.total += diff_input + diff_output
|
||||
seen[msg.id] = (input_tokens, output_tokens)
|
||||
|
||||
if usage_accum.total <= 0:
|
||||
return None
|
||||
|
||||
fractions = [("total", usage_accum.total, self._config.max_tokens)]
|
||||
if self._config.max_input_tokens:
|
||||
fractions.append(("input", usage_accum.input, self._config.max_input_tokens))
|
||||
if self._config.max_output_tokens:
|
||||
fractions.append(("output", usage_accum.output, self._config.max_output_tokens))
|
||||
|
||||
highest_fraction = 0.0
|
||||
trigger_reason = ""
|
||||
trigger_used = 0
|
||||
trigger_budget = 0
|
||||
|
||||
for reason, used, limit in fractions:
|
||||
frac = used / limit
|
||||
if frac > highest_fraction:
|
||||
highest_fraction = frac
|
||||
trigger_reason = reason
|
||||
trigger_used = used
|
||||
trigger_budget = limit
|
||||
|
||||
if highest_fraction >= self._config.hard_stop_threshold:
|
||||
logger.warning("Token budget hard stop triggered for run %s: %s limit exceeded", run_id, trigger_reason)
|
||||
stop_text = _BUDGET_EXCEEDED_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget)
|
||||
return self._build_hard_stop_update(last_msg, stop_text)
|
||||
|
||||
if highest_fraction >= self._config.warn_threshold and not self._warned.get(run_id, False):
|
||||
self._warned[run_id] = True
|
||||
percent = highest_fraction * 100
|
||||
warn_text = _BUDGET_WARNING_MSG.format(reason=trigger_reason, used=trigger_used, budget=trigger_budget, percent=percent)
|
||||
logger.info("Token budget warning triggered for run %s: %s limit at %.1f%%", run_id, trigger_reason, percent)
|
||||
# queue warning for wrap_model_call
|
||||
warnings = self._pending_warnings.setdefault(run_id, [])
|
||||
warnings.append(warn_text)
|
||||
return None
|
||||
|
||||
return None
|
||||
|
||||
@override
|
||||
def after_model(self, state: AgentState, runtime: Runtime) -> dict | None:
|
||||
return self._apply(state, runtime)
|
||||
|
||||
@override
|
||||
async def aafter_model(self, state: AgentState, runtime: Runtime) -> dict | None:
|
||||
return self._apply(state, runtime)
|
||||
|
||||
def _drain_pending_warnings(self, runtime: Runtime) -> list[str]:
|
||||
if not self._config.enabled:
|
||||
return []
|
||||
|
||||
run_id = self._get_run_id(runtime)
|
||||
warnings = self._pending_warnings.pop(run_id, None)
|
||||
return warnings or []
|
||||
|
||||
def _inject_warnings(self, request: ModelRequest, warnings: list[str]) -> ModelRequest:
|
||||
if not warnings:
|
||||
return request
|
||||
|
||||
merged_text = "\n\n".join(warnings)
|
||||
warning_msg = HumanMessage(content=merged_text, name="budget_warning")
|
||||
|
||||
messages = getattr(request, "messages", [])
|
||||
new_messages = list(messages) + [warning_msg]
|
||||
return request.override(messages=new_messages)
|
||||
|
||||
@override
|
||||
def wrap_model_call(self, request: ModelRequest, handler: Callable[[ModelRequest], ModelResponse]) -> ModelCallResult:
|
||||
|
||||
warnings = self._drain_pending_warnings(request.runtime)
|
||||
request = self._inject_warnings(request, warnings)
|
||||
|
||||
return handler(request)
|
||||
|
||||
@override
|
||||
async def awrap_model_call(self, request: ModelRequest, handler: Callable[[ModelRequest], Awaitable[ModelResponse]]) -> ModelCallResult:
|
||||
warnings = self._drain_pending_warnings(request.runtime)
|
||||
request = self._inject_warnings(request, warnings)
|
||||
return await handler(request)
|
||||
@ -33,6 +33,7 @@ from deerflow.config.subagents_config import SubagentsAppConfig, load_subagents_
|
||||
from deerflow.config.suggestions_config import SuggestionsConfig
|
||||
from deerflow.config.summarization_config import SummarizationConfig, load_summarization_config_from_dict
|
||||
from deerflow.config.title_config import TitleConfig, load_title_config_from_dict
|
||||
from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
from deerflow.config.token_usage_config import TokenUsageConfig
|
||||
from deerflow.config.tool_config import ToolConfig, ToolGroupConfig
|
||||
from deerflow.config.tool_output_config import ToolOutputConfig
|
||||
@ -98,6 +99,7 @@ class AppConfig(BaseModel):
|
||||
),
|
||||
)
|
||||
token_usage: TokenUsageConfig = Field(default_factory=TokenUsageConfig, description="Token usage tracking configuration")
|
||||
token_budget: TokenBudgetConfig = Field(default_factory=TokenBudgetConfig, description="Token Budget tracking and limits configuration.")
|
||||
models: list[ModelConfig] = Field(default_factory=list, description="Available models")
|
||||
sandbox: SandboxConfig = Field(
|
||||
description=format_field_description(
|
||||
|
||||
@ -0,0 +1,21 @@
|
||||
"""Config for token budget middleware."""
|
||||
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
|
||||
|
||||
class TokenBudgetConfig(BaseModel):
|
||||
"""Configuration for per-run token budget enforcement."""
|
||||
|
||||
enabled: bool = Field(default=False, description="Whether to enable per-run token budget enforcement.")
|
||||
max_tokens: int = Field(default=200000, ge=1000, description="Maximum total tokens (input + output) allowed per run.")
|
||||
max_input_tokens: int | None = Field(default=None, ge=1, description="Optional separate limit for input tokens only.")
|
||||
max_output_tokens: int | None = Field(default=None, ge=1, description="Optional separate limit for output tokens only.")
|
||||
warn_threshold: float = Field(default=0.8, ge=0.0, le=1.0, description="Fraction of max_tokens at which a soft warning is injected. E.g., 0.8 means warn at 80% of max_tokens")
|
||||
hard_stop_threshold: float = Field(default=1.0, ge=0.0, le=1.0, description=("Fraction of max_tokens at which tool calls are stripped and the agent is forced to produce a final answer. E.g., 1.0 means stop at 100% of max_tokens."))
|
||||
|
||||
@model_validator(mode="after")
|
||||
def validate_thresholds(self) -> "TokenBudgetConfig":
|
||||
"""Ensure hard stop cannot trigger before the warning."""
|
||||
if self.hard_stop_threshold < self.warn_threshold:
|
||||
raise ValueError("hard_stop_threshold must be >= warn_threshold")
|
||||
return self
|
||||
@ -671,7 +671,7 @@ def test_loop_detection_custom_middleware(mock_create_agent):
|
||||
mw_types = [type(m).__name__ for m in middleware]
|
||||
# Default LoopDetectionMiddleware must not also appear.
|
||||
assert "LoopDetectionMiddleware" not in mw_types
|
||||
# Custom replacement still sits immediately before ClarificationMiddleware.
|
||||
# Custom replacement sits immediately before TokenBudgetMiddleware and ClarificationMiddleware.
|
||||
assert mw_types[-1] == "ClarificationMiddleware"
|
||||
assert mw_types[-2] == "MyLoopDetection"
|
||||
|
||||
|
||||
165
backend/tests/test_token_budget_middleware.py
Normal file
165
backend/tests/test_token_budget_middleware.py
Normal file
@ -0,0 +1,165 @@
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
|
||||
from deerflow.agents.middlewares.token_budget_middleware import TokenBudgetMiddleware
|
||||
from deerflow.config.token_budget_config import TokenBudgetConfig
|
||||
|
||||
|
||||
def _make_runtime(thread_id="test-thread", run_id="test-run"):
|
||||
runtime = MagicMock()
|
||||
runtime.context = {"thread_id": thread_id, "run_id": run_id}
|
||||
return runtime
|
||||
|
||||
|
||||
def _make_request(messages, runtime):
|
||||
request = MagicMock()
|
||||
request.messages = list(messages)
|
||||
request.runtime = runtime
|
||||
|
||||
def override_fn(messages=None, **kwags):
|
||||
new_req = MagicMock()
|
||||
new_req.messages = messages if messages is not None else request.messages
|
||||
new_req.runtime = request.runtime
|
||||
return new_req
|
||||
|
||||
request.override = override_fn
|
||||
return request
|
||||
|
||||
|
||||
def _capture_handler():
|
||||
captured: list = []
|
||||
|
||||
def handler(req):
|
||||
captured.append(req)
|
||||
return MagicMock()
|
||||
|
||||
return captured, handler
|
||||
|
||||
|
||||
def _make_state_with_usage(total: int, input_tk: int = 0, output_tk: int = 0, tool_calls=None, content=""):
|
||||
"""Build a state dict with a single AIMessage containing usage."""
|
||||
if input_tk == 0 and output_tk == 0:
|
||||
input_tk = total
|
||||
msg = AIMessage(id="test-msg", content=content, tool_calls=tool_calls or [], usage_metadata={"input_tokens": input_tk, "output_tokens": output_tk, "total_tokens": total})
|
||||
return {"messages": [msg]}
|
||||
|
||||
|
||||
class TestTokenBudgetTracking:
|
||||
def test_no_usage_metadata_returns_none(self):
|
||||
config = TokenBudgetConfig(max_tokens=1000, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
|
||||
state = {"messages": [AIMessage(content="hello", tool_calls=[])]}
|
||||
result = mw._apply(state, _make_runtime())
|
||||
assert result is None
|
||||
|
||||
def test_below_threshold_returns_none(self):
|
||||
config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
|
||||
state = _make_state_with_usage(total=50000)
|
||||
result = mw._apply(state, _make_runtime())
|
||||
assert result is None
|
||||
|
||||
def test_warning_threshold_injects_warning_and_returns_none(self):
|
||||
config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
|
||||
# history with multiple AIMessages that add up to 85000 tokens (>80%)
|
||||
msg1 = AIMessage(id="msg1", content="1", usage_metadata={"total_tokens": 45000, "input_tokens": 45000, "output_tokens": 0})
|
||||
msg2 = ToolMessage(content="ok", tool_call_id="call1")
|
||||
msg3 = AIMessage(id="msg3", content="3", usage_metadata={"total_tokens": 45000, "input_tokens": 45000, "output_tokens": 0})
|
||||
|
||||
state = {"messages": [msg1, msg2, msg3]}
|
||||
result = mw._apply(state, _make_runtime())
|
||||
|
||||
# should queue warning but not mutate state (return None)
|
||||
assert result is None
|
||||
assert len(mw._pending_warnings["test-run"]) == 1
|
||||
assert "TOKEN BUDGET WARNING" in mw._pending_warnings["test-run"][0]
|
||||
|
||||
|
||||
class TestTokenBudgetWarning:
|
||||
def test_warn_injected_at_next_model_call(self):
|
||||
config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
runtime = _make_runtime()
|
||||
|
||||
# trigger warning queue
|
||||
mw._apply(_make_state_with_usage(total=85000), runtime)
|
||||
|
||||
ai_msg = AIMessage(content="", tool_calls=[{"name": "test", "args": {}, "id": "1"}])
|
||||
tool_msg = ToolMessage(content="ok", tool_call_id="1")
|
||||
|
||||
request = _make_request([ai_msg, tool_msg], runtime)
|
||||
|
||||
captured, handler = _capture_handler()
|
||||
mw.wrap_model_call(request, handler)
|
||||
|
||||
sent = captured[0].messages
|
||||
|
||||
assert sent[0] is ai_msg
|
||||
assert sent[1] is tool_msg
|
||||
assert isinstance(sent[2], HumanMessage)
|
||||
assert sent[2].name == "budget_warning"
|
||||
assert "TOKEN BUDGET WARNING" in sent[2].content
|
||||
|
||||
def test_warn_only_once_per_run(self):
|
||||
config = TokenBudgetConfig(max_tokens=100000, warn_threshold=0.8, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
runtime = _make_runtime()
|
||||
|
||||
mw._apply(_make_state_with_usage(total=85000), runtime)
|
||||
|
||||
assert len(mw._pending_warnings["test-run"]) == 1
|
||||
|
||||
# call 2: still above threshold, but already warning -> no second enqueue
|
||||
mw._apply(_make_state_with_usage(total=90000), runtime)
|
||||
assert len(mw._pending_warnings["test-run"]) == 1
|
||||
|
||||
|
||||
class TestTokenBudgetHardStop:
|
||||
def test_hard_stop_strip_tool_calls(self):
|
||||
config = TokenBudgetConfig(max_tokens=100000, hard_stop_threshold=1.0, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
|
||||
tool_calls = [{"name": "bash", "args": {"command": "ls"}, "id": "call_1"}]
|
||||
state = _make_state_with_usage(total=105000, tool_calls=tool_calls, content="Thinking")
|
||||
|
||||
res = mw._apply(state, _make_runtime())
|
||||
|
||||
assert res is not None
|
||||
msgs = res["messages"]
|
||||
assert len(msgs) == 1
|
||||
|
||||
# tool calls must be stripped
|
||||
assert msgs[0].tool_calls == []
|
||||
# content must have the warning appended
|
||||
assert "Thinking" in msgs[0].content
|
||||
assert "TOKEN BUDGET EXCEEDED" in msgs[0].content
|
||||
|
||||
|
||||
class TestIndependentDimensions:
|
||||
def test_input_tokens_trigger_limit(self):
|
||||
config = TokenBudgetConfig(max_tokens=100000, max_input_tokens=10000, warn_threshold=0.8, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
|
||||
# total is safe (10k < 100k) but input is over limit (9k >= 8k)
|
||||
state = _make_state_with_usage(total=10000, input_tk=9000, output_tk=1000)
|
||||
mw._apply(state, _make_runtime())
|
||||
|
||||
warnings = mw._pending_warnings["test-run"]
|
||||
assert len(warnings) == 1
|
||||
assert "input token" in warnings[0]
|
||||
|
||||
def test_output_tokens_trigger_limit(self):
|
||||
config = TokenBudgetConfig(max_tokens=100_000, max_output_tokens=5_000, hard_stop_threshold=1.0, enabled=True)
|
||||
mw = TokenBudgetMiddleware.from_config(config)
|
||||
|
||||
# Total is safe (10k < 100k) but output is over hard limit (6k >= 5k)
|
||||
state = _make_state_with_usage(total=10_000, input_tk=4000, output_tk=6000)
|
||||
result = mw._apply(state, _make_runtime())
|
||||
|
||||
assert result is not None
|
||||
assert "output token" in result["messages"][0].content
|
||||
@ -15,7 +15,7 @@
|
||||
# ============================================================================
|
||||
# Bump this number when the config schema changes.
|
||||
# Run `make config-upgrade` to merge new fields into your local config.yaml.
|
||||
config_version: 14
|
||||
config_version: 15
|
||||
|
||||
# ============================================================================
|
||||
# Logging
|
||||
@ -32,6 +32,22 @@ log_level: info
|
||||
token_usage:
|
||||
enabled: true
|
||||
|
||||
# ============================================================================
|
||||
# Token Budget — Per-run token limits
|
||||
# ============================================================================
|
||||
# Prevents runaway API costs by enforcing hard token limits per run.
|
||||
# When warn_threshold is crossed, the agent receives an in-context warning.
|
||||
# When hard_stop_threshold is crossed, tool_calls are stripped and the agent
|
||||
# is forced to produce a final answer immediately.
|
||||
token_budget:
|
||||
enabled: false # Set to true to activate budget enforcement
|
||||
max_tokens: 200000 # Total token limit (input + output) per run
|
||||
max_input_tokens: null # Optional separate input-only limit
|
||||
max_output_tokens: null # Optional separate output-only limit
|
||||
warn_threshold: 0.8 # Warn at 80% of the budget
|
||||
hard_stop_threshold: 1.0 # Force stop at 100% of the budget
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Models Configuration
|
||||
# ============================================================================
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user