102 lines
3.9 KiB
Python
102 lines
3.9 KiB
Python
"""Token meter + budget ChatMiddleware fed from REAL UsageDetails (B4 / D6).
|
|
|
|
The hard token cap is enforced off the provider's own ``UsageDetails`` — **never** a
|
|
word-count proxy (research 03 Rec 3; the Fase 1 ``_word_tokens`` theatre is retired). A
|
|
``BudgetMiddleware`` reads ``response.usage_details["total_token_count"]`` after each chat
|
|
call and short-circuits with ``BudgetExceeded`` the moment the cap is crossed.
|
|
|
|
``strict_usage`` (default ``True``) makes a missing usage a HARD FAIL (``UsageUnavailable``):
|
|
a usage regression must never silently disable the cap. Test doubles that legitimately supply
|
|
a synthetic ``UsageDetails`` do not trip it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
|
|
from agent_framework import ChatContext, ChatMiddleware
|
|
|
|
|
|
class BudgetExceeded(RuntimeError):
|
|
"""Raised the moment a runtime cap (tokens or rounds) is crossed (B4).
|
|
|
|
Carries the breached ``kind`` ("tokens" | "rounds"), the ``limit`` set, and the
|
|
``observed`` value that crossed it — a structured stop event, never a silent hang.
|
|
"""
|
|
|
|
def __init__(self, kind: str, limit: int, observed: int) -> None:
|
|
self.kind = kind
|
|
self.limit = limit
|
|
self.observed = observed
|
|
super().__init__(f"budget exceeded: {kind} limit={limit} observed={observed}")
|
|
|
|
|
|
class UsageUnavailable(RuntimeError):
|
|
"""Raised when ``strict_usage`` is on and a response carries no usable usage — the cap
|
|
must fail closed rather than silently stop counting (research 03 Rec 3)."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Budget:
|
|
"""Hard token + round/iteration caps, required at startup (A4 / D6).
|
|
|
|
Refuses to construct without positive caps — fail-fast, never an unbounded loop.
|
|
"""
|
|
|
|
max_tokens: int
|
|
max_rounds: int
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.max_tokens <= 0:
|
|
raise ValueError(f"max_tokens must be positive, got {self.max_tokens}")
|
|
if self.max_rounds <= 0:
|
|
raise ValueError(f"max_rounds must be positive, got {self.max_rounds}")
|
|
|
|
|
|
class TokenMeter:
|
|
"""Accumulates token and round usage against a ``Budget``; raises the moment a cap is
|
|
crossed."""
|
|
|
|
def __init__(self, budget: Budget) -> None:
|
|
self.budget = budget
|
|
self.tokens = 0
|
|
self.rounds = 0
|
|
|
|
def charge(self, tokens: int) -> int:
|
|
"""Add ``tokens`` to the running total; raise ``BudgetExceeded`` if over cap."""
|
|
self.tokens += tokens
|
|
if self.tokens > self.budget.max_tokens:
|
|
raise BudgetExceeded("tokens", self.budget.max_tokens, self.tokens)
|
|
return self.tokens
|
|
|
|
def tick_round(self) -> int:
|
|
"""Increment the round counter; raise ``BudgetExceeded`` if over cap."""
|
|
self.rounds += 1
|
|
if self.rounds > self.budget.max_rounds:
|
|
raise BudgetExceeded("rounds", self.budget.max_rounds, self.rounds)
|
|
return self.rounds
|
|
|
|
|
|
class BudgetMiddleware(ChatMiddleware):
|
|
"""Chat middleware that charges a ``TokenMeter`` from each response's real
|
|
``UsageDetails`` and short-circuits when the cap is crossed."""
|
|
|
|
def __init__(self, meter: TokenMeter, *, strict_usage: bool = True) -> None:
|
|
self._meter = meter
|
|
self._strict = strict_usage
|
|
|
|
async def process(
|
|
self, context: ChatContext, call_next: Callable[[], Awaitable[None]]
|
|
) -> None:
|
|
await call_next()
|
|
usage = getattr(context.result, "usage_details", None)
|
|
total = usage.get("total_token_count") if usage is not None else None
|
|
if total is None:
|
|
if self._strict:
|
|
raise UsageUnavailable(
|
|
"usage_details missing total_token_count; the cap must fail closed "
|
|
"(set strict_usage=False only for non-counting paths)"
|
|
)
|
|
return
|
|
self._meter.charge(int(total)) # raises BudgetExceeded if over cap
|