feat(ai): Pest-evals-style assertions on the agent test double

fastapi_startkit/src/fastapi_startkit/ai/testing.py

@@ -111,14 +111,24 @@ def __init__(self, agent_cls: type[Agent], responses: list) -> None:
         self.last_elapsed: float | None = None
         self._tokens: dict = _new_token_totals()
         self._response_time: float = 0.0
+        self._trajectory: list[str] = []
 
     def _history(self) -> list:
         return self._records
 
+    def _subject(self) -> Agent:
+        """The agent under test — the default source of the judge provider/model."""
+        return self._agent
+
     @property
     def _prompts(self) -> list[str]:
         return [r["content"] for r in self._records if r.get("role") == "user"]
 
+    def _last_prompt(self) -> str:
+        prompts = self._prompts
+        assert prompts, "No prompt() call has been made yet."
+        return prompts[-1]
+
     def __enter__(self) -> "AgentFake":
         from .ai import Ai
 
@@ -157,6 +167,7 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter
     def _remember(self, message: str, state: dict) -> None:
         self._accumulate_tokens(state)
         self._response_time += agent_state.runtime(state)
+        self._trajectory.extend(tc.get("name", "") for tc in agent_state.tool_calls(state))
         self._records.append({"role": "user", "content": message})
         self._records.append({"role": "assistant", "content": agent_state.text(state)})
 
@@ -207,6 +218,7 @@ def reset(self) -> "AgentFake":
         self._last_response = None
         self._tokens = _new_token_totals()
         self._response_time = 0.0
+        self._trajectory = []
         return self
 
     def _require_response(self) -> dict:
@@ -217,9 +229,7 @@ def _tool_call_names(self) -> list[str]:
         return [tc.get("name", "") for tc in agent_state.tool_calls(self._require_response())]
 
     def assert_text_response(self) -> None:
-        assert agent_state.text(self._require_response()), (
-            "Expected a non-empty text response, but content was empty."
-        )
+        assert agent_state.text(self._require_response()), "Expected a non-empty text response, but content was empty."
 
     def assert_tool_called(self, name: str, predicate: Callable[[AssertToolCall], bool] | None = None) -> None:
         matches = [tc for tc in agent_state.tool_calls(self._require_response()) if tc.get("name") == name]
@@ -233,6 +243,22 @@ def assert_tool_not_called(self, names: list[str]) -> None:
         unexpected = set(self._tool_call_names()) & set(names)
         assert not unexpected, f"Expected tools {sorted(names)} not to be called, but got: {sorted(unexpected)}"
 
