mirror of
https://github.com/HKUDS/OpenSpace.git
synced 2026-08-28 05:15:00 +00:00
214 lines
6.7 KiB
Python
214 lines
6.7 KiB
Python
"""Token budget continuation tracking.
|
|
|
|
Implementations:
|
|
- query/tokenBudget.ts
|
|
- utils/tokenBudget.ts
|
|
|
|
This module intentionally tracks *turn output tokens*, not context-window
|
|
tokens. Context-window pressure remains owned by services.compact
|
|
auto-compact. Token budget is the OpenSpace "+500k / spend 2m tokens" feature: when
|
|
the model naturally stops before spending the requested output budget, the loop
|
|
injects a meta user nudge and continues; when the budget is reached, or progress
|
|
has diminishing returns, the natural stop is allowed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
import sys
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import Literal
|
|
|
|
from openspace.llm.types import TokenUsage
|
|
|
|
COMPLETION_THRESHOLD = 0.9
|
|
DIMINISHING_THRESHOLD = 500
|
|
|
|
_SHORTHAND_START_RE = re.compile(r"^\s*\+(\d+(?:\.\d+)?)\s*(k|m|b)\b", re.I)
|
|
_SHORTHAND_END_RE = re.compile(r"\s\+(\d+(?:\.\d+)?)\s*(k|m|b)\s*[.!?]?\s*$", re.I)
|
|
_VERBOSE_RE = re.compile(
|
|
r"\b(?:use|spend)\s+(\d+(?:\.\d+)?)\s*(k|m|b)\s*tokens?\b",
|
|
re.I,
|
|
)
|
|
_VERBOSE_RE_G = re.compile(_VERBOSE_RE.pattern, re.I)
|
|
|
|
_MULTIPLIERS = {
|
|
"k": 1_000,
|
|
"m": 1_000_000,
|
|
"b": 1_000_000_000,
|
|
}
|
|
_SLOTS_KW = {"slots": True} if sys.version_info >= (3, 10) else {}
|
|
|
|
TokenBudgetAction = Literal["continue", "stop"]
|
|
|
|
|
|
@dataclass(**_SLOTS_KW)
|
|
class TokenBudgetCompletionEvent:
|
|
continuation_count: int
|
|
pct: int
|
|
turn_tokens: int
|
|
budget: int
|
|
diminishing_returns: bool
|
|
duration_ms: int
|
|
|
|
def to_dict(self) -> dict[str, int | bool]:
|
|
return {
|
|
"continuationCount": self.continuation_count,
|
|
"pct": self.pct,
|
|
"turnTokens": self.turn_tokens,
|
|
"budget": self.budget,
|
|
"diminishingReturns": self.diminishing_returns,
|
|
"durationMs": self.duration_ms,
|
|
# snake_case mirrors the rest of OS runtime event payloads.
|
|
"continuation_count": self.continuation_count,
|
|
"turn_tokens": self.turn_tokens,
|
|
"diminishing_returns": self.diminishing_returns,
|
|
"duration_ms": self.duration_ms,
|
|
}
|
|
|
|
|
|
@dataclass(**_SLOTS_KW)
|
|
class TokenBudgetDecision:
|
|
action: TokenBudgetAction
|
|
nudge_message: str | None = None
|
|
continuation_count: int = 0
|
|
pct: int = 0
|
|
turn_tokens: int = 0
|
|
budget: int = 0
|
|
completion_event: TokenBudgetCompletionEvent | None = None
|
|
|
|
|
|
@dataclass(**_SLOTS_KW)
|
|
class BudgetTracker:
|
|
continuation_count: int = 0
|
|
last_delta_tokens: int = 0
|
|
last_global_turn_tokens: int = 0
|
|
started_at: float = field(default_factory=lambda: time.time() * 1000)
|
|
total_output_tokens: int = 0
|
|
|
|
def record_usage(self, usage: TokenUsage | None) -> None:
|
|
if usage is None:
|
|
return
|
|
self.total_output_tokens += max(0, int(usage.output_tokens or 0))
|
|
|
|
def check(
|
|
self,
|
|
*,
|
|
agent_id: str | None,
|
|
budget: int | None,
|
|
global_turn_tokens: int | None = None,
|
|
) -> TokenBudgetDecision:
|
|
return check_token_budget(
|
|
self,
|
|
agent_id,
|
|
budget,
|
|
self.total_output_tokens if global_turn_tokens is None else global_turn_tokens,
|
|
)
|
|
|
|
|
|
def create_budget_tracker() -> BudgetTracker:
|
|
return BudgetTracker()
|
|
|
|
|
|
def _parse_budget_match(value: str, suffix: str) -> int:
|
|
return int(float(value) * _MULTIPLIERS[suffix.lower()])
|
|
|
|
|
|
def parse_token_budget(text: str) -> int | None:
|
|
start_match = _SHORTHAND_START_RE.search(text)
|
|
if start_match:
|
|
return _parse_budget_match(start_match.group(1), start_match.group(2))
|
|
|
|
end_match = _SHORTHAND_END_RE.search(text)
|
|
if end_match:
|
|
return _parse_budget_match(end_match.group(1), end_match.group(2))
|
|
|
|
verbose_match = _VERBOSE_RE.search(text)
|
|
if verbose_match:
|
|
return _parse_budget_match(verbose_match.group(1), verbose_match.group(2))
|
|
|
|
return None
|
|
|
|
|
|
def find_token_budget_positions(text: str) -> list[dict[str, int]]:
|
|
positions: list[dict[str, int]] = []
|
|
start_match = _SHORTHAND_START_RE.search(text)
|
|
if start_match:
|
|
offset = start_match.start() + len(start_match.group(0)) - len(start_match.group(0).lstrip())
|
|
positions.append({"start": offset, "end": start_match.end()})
|
|
|
|
end_match = _SHORTHAND_END_RE.search(text)
|
|
if end_match:
|
|
end_start = end_match.start() + 1
|
|
already_covered = any(
|
|
end_start >= pos["start"] and end_start < pos["end"]
|
|
for pos in positions
|
|
)
|
|
if not already_covered:
|
|
positions.append({"start": end_start, "end": end_match.end()})
|
|
|
|
for match in _VERBOSE_RE_G.finditer(text):
|
|
positions.append({"start": match.start(), "end": match.end()})
|
|
|
|
return positions
|
|
|
|
|
|
def get_budget_continuation_message(pct: int, turn_tokens: int, budget: int) -> str:
|
|
return (
|
|
f"Stopped at {pct}% of token target "
|
|
f"({turn_tokens:,} / {budget:,}). Keep working \u2014 do not summarize."
|
|
)
|
|
|
|
|
|
def check_token_budget(
|
|
tracker: BudgetTracker,
|
|
agent_id: str | None,
|
|
budget: int | None,
|
|
global_turn_tokens: int,
|
|
) -> TokenBudgetDecision:
|
|
if agent_id or budget is None or budget <= 0:
|
|
return TokenBudgetDecision(action="stop")
|
|
|
|
turn_tokens = max(0, int(global_turn_tokens))
|
|
pct = round((turn_tokens / budget) * 100)
|
|
delta_since_last_check = turn_tokens - tracker.last_global_turn_tokens
|
|
|
|
is_diminishing = (
|
|
tracker.continuation_count >= 3
|
|
and delta_since_last_check < DIMINISHING_THRESHOLD
|
|
and tracker.last_delta_tokens < DIMINISHING_THRESHOLD
|
|
)
|
|
|
|
if not is_diminishing and turn_tokens < budget * COMPLETION_THRESHOLD:
|
|
tracker.continuation_count += 1
|
|
tracker.last_delta_tokens = delta_since_last_check
|
|
tracker.last_global_turn_tokens = turn_tokens
|
|
return TokenBudgetDecision(
|
|
action="continue",
|
|
nudge_message=get_budget_continuation_message(pct, turn_tokens, budget),
|
|
continuation_count=tracker.continuation_count,
|
|
pct=pct,
|
|
turn_tokens=turn_tokens,
|
|
budget=budget,
|
|
)
|
|
|
|
if is_diminishing or tracker.continuation_count > 0:
|
|
event = TokenBudgetCompletionEvent(
|
|
continuation_count=tracker.continuation_count,
|
|
pct=pct,
|
|
turn_tokens=turn_tokens,
|
|
budget=budget,
|
|
diminishing_returns=is_diminishing,
|
|
duration_ms=max(0, int(time.time() * 1000 - tracker.started_at)),
|
|
)
|
|
return TokenBudgetDecision(
|
|
action="stop",
|
|
continuation_count=tracker.continuation_count,
|
|
pct=pct,
|
|
turn_tokens=turn_tokens,
|
|
budget=budget,
|
|
completion_event=event,
|
|
)
|
|
|
|
return TokenBudgetDecision(action="stop")
|