From 7bd3618d18c313c16bd09b6e532f80992130c244 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sat, 25 Jul 2026 02:42:53 -0700 Subject: [PATCH 1/4] feat(ai): add agent.assert_tokens() for accumulated token limits agent.assert_tokens(lambda x: x.where('input', '<=', 5000).where('output', '<=', 5000)) Tokens accumulate across every prompt()/stream() in a session (summing each ai message's token uses from the transcript, with a usage fallback), and the fluent TokenQuery evaluates where-clauses over input/output/cache/total. Works on both live-record and replay; to_response() now carries the transcript so replayed turns expose their per-message uses. --- .../tests/units/agents/test_router_agent.py | 3 + .../src/fastapi_startkit/ai/recording.py | 22 ++++- .../src/fastapi_startkit/ai/testing.py | 69 ++++++++++++++ .../tests/ai/test_assert_tokens.py | 93 +++++++++++++++++++ 4 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 fastapi_startkit/tests/ai/test_assert_tokens.py diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py index 3b966ea3..0925ca12 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -20,6 +20,9 @@ def assert_tool_calls(tool: AssertToolCall): await agent.prompt("suggest python developer jobs") agent.assert_tool_called("job_search_tool", assert_tool_calls) + # Tokens accumulate across both turns of the session. + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + async def test_the_router_with_initial_messages(self): with ChatAgent.record( "record_stream.json", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/recording.py b/fastapi_startkit/src/fastapi_startkit/ai/recording.py index 908dd90f..39299a24 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/recording.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/recording.py @@ -132,7 +132,27 @@ def to_response(transcript: list[dict]) -> AgentResponse: content = entry.get("content", "") runtime += (entry.get("response_time") or 0) / 1000 - return AgentResponse(content=content, tool_calls=tool_calls, usage=usage, tool_events=tool_events, runtime=runtime) + return AgentResponse( + content=content, + tool_calls=tool_calls, + usage=usage, + tool_events=tool_events, + runtime=runtime, + transcript=list(transcript), + ) + + +def accumulate_uses(transcript: list[dict], totals: dict) -> None: + """Add every ``ai`` entry's token ``uses`` in ``transcript`` into ``totals`` + (keys: input, output, cache, total).""" + for entry in transcript or []: + if entry.get("type") != "ai": + continue + uses = entry.get("uses") or {} + totals["input"] += uses.get("input_token", 0) + totals["output"] += uses.get("output_token", 0) + totals["cache"] += uses.get("cache_token", 0) + totals["total"] += uses.get("total_token", 0) def chunks_from_transcript(transcript: list[dict]) -> list[str]: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 6f651895..bd78fb79 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -4,6 +4,7 @@ import hashlib import inspect import json +import operator import sys import time from collections.abc import AsyncIterator @@ -21,6 +22,48 @@ def _joined(value: Any) -> str: return "".join(value) if isinstance(value, list) else value +def _new_token_totals() -> dict: + return {"input": 0, "output": 0, "cache": 0, "total": 0} + + +class TokenQuery: + """Fluent predicate over accumulated token totals, e.g.:: + + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + + Fields: ``input`` / ``output`` / ``cache`` / ``total``. + """ + + _OPS = { + "<=": operator.le, + ">=": operator.ge, + "<": operator.lt, + ">": operator.gt, + "==": operator.eq, + "=": operator.eq, + "!=": operator.ne, + } + + def __init__(self, totals: dict) -> None: + self._totals = totals + self._checks: list[tuple[str, str, float]] = [] + + def where(self, field: str, op: str, value: float) -> "TokenQuery": + self._checks.append((field, op, value)) + return self + + def failures(self) -> list[str]: + problems: list[str] = [] + for field, op, value in self._checks: + compare = self._OPS.get(op) + if compare is None: + raise ValueError(f"Unsupported token operator {op!r}") + actual = self._totals.get(field, 0) + if not compare(actual, value): + problems.append(f"{field}={actual} is not {op} {value}") + return problems + + class AgentFake: def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._agent_cls = agent_cls @@ -30,6 +73,7 @@ def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._records: list[dict] = [] self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None + self._tokens: dict = _new_token_totals() def _history(self) -> list: return self._records @@ -73,9 +117,31 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter self._remember(message, self._last_response) def _remember(self, message: str, response: AgentResponse) -> None: + self._accumulate_tokens(response) self._records.append({"role": "user", "content": message}) self._records.append({"role": "assistant", "content": response.content}) + def _accumulate_tokens(self, response: AgentResponse) -> None: + """Add a turn's token usage to the running totals. Prefer the per-message + ``uses`` on the transcript; fall back to the response's summary usage.""" + from . import recording # noqa: PLC0415 + + if response.transcript: + recording.accumulate_uses(response.transcript, self._tokens) + else: + self._tokens["input"] += response.usage.get("input", 0) + self._tokens["output"] += response.usage.get("output", 0) + + def assert_tokens(self, predicate: Callable[[TokenQuery], Any]) -> None: + """Assert on the tokens accumulated across every prompt()/stream() so far. + + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + """ + query = TokenQuery(dict(self._tokens)) + predicate(query) + failures = query.failures() + assert not failures, f"Token assertion failed: {'; '.join(failures)}. Accumulated tokens: {self._tokens}" + def assert_prompt(self, expected: str | Callable[[str], bool]) -> None: if callable(expected): assert any(expected(p) for p in self._prompts), ( @@ -107,6 +173,7 @@ def assert_not_prompted(self) -> None: def reset(self) -> "AgentFake": self._records.clear() self._last_response = None + self._tokens = _new_token_totals() return self def _require_response(self) -> AgentResponse: @@ -194,6 +261,7 @@ def __init__(self, real: Agent, cassette: str | None = None, messages: list | No self._real.messages = self._history # type: ignore[method-assign] self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None + self._tokens: dict = _new_token_totals() def _history(self) -> list: return self._seed_messages + self._records @@ -263,6 +331,7 @@ def _response_from_cache(value: Any) -> AgentResponse: return AgentResponse(content=_joined(value)) def _remember_turn(self, message: str, response: AgentResponse) -> None: + self._accumulate_tokens(response) self._records.append({"role": "user", "content": message}) turn: dict[str, Any] = {"role": "assistant", "content": response.content} if response.tool_calls: diff --git a/fastapi_startkit/tests/ai/test_assert_tokens.py b/fastapi_startkit/tests/ai/test_assert_tokens.py new file mode 100644 index 00000000..8cfe9447 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_assert_tokens.py @@ -0,0 +1,93 @@ +"""Tests for agent.assert_tokens() — assert on accumulated token usage (task #1251). + + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + +Tokens accumulate across every recorded turn (and every ai message within a +turn); the predicate builds where-clauses that must all hold. +""" + +import os +import tempfile +import unittest +from unittest import mock + +from fastapi_startkit.ai import recording +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.response import AgentResponse + + +class TokenAgent(Agent): + pass + + +def _ai(input_tokens: int, output_tokens: int) -> dict: + return recording.ai( + content="reply", + uses={ + "input_token": input_tokens, + "output_token": output_tokens, + "cache_token": 0, + "total_token": input_tokens + output_tokens, + }, + response_time=1.0, + ) + + +def _fake_prompt(responses: list): + queue = list(responses) + + async def prompt(agent_self, message, **kwargs): + return queue.pop(0) + + return mock.patch.object(TokenAgent, "prompt", prompt) + + +class TestAssertTokens(unittest.IsolatedAsyncioTestCase): + async def test_passes_when_accumulated_tokens_are_within_limits(self): + responses = [ + AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=[_ai(100, 20)]), + AgentResponse(content="b", usage={"input": 200, "output": 30}, transcript=[_ai(200, 30)]), + ] + with tempfile.TemporaryDirectory() as tmp: + with _fake_prompt(responses): + with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("a") + await agent.prompt("b") + # accumulated: input=300, output=50 + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + + async def test_fails_when_a_limit_is_exceeded(self): + responses = [AgentResponse(content="a", usage={"input": 400, "output": 20}, transcript=[_ai(400, 20)])] + with tempfile.TemporaryDirectory() as tmp: + with _fake_prompt(responses): + with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("a") + with self.assertRaises(AssertionError): + agent.assert_tokens(lambda x: x.where("input", "<=", 300)) + + async def test_supports_cache_and_total_fields(self): + transcript = [ + recording.ai( + content="a", + uses={"input_token": 100, "output_token": 20, "cache_token": 40, "total_token": 160}, + response_time=1.0, + ) + ] + with tempfile.TemporaryDirectory() as tmp: + with _fake_prompt([AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=transcript)]): + with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("a") + agent.assert_tokens(lambda x: x.where("cache", "<=", 40).where("total", ">=", 160)) + + async def test_accumulates_from_the_cassette_on_replay(self): + responses = [AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=[_ai(100, 20)])] + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with _fake_prompt(responses): + with TokenAgent.record(cassette) as agent: + await agent.prompt("a") + + # Replay: no live prompt available; tokens must come from the cassette. + with TokenAgent.record(cassette) as agent: + await agent.prompt("a") + agent.assert_tokens(lambda x: x.where("input", "==", 100).where("output", "==", 20)) From 3659e42670f269fef2f6e5d6ec82b8fd3763843a Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sat, 25 Jul 2026 03:05:33 -0700 Subject: [PATCH 2/4] style(ai): ruff-format assert_tokens docstring --- fastapi_startkit/src/fastapi_startkit/ai/testing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index bd78fb79..46918704 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -135,7 +135,7 @@ def _accumulate_tokens(self, response: AgentResponse) -> None: def assert_tokens(self, predicate: Callable[[TokenQuery], Any]) -> None: """Assert on the tokens accumulated across every prompt()/stream() so far. - agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) """ query = TokenQuery(dict(self._tokens)) predicate(query) From cd1dee897da6e6c3230d052138e6e4a79dfa64ed Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sat, 25 Jul 2026 12:35:00 -0700 Subject: [PATCH 3/4] style(ai): ruff-format runner.py and test_record_cassette_format.py CI runs 'ruff format --check .' repo-wide; these two files (carried in from #199) were unformatted and failed the check. --- .../src/fastapi_startkit/ai/runner.py | 4 ++- .../tests/ai/test_record_cassette_format.py | 25 ++++++++++++------- 2 files changed, 19 insertions(+), 10 deletions(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index b078b1c4..2d846f19 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -53,7 +53,9 @@ def _record_ai_message( tool_calls = list(getattr(message, "tool_calls", None) or []) uses = recording.uses_from_usage_metadata(getattr(message, "usage_metadata", None)) self._transcript.append( - recording.ai(content=content, tool_calls=tool_calls, uses=uses, response_time=response_time_ms, chunks=chunks) + recording.ai( + content=content, tool_calls=tool_calls, uses=uses, response_time=response_time_ms, chunks=chunks + ) ) if tool_calls: self._requested_tool_calls = tool_calls diff --git a/fastapi_startkit/tests/ai/test_record_cassette_format.py b/fastapi_startkit/tests/ai/test_record_cassette_format.py index e5122e22..92e0b4b9 100644 --- a/fastapi_startkit/tests/ai/test_record_cassette_format.py +++ b/fastapi_startkit/tests/ai/test_record_cassette_format.py @@ -30,12 +30,15 @@ class TestNewCassetteFormat(unittest.IsolatedAsyncioTestCase): async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): response = AgentResponse( transcript=[ - {"type": "ai", "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - "uses": {"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, - "response_time": 12.0}, - {"type": "tool_response", "content_type": "json", "content": "[{\"id\": 2}]", "response_time": 5.0}, + { + "type": "ai", + "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "uses": {"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, + "response_time": 12.0, + }, + {"type": "tool_response", "content_type": "json", "content": '[{"id": 2}]', "response_time": 5.0}, ], - content="[{\"id\": 2}]", + content='[{"id": 2}]', tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], ) with tempfile.TemporaryDirectory() as tmp: @@ -57,11 +60,15 @@ async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): async def test_replaying_a_new_format_cassette_reconstructs_tool_calls(self): response = AgentResponse( transcript=[ - {"type": "ai", "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - "uses": {"input_token": 117, "output_token": 22}, "response_time": 12.0}, - {"type": "tool_response", "content_type": "json", "content": "[{\"id\": 2}]", "response_time": 5.0}, + { + "type": "ai", + "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "uses": {"input_token": 117, "output_token": 22}, + "response_time": 12.0, + }, + {"type": "tool_response", "content_type": "json", "content": '[{"id": 2}]', "response_time": 5.0}, ], - content="[{\"id\": 2}]", + content='[{"id": 2}]', tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], ) with tempfile.TemporaryDirectory() as tmp: From aba408fab81b4b73dd42bff7a2d49befee764010 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sat, 25 Jul 2026 12:44:36 -0700 Subject: [PATCH 4/4] test(orm): fix polymorphic relation tests to use the record relationship The Like model's MorphTo relationship is named 'record'; the tests referenced a nonexistent 'like.log' attribute, which returned None and failed to await / assert. Aligns with the eager-load case that already uses with_('record'). --- .../sqlite/relationships/test_sqlite_polymorphic.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py b/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py index f86fc1ef..675cf79a 100644 --- a/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py +++ b/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py @@ -9,10 +9,10 @@ class TestRelationships(TestCase): async def test_can_get_polymorphic_relation(self): likes = await Like.get() for like in likes: - record = await like.log + record = await like.record assert isinstance(record, (Articles, Product)) async def test_can_get_eager_load_polymorphic_relation(self): likes = await Like.with_("record").get() for like in likes: - assert isinstance(like.log, (Articles, Product)) + assert isinstance(like.record, (Articles, Product))