99 lines
3.9 KiB
Python
99 lines
3.9 KiB
Python
"""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]
|
|
|
|
|
|
async def test_budget_middleware_fires_on_real_agent_chat(make_client_factory) -> None:
|
|
"""F8 (real-client half): registering BudgetMiddleware on a real Agent (layered chat client
|
|
that carries ChatMiddlewareLayer) and running an actual chat call short-circuits with
|
|
BudgetExceeded — the middleware<->client integration the prior suite never exercised (the
|
|
minimal-base stand-in silently no-ops the middleware)."""
|
|
from agent_framework import Agent
|
|
|
|
client = make_client_factory("ok", tokens=8)("proposer") # 8 tokens/reply > cap 5
|
|
meter = TokenMeter(Budget(max_tokens=5, max_rounds=10))
|
|
agent = Agent(
|
|
client, "propose a measure", name="proposer", middleware=[BudgetMiddleware(meter)]
|
|
)
|
|
with pytest.raises(BudgetExceeded) as exc:
|
|
await agent.run("hi")
|
|
assert exc.value.kind == "tokens"
|
|
assert meter.tokens == 8 # charged from the synthetic UsageDetails via the middleware
|
|
|
|
|
|
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}"
|