refactor(ai): back assert_response_judged with a JudgeAgent (#185)

* refactor(ai): back assert_response_judged with a JudgeAgent

Replace RecordingAgent._judge_live's hand-rolled init_chat_model()/invoke()
call with JudgeAgent, an Agent subclass. The judge now goes through the same
provider/model resolution pipeline as any other agent, and is fakeable via
JudgeAgent.fake() and replayable via JudgeAgent.record() instead of only
being testable by mocking langchain directly.

assert_response_judged() is now async (the judge call is a real Agent.prompt()
under the hood) and accepts an optional provider kwarg, forwarded through to
JudgeAgent and folded into the cached verdict's cache key.

* refactor(ai): drop JudgeAgent's constructor, use plain Agent attributes

model/provider are never set via constructor args anywhere else in the
framework -- always class attributes or the @model()/@Provider() decorators.
Match that: JudgeAgent has no __init__ override, and _judge_live() sets
.model/.provider directly on the instance, same as any other Agent.

* refactor(ai): parse JudgeAgent verdicts via schema(), not hand-rolled JSON

Give JudgeAgent a Verdict pydantic schema() and put the grading rubric in
instructions() -- the JSON reply is now turned into a typed result through
the same structured-output path any Agent gets from schema()/response.parsed.
Drops the bespoke _build_prompt()/_parse_verdict() helpers.

* feat(ai): enforce schema() via with_structured_output on real model calls

When an Agent declares schema() and has no tools, build the model with
chat_model.with_structured_output(schema, include_raw=True) so the provider
enforces the shape, instead of relying on prompt instructions + post-hoc JSON
parsing. include_raw keeps the raw message so content/usage/cassettes still
work; the Runner passes the structured result through without executing the
synthetic tool call, and _to_agent_response unwraps it into response.parsed.

Streaming opts out (needs raw token chunks), tools take precedence over a
schema in a single call, and the fake/record paths are unchanged (they keep
exercising the JSON-string parse path for deterministic replay).

Also drops two stale docstrings from Ai.

* feat(ai): pass tools and schema in one payload; model picks per turn

When an agent declares both tools() and schema(), bind them together via
bind_tools([*tools, schema]) so a single model call offers both. The model
returns either a real tool call (which the Runner executes) or the schema as
its structured answer (which the Runner parses into response.parsed). Schema
without tools still uses with_structured_output() for enforcement.

_apply_schema no longer coerces content when the agent has tools, so a tool's
plain-text output is not force-parsed into the schema.

* refactor(ai): always bind schema as a tool, drop with_structured_output branch

schema() is just appended to tools() and the whole set is bound in one call.
This collapses the schema-only special case: with_structured_output only added
tool_choice="any" enforcement, which is redundant now that _apply_schema parses
plain JSON text for tool-less agents -- so a schema-only agent gets its parsed
result whether the model emits the schema tool call or replies in JSON text.

Also drops the now-dead structured-dict passthrough in Runner.run.

* refactor(ai): use with_structured_output for schema; leave tool binding as-is

Bind tools exactly as before, then wrap the model with
with_structured_output(schema, include_raw=True) when a schema is declared.
The wrapped model returns {raw, parsed, parsing_error}; the Runner passes it
through and _to_agent_response unwraps it into response.parsed.

Reverts the schema-as-a-tool detection and the _apply_schema tools guard.

Author: tmgbedu

example/agents/tests/units/agents/test_router_agent.py

@@ -1,16 +1,18 @@
 from langchain_core.messages import AIMessage, HumanMessage
 
 from app.agents.chat import RouterAgent
+from tests.test_case import TestCase
 
 
-class TestRouterAgent:
+class TestRouterAgent(TestCase):
     async def test_the_router_agent(self):
         with RouterAgent.record("record_stream.json") as agent:
             await agent.prompt("hello")
             agent.assert_text_response()
             agent.assert_tool_not_called(["job_search_tool"])
-            agent.assert_response_judged(
-                model="gpt-3.5-turbo",
+            await agent.assert_response_judged(
+                model="gemini-3.5-flash-lite",
+                provider="google",
                 expectation="The llm should respond with greetings",
             )
             agent.assert_response_time_lt(5)

fastapi_startkit/src/fastapi_startkit/ai/__init__.py

@@ -10,6 +10,7 @@
 from .image import Image, ImageResponse
 from .image_factory import ImageFactory
 from .ai import Ai
+from .judge import JudgeAgent
 from .providers.ai_provider import AIProvider
 from .response import AgentResponse, AgentSnapshot
 from .testing import AgentBinding, AgentModelFake, RecordingAgent, ToolCallView
@@ -25,6 +26,7 @@
     "AIConfig",
     "AIProvider",
     "AnthropicConfig",
+    "JudgeAgent",
     "RecordingAgent",
     "ToolCallView",
     "Audio",

fastapi_startkit/src/fastapi_startkit/ai/agent.py

@@ -188,27 +188,37 @@ def _build_messages(
 
         return messages
 
-    def _build_model(self, model: str | None = None, provider_options: dict | None = None) -> Any:
+    def _build_model(
+        self, model: str | None = None, provider_options: dict | None = None, structured: bool = True
+    ) -> Any:
         from .ai import Ai  # noqa: PLC0415
 
-        return Ai().get_model_for(self, model, provider_options)
+        return Ai().get_model_for(self, model, provider_options, structured)
 
     def _to_agent_response(self, result: Any) -> AgentResponse:
+        parsed = None
+        structured = isinstance(result, dict) and "parsed" in result and "raw" in result
+        if structured:
+            parsed = result.get("parsed")
+            result = result.get("raw")
+
         messages = result.get("messages", []) if isinstance(result, dict) else []
         final = messages[-1] if messages else result
 
         content = getattr(final, "content", "")
         if not isinstance(content, str):
             content = str(content)
+        if structured and not content and hasattr(parsed, "model_dump_json"):
+            content = parsed.model_dump_json()
 
-        tool_calls = list(getattr(final, "tool_calls", None) or [])
+        tool_calls = [] if structured else list(getattr(final, "tool_calls", None) or [])
 
         usage: dict[str, Any] = {}
         meta = getattr(final, "usage_metadata", None)
         if meta:
             usage = {"input": meta.get("input_tokens", 0), "output": meta.get("output_tokens", 0)}
 
-        return AgentResponse(content=content, tool_calls=tool_calls, usage=usage, raw=result)
+        return AgentResponse(content=content, tool_calls=tool_calls, usage=usage, raw=result, parsed=parsed)
 
     def _apply_schema(self, response: AgentResponse) -> AgentResponse:
         schema = self.schema()
@@ -253,7 +263,7 @@ async def _stream(
         from .runner import StreamRunner  # noqa: PLC0415
 
         messages = self._build_messages(message)
-        chat_model = self._build_model(model, provider_options)
+        chat_model = self._build_model(model, provider_options, structured=False)
         chain = list(self.middleware())
 
         def core(m: Any) -> Response:

fastapi_startkit/src/fastapi_startkit/ai/ai.py

@@ -21,11 +21,6 @@ def _key(agent: "Agent | str") -> str:
 
     @classmethod
     def fake(cls, agent: "Agent | str", messages: list) -> Any:
-        """Register a deterministic stand-in chat model for ``agent``.
-
-        Replays ``messages`` in order via a GenericFakeChatModel — no live
-        LLM call. Plain strings are coerced into ``AIMessage(content=...)``.
-        """
         from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
         from langchain_core.messages import AIMessage
 
@@ -51,14 +46,24 @@ def reset_fakes(cls) -> None:
         cls.fake_agent_models.clear()
         cls.fake_agent_responses.clear()
 
-    def get_model_for(self, agent: "Agent", model: str | None = None, provider_options: dict | None = None) -> Any:
-        """Resolve the model to run: a registered fake if one exists for
-        ``agent``, otherwise a freshly-built provider model."""
+    def get_model_for(
+        self,
+        agent: "Agent",
+        model: str | None = None,
+        provider_options: dict | None = None,
+        structured: bool = True,
+    ) -> Any:
         if self.has_fake_model_for(agent):
             return self.get_fake_model_for(agent)
-        return self.build(agent, model, provider_options)
-
-    def build(self, agent: "Agent", model: str | None = None, provider_options: dict | None = None) -> Any:
+        return self.build(agent, model, provider_options, structured)
+
+    def build(
+        self,
+        agent: "Agent",
+        model: str | None = None,
+        provider_options: dict | None = None,
+        structured: bool = True,
+    ) -> Any:
         from langchain.chat_models import init_chat_model  # noqa: PLC0415
 
         lab = Lab.get_provider(agent.provider)
@@ -79,7 +84,13 @@ def build(self, agent: "Agent", model: str | None = None, provider_options: dict
         chat_model = init_chat_model(self._resolve_model(agent, model), **kwargs)
 
         tools = list(agent.tools())
-        return chat_model.bind_tools(tools) if tools else chat_model
+        chat_model = chat_model.bind_tools(tools) if tools else chat_model
+
+        schema = agent.schema()
+        if structured and schema is not None:
+            chat_model = chat_model.with_structured_output(schema, include_raw=True)
+
+        return chat_model
 
     def _resolve_model(self, agent: "Agent", override: str | None = None) -> str:
         return Lab.get_provider(agent.provider).get_model(override or agent.model or None)

fastapi_startkit/src/fastapi_startkit/ai/judge.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+from .agent import Agent
+
+
+class Verdict(BaseModel):
+    passed: bool
+    reasoning: str = ""
+
+
+class JudgeAgent(Agent):
+    """Grades a response against a natural-language expectation.
+
+    A plain ``Agent`` whose ``schema()`` is a ``Verdict`` model, so the JSON
+    reply is parsed into a typed result through the standard structured-output
+    path — no hand-rolled verdict parsing. Set ``.model``/``.provider`` like
+    any other agent; it's fakeable via ``fake()`` and replayable via
+    ``record()`` for free.
+    """
+
+    def instructions(self) -> str:
+        return (
+            "You are grading whether an AI agent's response satisfies an expectation. "
+            'Reply with strict JSON only, no prose: {"passed": true|false, "reasoning": "<one sentence>"}'
+        )
+
+    def schema(self):
+        return Verdict
+
+    async def judge(self, expectation: str, content: str) -> dict:
+        response = await self.prompt(f"Expectation: {expectation}\n\nResponse to grade:\n{content}")
+        return response.parsed.model_dump()

fastapi_startkit/src/fastapi_startkit/ai/runner.py

@@ -27,6 +27,9 @@ async def run(self, messages: Sequence[Message]) -> BaseMessage:
         history: list[Message] = list(messages)
         response: AIMessage = await self.model.ainvoke(history)  # type: ignore[assignment]
 
+        if isinstance(response, dict) and "parsed" in response:
+            return response  # type: ignore[return-value]
+
         if not response.tool_calls:
             return response
 

fastapi_startkit/src/fastapi_startkit/ai/testing.py

@@ -5,7 +5,6 @@
 import hashlib
 import inspect
 import json
-import re
 import sys
 import time
 from collections.abc import AsyncIterator
@@ -244,49 +243,38 @@ def assert_response_time_lt(self, seconds: float) -> None:
         assert self.last_elapsed is not None, "No prompt() call has been made yet."
         assert self.last_elapsed < seconds, f"Expected response time < {seconds}s, took {self.last_elapsed:.3f}s"
 
-    def assert_response_judged(self, *, model: str, expectation: str) -> None:
+    async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None:
         response = self._require_response()
-        verdict = self._judge(model, expectation, response.content)
+        verdict = await self._judge(model, expectation, response.content, provider)
         assert verdict.get("passed"), (
             f"Judge ({model}) rejected the response for expectation {expectation!r}: "
             f"{verdict.get('reasoning', '')!r} — response was {response.content!r}"
         )
 
-    def _judge(self, model: str, expectation: str, content: str) -> dict:
+    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)
+        key = self._judge_key(model, expectation, content, provider)
         if key in store:
             return store[key]
-        verdict = self._judge_live(model, expectation, content)
+        verdict = await self._judge_live(model, expectation, content, provider)
         self._save(cassette, store, key, verdict)
         return verdict
 
     @staticmethod
-    def _judge_key(model: str, expectation: str, content: str) -> str:
+    def _judge_key(model: str, expectation: str, content: str, provider: str | None = None) -> str:
         payload = json.dumps(
-            {"judge_model": model, "expectation": expectation, "content": content},
+            {"judge_model": model, "judge_provider": provider, "expectation": expectation, "content": content},
             sort_keys=True,
         )
         return "judge:" + hashlib.sha256(payload.encode()).hexdigest()
 
-    def _judge_live(self, model: str, expectation: str, content: str) -> dict:
-        from langchain.chat_models import init_chat_model  # noqa: PLC0415
+    async def _judge_live(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict:
+        from .judge import JudgeAgent  # noqa: PLC0415
 
-        prompt = (
-            "You are grading whether an AI agent's response satisfies an expectation.\n"
-            f"Expectation: {expectation}\n"
-            f"Response: {content}\n\n"
-            'Reply with strict JSON only, no prose: {"passed": true|false, "reasoning": "<one sentence>"}'
-        )
-        chat_model = init_chat_model(model)
-        result = chat_model.invoke(prompt)
-        return self._parse_verdict(result.content)
-
-    @staticmethod
-    def _parse_verdict(raw: str) -> dict:
-        match = re.search(r"\{.*\}", raw, re.DOTALL)
-        data = json.loads(match.group(0) if match else raw)
-        return {"passed": bool(data.get("passed")), "reasoning": data.get("reasoning", "")}
+        judge = JudgeAgent()
+        judge.model = model
+        judge.provider = provider
+        return await judge.judge(expectation, content)
 
 
 class AgentBinding:

fastapi_startkit/tests/ai/test_agent_record_fluent.py

@@ -20,7 +20,6 @@
 import unittest
 from unittest import mock
 
-import langchain.chat_models as chat_models
 from langchain_core.messages import AIMessage, HumanMessage
 
 from fastapi_startkit.ai.agent import Agent
@@ -255,86 +254,79 @@ async def test_passes_when_judge_approves(self):
         self.setup_agent("Hello there, welcome!")
         with tempfile.TemporaryDirectory() as tmp:
             with mock.patch.object(
-                RecordingAgent, "_judge_live", return_value={"passed": True, "reasoning": "greets the user"}
+                RecordingAgent,
+                "_judge_live",
+                mock.AsyncMock(return_value={"passed": True, "reasoning": "greets the user"}),
             ):
                 with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent:
                     await agent.prompt("hello")
-                    agent.assert_response_judged(
+                    await agent.assert_response_judged(
                         model="gpt-3.5-turbo", expectation="The llm should respond with greetings"
                     )
 
     async def test_fails_when_judge_rejects(self):
         self.setup_agent("Completely unrelated content")
         with tempfile.TemporaryDirectory() as tmp:
             with mock.patch.object(
-                RecordingAgent, "_judge_live", return_value={"passed": False, "reasoning": "not a greeting"}
+                RecordingAgent,
+                "_judge_live",
+                mock.AsyncMock(return_value={"passed": False, "reasoning": "not a greeting"}),
             ):
                 with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent:
                     await agent.prompt("hello")
                     with self.assertRaises(AssertionError):
-                        agent.assert_response_judged(
+                        await agent.assert_response_judged(
                             model="gpt-3.5-turbo", expectation="The llm should respond with greetings"
                         )
 
     async def test_verdict_is_cached_in_the_cassette_and_not_re_judged(self):
         self.setup_agent("Hello there!")
-        judge = mock.Mock(return_value={"passed": True, "reasoning": "ok"})
+        judge = mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"})
         with tempfile.TemporaryDirectory() as tmp:
             cassette = os.path.join(tmp, "c.json")
             with mock.patch.object(RecordingAgent, "_judge_live", judge):
                 with SimpleAgent.record(cassette) as agent:
                     await agent.prompt("hello")
-                    agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
-                    agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
+                    await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
+                    await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
 
         judge.assert_called_once()
 
     async def test_verdict_persists_to_disk_for_a_later_replay(self):
         self.setup_agent("Hello there!")
         with tempfile.TemporaryDirectory() as tmp:
             cassette = os.path.join(tmp, "c.json")
-            with mock.patch.object(RecordingAgent, "_judge_live", return_value={"passed": True, "reasoning": "ok"}):
+            with mock.patch.object(
+                RecordingAgent, "_judge_live", mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"})
+            ):
                 with SimpleAgent.record(cassette) as agent:
                     await agent.prompt("hello")
-                    agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
+                    await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
 
-            judge = mock.Mock(side_effect=AssertionError("must not be called on replay"))
+            judge = mock.AsyncMock(side_effect=AssertionError("must not be called on replay"))
             with mock.patch.object(RecordingAgent, "_judge_live", judge):
                 with SimpleAgent.record(cassette) as agent:
                     await agent.prompt("hello")
-                    agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
+                    await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
 
             judge.assert_not_called()
 
     async def test_fails_when_no_prompt_has_been_made(self):
         with tempfile.TemporaryDirectory() as tmp:
             with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent:
                 with self.assertRaises(AssertionError):
-                    agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
-
-
-class TestJudgeLiveModelCall(unittest.TestCase):
-    def test_calls_init_chat_model_and_parses_json_verdict(self):
-        captured = {}
-
-        class FakeResult:
-            content = '{"passed": true, "reasoning": "Greets the user politely."}'
-
-        class FakeModel:
-            def invoke(self, prompt):
-                captured["prompt"] = prompt
-                return FakeResult()
-
-        patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: FakeModel())
-        patcher.start()
-        self.addCleanup(patcher.stop)
+                    await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet")
 
-        agent = RecordingAgent(SimpleAgent())
-        verdict = agent._judge_live("gpt-3.5-turbo", "The llm should respond with greetings", "Hello there!")
+    async def test_provider_is_forwarded_to_the_judge(self):
+        self.setup_agent("Hello there!")
+        judge = mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"})
+        with tempfile.TemporaryDirectory() as tmp:
+            with mock.patch.object(RecordingAgent, "_judge_live", judge):
+                with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent:
+                    await agent.prompt("hello")
+                    await agent.assert_response_judged(model="gpt-3.5-turbo", provider="openai", expectation="greet")
 
-        self.assertTrue(verdict["passed"])
-        self.assertIn("Greets", verdict["reasoning"])
-        self.assertIn("Hello there!", captured["prompt"])
+        judge.assert_called_once_with("gpt-3.5-turbo", "greet", "Hello there!", "openai")
 
 
 class TestExistingRecordApiIsUnaffected(unittest.IsolatedAsyncioTestCase):