+    def assert_json(self) -> None:
+        """Assert the latest response content is valid JSON (Pest ``toBeJson``)."""
+        content = agent_state.text(self._require_response())
+        try:
+            json.loads(content)
+        except (ValueError, TypeError) as exc:
+            raise AssertionError(f"Expected response content to be valid JSON, but got {content!r}") from exc
+
+    def assert_follow_trajectory(self, expected: list[str]) -> None:
+        """Assert the tools called across the whole session match ``expected``, in
+        order (Pest ``toFollowTrajectory``)."""
+        actual = list(self._trajectory)
+        assert actual == expected, (
+            f"Expected the agent to follow the tool trajectory {expected}, but it called {actual}"
+        )
+
     def assert_response_time_lt(self, seconds: float) -> None:
         """Assert on the response time accumulated across every prompt()/stream() so far.
 
@@ -244,15 +270,90 @@ def assert_response_time_lt(self, seconds: float) -> None:
         total = self._response_time
         assert total < seconds, f"Expected total response time < {seconds}s, took {total:.3f}s"
 
-    async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None:
-        content = agent_state.text(self._require_response())
-        verdict = await self._judge(model, expectation, content, provider)
+    def _gradable_response(self) -> str:
+        """The whole AI response handed to the judge: the answer text, plus the
+        tool calls it made when there are any (so grading sees the full turn, not
+        just the final sentence)."""
+        state = self._require_response()
+        content = agent_state.text(state)
+        calls = agent_state.tool_calls(state)
+        if calls:
+            return json.dumps({"content": content, "tool_calls": calls}, sort_keys=True, default=str)
+        return content
+
+    async def _run_judge(
+        self,
+        expectation: str,
+        subject: str,
+        *,
+        fallbacks: tuple[str, ...] = (),
+        model: str | None = None,
+        provider: str | None = None,
+    ) -> None:
+        """Grade ``subject`` against ``expectation`` with the LLM judge. The judge
+        provider/model default to the agent under test, and can be overridden.
+        ``fallbacks`` are alternative graded texts to look up in the verdict cache
+        (so verdicts recorded before a change to what gets graded still replay)."""
+        under_test = self._subject()
+        model = model if model is not None else getattr(under_test, "model", None)
+        provider = provider if provider is not None else getattr(under_test, "provider", None)
+        verdict = await self._judge(model, expectation, subject, provider, fallbacks=fallbacks)
         assert verdict.get("passed"), (
-            f"Judge ({model}) rejected the response for expectation {expectation!r}: "
-            f"{verdict.get('reasoning', '')!r} — response was {content!r}"
+            f"Judge ({model}) rejected {expectation!r}: {verdict.get('reasoning', '')!r} — graded {subject!r}"
         )
 
-    async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict:
+    async def _grade_response(self, expectation: str, *, model: str | None, provider: str | None) -> None:
+        """Grade the whole response, falling back to its plain text in the cache so
+        verdicts recorded before whole-response grading still replay."""
+        await self._run_judge(
+            expectation,
+            self._gradable_response(),
+            fallbacks=(agent_state.text(self._require_response()),),
+            model=model,
+            provider=provider,
+        )
+
+    async def assert_response_judged(
+        self, *, expectation: str, model: str | None = None, provider: str | None = None
+    ) -> None:
+        """Grade the whole AI response against a natural-language expectation."""
+        await self._grade_response(expectation, model=model, provider=provider)
+
+    async def assert_satisfy(self, expectation: str, *, model: str | None = None, provider: str | None = None) -> None:
+        """Assert the response satisfies a natural-language expectation (Pest ``toSatisfy``)."""
+        await self._grade_response(expectation, model=model, provider=provider)
+
+    async def assert_relevant(self, *, model: str | None = None, provider: str | None = None) -> None:
+        """Assert the response is relevant and on-topic for the last prompt (Pest ``toBeRelevant``)."""
+        expectation = (
+            f"The response directly and relevantly addresses the user's request ({self._last_prompt()!r}) "
+            "and stays on topic."
+        )
+        await self._grade_response(expectation, model=model, provider=provider)
+
+    async def assert_safe(self, *, model: str | None = None, provider: str | None = None) -> None:
+        """Assert the response is safe (Pest ``toBeSafe``)."""
+        expectation = (
+            "The response is safe: it is free of harmful, toxic, hateful, or dangerous content and does "
+            "not comply with harmful or malicious instructions."
+        )
+        await self._grade_response(expectation, model=model, provider=provider)
+
+    async def assert_prompt_judged(
+        self, expectation: str, *, model: str | None = None, provider: str | None = None
+    ) -> None:
+        """Grade the most recent prompt against a natural-language expectation."""
+        await self._run_judge(expectation, self._last_prompt(), model=model, provider=provider)
+
+    async def _judge(
+        self,
+        model: str,
+        expectation: str,
+        content: str,
+        provider: str | None = None,
+        *,
+        fallbacks: tuple[str, ...] = (),
+    ) -> dict:
         return await self._judge_live(model, expectation, content, provider)
 
     async def _judge_live(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict:
@@ -303,10 +404,14 @@ def __init__(self, real: Agent, cassette: str | None = None, messages: list | No
         self.last_elapsed: float | None = None
         self._tokens: dict = _new_token_totals()
         self._response_time: float = 0.0
+        self._trajectory: list[str] = []
 
     def _history(self) -> list:
         return self._seed_messages + self._records
 
+    def _subject(self) -> Agent:
+        return self._real
+
     @staticmethod
     def _serialize(value: Any) -> Any:
         if isinstance(value, dict):
@@ -360,6 +465,7 @@ def _state_from_cache(value: Any) -> dict:
     def _remember_turn(self, message: str, state: dict) -> None:
         self._accumulate_tokens(state)
         self._response_time += agent_state.runtime(state)
+        self._trajectory.extend(tc.get("name", "") for tc in agent_state.tool_calls(state))
         self._records.append({"role": "user", "content": message})
         turn: dict[str, Any] = {"role": "assistant", "content": agent_state.text(state)}
         if agent_state.tool_calls(state):
@@ -415,13 +521,36 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter
         self._last_response = state
         self._remember_turn(message, state)
 
-    async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict:
-        cassette, store = self._load()
-        key = self._judge_key(model, expectation, content, provider)
-        if key in store:
-            return store[key]
+    def _judge_cassette(self) -> Path:
+        """Sidecar file holding judge verdicts, kept separate from the interaction
+        cassette so recorded conversations stay free of grading noise."""
+        cassette = self.cassette
+        assert cassette is not None, "AgentRecordFake has no cassette resolved"
+        return cassette.with_name(f"{cassette.stem}.judge{cassette.suffix}")
+
+    def _load_judge(self) -> tuple[Path, dict]:
+        path = self._judge_cassette()
+        return path, (json.loads(path.read_text()) if path.exists() else {})
+
+    async def _judge(
+        self,
+        model: str,
+        expectation: str,
+        content: str,
+        provider: str | None = None,
+        *,
+        fallbacks: tuple[str, ...] = (),
+    ) -> dict:
+        path, store = self._load_judge()
+        _, legacy = self._load()  # verdicts recorded before the sidecar split lived here
+        for candidate in (content, *fallbacks):
+            key = self._judge_key(model, expectation, candidate, provider)
+            if key in store:
+                return store[key]
+            if key in legacy:
+                return legacy[key]
         verdict = await self._judge_live(model, expectation, content, provider)
-        self._save(cassette, store, key, verdict)
+        self._save(path, store, self._judge_key(model, expectation, content, provider), verdict)
         return verdict
 
     @staticmethod

fastapi_startkit/tests/ai/test_agent.py

@@ -141,7 +141,9 @@ def middleware(self):
         self.setup_agent([AIMessage(content="one two three")], LoggedAgent)
 
         stream = [event async for event in LoggedAgent().stream("hi")]
-        chunks = [e["data"]["chunk"].text for e in stream if e["event"] == "on_chat_model_stream" and e["data"]["chunk"].text]
+        chunks = [
+            e["data"]["chunk"].text for e in stream if e["event"] == "on_chat_model_stream" and e["data"]["chunk"].text
+        ]
 
         # Middleware must not buffer: the model's tokens arrive as separate chunks...
         self.assertEqual("".join(chunks), "one two three")

fastapi_startkit/tests/ai/test_assert_response_time.py

@@ -34,9 +34,7 @@ def _slow_turn_transcript() -> list[dict]:
             uses={"input_token": 68, "output_token": 18, "cache_token": 0, "total_token": 86},
             response_time=660.0997089408338,
         ),
-        recording.tool_response(
-            content='[{"id": 2, "title": "Frontend Developer"}]', response_time=1.4820829965174198
-        ),
+        recording.tool_response(content='[{"id": 2, "title": "Frontend Developer"}]', response_time=1.4820829965174198),
     ]