1"""Token and USD accounting with a hard cap. One tracker per run (and a global one per process)."""2from __future__ import annotations34import threading5from dataclasses import dataclass, field678class BudgetExceeded(RuntimeError):9 pass101112@dataclass13class Usage:14 calls: int = 015 cached_hits: int = 016 input_tokens: int = 017 cached_input_tokens: int = 018 output_tokens: int = 019 reasoning_tokens: int = 020 usd: float = 0.021 failures: int = 022 invalid: int = 023 latency_s: float = 0.024 by_model: dict[str, dict[str, float]] = field(default_factory=dict)2526 def as_dict(self) -> dict:27 d = {k: getattr(self, k) for k in ("calls", "cached_hits", "input_tokens", "cached_input_tokens",28 "output_tokens", "reasoning_tokens", "usd", "failures", "invalid", "latency_s")}29 d["by_model"] = self.by_model30 return d313233class BudgetTracker:34 def __init__(self, usd_cap: float, token_cap: int | None, prices: dict[str, dict[str, float]]) -> None:35 self.usd_cap = usd_cap36 self.token_cap = token_cap37 self.prices = prices38 self.usage = Usage()39 self._lock = threading.Lock()4041 def price(self, model: str, input_tokens: int, cached_tokens: int, output_tokens: int) -> float:42 p = self.prices.get(model)43 if p is None:44 # unknown model: use the most expensive known price so the cap is conservative45 p = max(self.prices.values(), key=lambda d: d["output"]) if self.prices else {"input": 5, "cached": 0.5, "output": 25}46 uncached = max(input_tokens - cached_tokens, 0)47 return (uncached * p["input"] + cached_tokens * p["cached"] + output_tokens * p["output"]) / 1e64849 def check(self) -> None:50 if self.usage.usd >= self.usd_cap:51 raise BudgetExceeded(f"USD cap reached: {self.usage.usd:.4f} >= {self.usd_cap}")52 if self.token_cap is not None and self.usage.input_tokens + self.usage.output_tokens >= self.token_cap:53 raise BudgetExceeded("token cap reached")5455 def record(self, model: str, input_tokens: int, cached_tokens: int, output_tokens: int,56 reasoning_tokens: int = 0, latency_s: float = 0.0, cached_hit: bool = False) -> float:57 with self._lock:58 u = self.usage59 u.calls += 160 if cached_hit:61 u.cached_hits += 162 return 0.063 cost = self.price(model, input_tokens, cached_tokens, output_tokens)64 u.input_tokens += input_tokens65 u.cached_input_tokens += cached_tokens66 u.output_tokens += output_tokens67 u.reasoning_tokens += reasoning_tokens68 u.usd += cost69 u.latency_s += latency_s70 m = u.by_model.setdefault(model, {"calls": 0, "input_tokens": 0, "output_tokens": 0, "usd": 0.0})71 m["calls"] += 172 m["input_tokens"] += input_tokens73 m["output_tokens"] += output_tokens74 m["usd"] += cost75 return cost76