From 3104e18b0742d11eebd58e0262166ee290a219ef Mon Sep 17 00:00:00 2001 From: Kjell Tore Guttormsen Date: Wed, 24 Jun 2026 13:31:39 +0200 Subject: [PATCH] feat(fase2): real-UsageDetails token meter + budget middleware --- src/portfolio_optimiser/budget.py | 102 ++++++++++++++++++++++++++++++ tests/test_budget.py | 83 ++++++++++++++++++++++++ 2 files changed, 185 insertions(+) create mode 100644 src/portfolio_optimiser/budget.py create mode 100644 tests/test_budget.py diff --git a/src/portfolio_optimiser/budget.py b/src/portfolio_optimiser/budget.py new file mode 100644 index 0000000..58b0790 --- /dev/null +++ b/src/portfolio_optimiser/budget.py @@ -0,0 +1,102 @@ +"""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 diff --git a/tests/test_budget.py b/tests/test_budget.py new file mode 100644 index 0000000..058776e --- /dev/null +++ b/tests/test_budget.py @@ -0,0 +1,83 @@ +"""Step 4 tests — token meter + budget middleware off REAL UsageDetails (no LLM). + +Usage is hand-built (synthetic ``UsageDetails``), so these run in CI without an endpoint. +The strict-usage guard hard-fails on a missing usage; the meta-guard proves the meter is +fed from ``UsageDetails`` and never a word-count proxy. Pattern: tests/spikes/test_harness.py. +""" + +import re +from pathlib import Path + +import pytest +from agent_framework import ChatResponse, Message, UsageDetails + +from portfolio_optimiser.budget import ( + Budget, + BudgetExceeded, + BudgetMiddleware, + TokenMeter, + UsageUnavailable, +) + + +class _Ctx: + """Minimal ChatContext stand-in carrying the post-call ``result``.""" + + def __init__(self, result: object) -> None: + self.result = result + + +async def _noop() -> None: + return None + + +def _resp(total: int | None) -> ChatResponse: + usage = UsageDetails(total_token_count=total) if total is not None else None + return ChatResponse( + messages=[Message(role="assistant", contents=["x"])], usage_details=usage + ) + + +async def test_meter_reads_total_token_count_and_accumulates() -> None: + meter = TokenMeter(Budget(max_tokens=1000, max_rounds=10)) + mw = BudgetMiddleware(meter) + await mw.process(_Ctx(_resp(40)), _noop) # type: ignore[arg-type] + await mw.process(_Ctx(_resp(50)), _noop) # type: ignore[arg-type] + assert meter.tokens == 90 # read from UsageDetails, accumulated across calls + + +async def test_cap_crossed_raises_budget_exceeded() -> None: + meter = TokenMeter(Budget(max_tokens=50, max_rounds=10)) + mw = BudgetMiddleware(meter) + with pytest.raises(BudgetExceeded) as exc: + await mw.process(_Ctx(_resp(60)), _noop) # type: ignore[arg-type] + assert exc.value.kind == "tokens" + assert exc.value.limit == 50 + + +def test_non_positive_cap_rejected() -> None: + with pytest.raises(ValueError): + Budget(max_tokens=0, max_rounds=5) + with pytest.raises(ValueError): + Budget(max_tokens=5, max_rounds=0) + + +async def test_strict_usage_none_hard_fails() -> None: + meter = TokenMeter(Budget(max_tokens=100, max_rounds=10)) + mw = BudgetMiddleware(meter, strict_usage=True) + with pytest.raises(UsageUnavailable): + await mw.process(_Ctx(_resp(None)), _noop) # type: ignore[arg-type] + + +def test_no_word_count_token_proxy_in_src() -> None: + # The meter is fed from real UsageDetails, never a len(text.split()) word-count proxy + # (research 03 Rec 3 — the Fase 1 _word_tokens anti-pattern is retired). NOTE: a bare + # `.split()` grep is wrong — retrieval.py legitimately uses .split() for KEYWORD scoring + # (a [0,1] overlap ratio, not a token count). So guard the specific proxy signature. + proxy = re.compile(r"len\(\s*[^)]*\.split\(\)\s*\)") + offenders = [ + py.name + for py in Path("src/portfolio_optimiser").rglob("*.py") + if proxy.search(py.read_text(encoding="utf-8")) + ] + assert offenders == [], f"word-count token proxy found in: {offenders}"