AI: add agent.assert_tokens() for accumulated token limits (#200)

* 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.

* style(ai): ruff-format assert_tokens docstring

* 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.

* 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').

Author: tmgbedu

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",

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]:

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

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:

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))

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:

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))