From a3ed98bbb96ec5be39735c81ac1655adf287f368 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 21:30:37 -0700 Subject: [PATCH 1/7] 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. --- .../tests/units/agents/test_router_agent.py | 8 ++- .../src/fastapi_startkit/ai/__init__.py | 2 + .../src/fastapi_startkit/ai/judge.py | 42 +++++++++++ .../src/fastapi_startkit/ai/testing.py | 35 +++------ .../tests/ai/test_agent_record_fluent.py | 62 +++++++--------- fastapi_startkit/tests/ai/test_judge_agent.py | 72 +++++++++++++++++++ 6 files changed, 158 insertions(+), 63 deletions(-) create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/judge.py create mode 100644 fastapi_startkit/tests/ai/test_judge_agent.py diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py index 26ddfae7..17a71d18 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/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) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index 1e8f8b8c..47abea1b 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/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", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/judge.py b/fastapi_startkit/src/fastapi_startkit/ai/judge.py new file mode 100644 index 00000000..6bc19589 --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/judge.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import json +import re + +from .agent import Agent + + +class JudgeAgent(Agent): + """Grades a response against a natural-language expectation. + + Just an ``Agent`` — the verdict call goes through the same model + resolution, provider handling, and pipeline as any other agent, so it's + fakeable via ``JudgeAgent.fake()`` and replayable via + ``JudgeAgent.record()`` instead of hand-rolling a separate langchain call. + """ + + def __init__(self, model: str | None = None, provider: str | None = None) -> None: + super().__init__() + if model is not None: + self.model = model + if provider is not None: + self.provider = provider + + async def judge(self, expectation: str, content: str) -> dict: + response = await self.prompt(self._build_prompt(expectation, content)) + return self._parse_verdict(response.content) + + @staticmethod + def _build_prompt(expectation: str, content: str) -> str: + return ( + "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": ""}' + ) + + @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", "")} diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 639fda71..e97c9812 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/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,35 @@ 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": ""}' - ) - 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", "")} + return await JudgeAgent(model=model, provider=provider).judge(expectation, content) class AgentBinding: diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py index 30d12989..0f2c42b2 100644 --- a/fastapi_startkit/tests/ai/test_agent_record_fluent.py +++ b/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,11 +254,13 @@ 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" ) @@ -267,25 +268,27 @@ 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() @@ -293,16 +296,18 @@ 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() @@ -310,31 +315,18 @@ 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): diff --git a/fastapi_startkit/tests/ai/test_judge_agent.py b/fastapi_startkit/tests/ai/test_judge_agent.py new file mode 100644 index 00000000..9a06f20c --- /dev/null +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -0,0 +1,72 @@ +"""Tests for JudgeAgent — the built-in Agent subclass that grades a +response against a natural-language expectation. + +Because it's just an Agent, it gets model/provider resolution, ``fake()``, +and ``record()`` for free instead of RecordingAgent hand-rolling its own +langchain call. +""" + +import json +import os +import tempfile +import unittest +from unittest import mock + +from fastapi_startkit.ai.judge import JudgeAgent +from fastapi_startkit.ai.response import AgentResponse + + +class TestJudgeAgent(unittest.IsolatedAsyncioTestCase): + async def test_judge_parses_a_passing_verdict_from_the_model_response(self): + verdict = '{"passed": true, "reasoning": "Greets the user politely."}' + with JudgeAgent.fake([verdict]): + judge = JudgeAgent(model="gpt-3.5-turbo") + result = await judge.judge("The llm should respond with greetings", "Hello there!") + + self.assertEqual(result, {"passed": True, "reasoning": "Greets the user politely."}) + + async def test_judge_parses_a_failing_verdict(self): + verdict = '{"passed": false, "reasoning": "Not a greeting."}' + with JudgeAgent.fake([verdict]): + judge = JudgeAgent(model="gpt-3.5-turbo") + result = await judge.judge("The llm should respond with greetings", "Completely unrelated") + + self.assertFalse(result["passed"]) + + async def test_judge_tolerates_prose_around_the_json_block(self): + verdict = 'Sure! Here is the verdict:\n{"passed": true, "reasoning": "ok"}\nHope that helps.' + with JudgeAgent.fake([verdict]): + judge = JudgeAgent(model="gpt-3.5-turbo") + result = await judge.judge("greet", "Hello!") + + self.assertTrue(result["passed"]) + + def test_prompt_includes_expectation_and_response(self): + prompt = JudgeAgent._build_prompt("The llm should respond with greetings", "Hello there!") + + self.assertIn("The llm should respond with greetings", prompt) + self.assertIn("Hello there!", prompt) + + async def test_model_and_provider_are_set_from_constructor(self): + judge = JudgeAgent(model="gpt-4o-mini", provider="openai") + self.assertEqual(judge.model, "gpt-4o-mini") + self.assertEqual(judge.provider, "openai") + + async def test_constructor_args_are_optional(self): + judge = JudgeAgent() + self.assertIsNone(judge.model) + self.assertIsNone(judge.provider) + + async def test_judge_is_usable_via_the_record_fluent_dsl(self): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content='{"passed": true, "reasoning": "ok"}') + + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "judge.json") + with mock.patch.object(JudgeAgent, "_run", fake_run): + with JudgeAgent.record(cassette) as agent: + response = await agent.prompt("grade this") + + self.assertIn('"passed"', response.content) + self.assertTrue(os.path.exists(cassette)) + self.assertEqual(len(json.loads(open(cassette).read())), 1) From c3114bef86fcd3713c6cee1abff776a159141080 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 21:46:50 -0700 Subject: [PATCH 2/7] 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. --- .../src/fastapi_startkit/ai/judge.py | 16 +++++---------- .../src/fastapi_startkit/ai/testing.py | 5 ++++- fastapi_startkit/tests/ai/test_judge_agent.py | 20 +++++++++++++------ 3 files changed, 23 insertions(+), 18 deletions(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/judge.py b/fastapi_startkit/src/fastapi_startkit/ai/judge.py index 6bc19589..3d320227 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/judge.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/judge.py @@ -9,19 +9,13 @@ class JudgeAgent(Agent): """Grades a response against a natural-language expectation. - Just an ``Agent`` — the verdict call goes through the same model - resolution, provider handling, and pipeline as any other agent, so it's - fakeable via ``JudgeAgent.fake()`` and replayable via - ``JudgeAgent.record()`` instead of hand-rolling a separate langchain call. + Just an ``Agent`` — set ``.model``/``.provider`` like any other agent + and the verdict call goes through the same model resolution, provider + handling, and pipeline as any other agent, so it's fakeable via + ``JudgeAgent.fake()`` and replayable via ``JudgeAgent.record()`` instead + of hand-rolling a separate langchain call. """ - def __init__(self, model: str | None = None, provider: str | None = None) -> None: - super().__init__() - if model is not None: - self.model = model - if provider is not None: - self.provider = provider - async def judge(self, expectation: str, content: str) -> dict: response = await self.prompt(self._build_prompt(expectation, content)) return self._parse_verdict(response.content) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index e97c9812..c2fed401 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -271,7 +271,10 @@ def _judge_key(model: str, expectation: str, content: str, provider: str | None async def _judge_live(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: from .judge import JudgeAgent # noqa: PLC0415 - return await JudgeAgent(model=model, provider=provider).judge(expectation, content) + judge = JudgeAgent() + judge.model = model + judge.provider = provider + return await judge.judge(expectation, content) class AgentBinding: diff --git a/fastapi_startkit/tests/ai/test_judge_agent.py b/fastapi_startkit/tests/ai/test_judge_agent.py index 9a06f20c..f4f21923 100644 --- a/fastapi_startkit/tests/ai/test_judge_agent.py +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -20,7 +20,8 @@ class TestJudgeAgent(unittest.IsolatedAsyncioTestCase): async def test_judge_parses_a_passing_verdict_from_the_model_response(self): verdict = '{"passed": true, "reasoning": "Greets the user politely."}' with JudgeAgent.fake([verdict]): - judge = JudgeAgent(model="gpt-3.5-turbo") + judge = JudgeAgent() + judge.model = "gpt-3.5-turbo" result = await judge.judge("The llm should respond with greetings", "Hello there!") self.assertEqual(result, {"passed": True, "reasoning": "Greets the user politely."}) @@ -28,7 +29,8 @@ async def test_judge_parses_a_passing_verdict_from_the_model_response(self): async def test_judge_parses_a_failing_verdict(self): verdict = '{"passed": false, "reasoning": "Not a greeting."}' with JudgeAgent.fake([verdict]): - judge = JudgeAgent(model="gpt-3.5-turbo") + judge = JudgeAgent() + judge.model = "gpt-3.5-turbo" result = await judge.judge("The llm should respond with greetings", "Completely unrelated") self.assertFalse(result["passed"]) @@ -36,7 +38,8 @@ async def test_judge_parses_a_failing_verdict(self): async def test_judge_tolerates_prose_around_the_json_block(self): verdict = 'Sure! Here is the verdict:\n{"passed": true, "reasoning": "ok"}\nHope that helps.' with JudgeAgent.fake([verdict]): - judge = JudgeAgent(model="gpt-3.5-turbo") + judge = JudgeAgent() + judge.model = "gpt-3.5-turbo" result = await judge.judge("greet", "Hello!") self.assertTrue(result["passed"]) @@ -47,13 +50,18 @@ def test_prompt_includes_expectation_and_response(self): self.assertIn("The llm should respond with greetings", prompt) self.assertIn("Hello there!", prompt) - async def test_model_and_provider_are_set_from_constructor(self): - judge = JudgeAgent(model="gpt-4o-mini", provider="openai") + def test_model_and_provider_are_plain_agent_attributes(self): + """No custom constructor — set like any other Agent's model/provider.""" + judge = JudgeAgent() + judge.model = "gpt-4o-mini" + judge.provider = "openai" + self.assertEqual(judge.model, "gpt-4o-mini") self.assertEqual(judge.provider, "openai") - async def test_constructor_args_are_optional(self): + def test_model_and_provider_default_to_agent_defaults(self): judge = JudgeAgent() + self.assertIsNone(judge.model) self.assertIsNone(judge.provider) From 6d787553fe0ced91f6c23fdc25c5c2021e4f825b Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 22:28:15 -0700 Subject: [PATCH 3/7] 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. --- .../src/fastapi_startkit/ai/judge.py | 40 ++++++------ fastapi_startkit/tests/ai/test_judge_agent.py | 61 +++++++++---------- 2 files changed, 49 insertions(+), 52 deletions(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/judge.py b/fastapi_startkit/src/fastapi_startkit/ai/judge.py index 3d320227..6923bd3d 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/judge.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/judge.py @@ -1,36 +1,34 @@ from __future__ import annotations -import json -import re +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. - Just an ``Agent`` — set ``.model``/``.provider`` like any other agent - and the verdict call goes through the same model resolution, provider - handling, and pipeline as any other agent, so it's fakeable via - ``JudgeAgent.fake()`` and replayable via ``JudgeAgent.record()`` instead - of hand-rolling a separate langchain call. + 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. """ - async def judge(self, expectation: str, content: str) -> dict: - response = await self.prompt(self._build_prompt(expectation, content)) - return self._parse_verdict(response.content) - - @staticmethod - def _build_prompt(expectation: str, content: str) -> str: + def instructions(self) -> str: return ( - "You are grading whether an AI agent's response satisfies an expectation.\n" - f"Expectation: {expectation}\n" - f"Response: {content}\n\n" + "You are grading whether an AI agent's response satisfies an expectation. " 'Reply with strict JSON only, no prose: {"passed": true|false, "reasoning": ""}' ) - @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", "")} + 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() diff --git a/fastapi_startkit/tests/ai/test_judge_agent.py b/fastapi_startkit/tests/ai/test_judge_agent.py index f4f21923..e1d5a964 100644 --- a/fastapi_startkit/tests/ai/test_judge_agent.py +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -1,54 +1,54 @@ -"""Tests for JudgeAgent — the built-in Agent subclass that grades a -response against a natural-language expectation. +"""Tests for JudgeAgent — grades a response against an expectation. -Because it's just an Agent, it gets model/provider resolution, ``fake()``, -and ``record()`` for free instead of RecordingAgent hand-rolling its own -langchain call. +It's a plain Agent whose ``schema()`` is a ``Verdict`` model, so the model's +JSON reply is turned into a typed result through the standard structured-output +path (``response.parsed``) — no hand-rolled verdict parsing. """ -import json import os import tempfile import unittest from unittest import mock -from fastapi_startkit.ai.judge import JudgeAgent +from langchain_core.messages import AIMessage + +from fastapi_startkit.ai.judge import JudgeAgent, Verdict from fastapi_startkit.ai.response import AgentResponse class TestJudgeAgent(unittest.IsolatedAsyncioTestCase): - async def test_judge_parses_a_passing_verdict_from_the_model_response(self): - verdict = '{"passed": true, "reasoning": "Greets the user politely."}' - with JudgeAgent.fake([verdict]): - judge = JudgeAgent() - judge.model = "gpt-3.5-turbo" - result = await judge.judge("The llm should respond with greetings", "Hello there!") + async def test_judge_returns_a_passing_verdict_dict(self): + with JudgeAgent.fake(['{"passed": true, "reasoning": "Greets the user politely."}']): + result = await JudgeAgent().judge("The llm should respond with greetings", "Hello there!") self.assertEqual(result, {"passed": True, "reasoning": "Greets the user politely."}) - async def test_judge_parses_a_failing_verdict(self): - verdict = '{"passed": false, "reasoning": "Not a greeting."}' - with JudgeAgent.fake([verdict]): - judge = JudgeAgent() - judge.model = "gpt-3.5-turbo" - result = await judge.judge("The llm should respond with greetings", "Completely unrelated") + async def test_judge_returns_a_failing_verdict(self): + with JudgeAgent.fake(['{"passed": false, "reasoning": "Not a greeting."}']): + result = await JudgeAgent().judge("greet", "Completely unrelated") self.assertFalse(result["passed"]) - async def test_judge_tolerates_prose_around_the_json_block(self): - verdict = 'Sure! Here is the verdict:\n{"passed": true, "reasoning": "ok"}\nHope that helps.' - with JudgeAgent.fake([verdict]): - judge = JudgeAgent() - judge.model = "gpt-3.5-turbo" - result = await judge.judge("greet", "Hello!") + def test_schema_is_the_verdict_model(self): + self.assertIs(JudgeAgent().schema(), Verdict) + + async def test_judge_feeds_expectation_and_response_to_the_model(self): + seen = {} + + class Capturing: + def bind_tools(self, tools, **kwargs): + return self - self.assertTrue(result["passed"]) + async def ainvoke(self, messages): + seen["messages"] = messages + return AIMessage(content='{"passed": true, "reasoning": "ok"}') - def test_prompt_includes_expectation_and_response(self): - prompt = JudgeAgent._build_prompt("The llm should respond with greetings", "Hello there!") + with mock.patch.object(JudgeAgent, "_build_model", lambda self, *a, **k: Capturing()): + await JudgeAgent().judge("The llm should respond with greetings", "Hello there!") - self.assertIn("The llm should respond with greetings", prompt) - self.assertIn("Hello there!", prompt) + blob = " ".join(str(getattr(m, "content", m)) for m in seen["messages"]) + self.assertIn("The llm should respond with greetings", blob) + self.assertIn("Hello there!", blob) def test_model_and_provider_are_plain_agent_attributes(self): """No custom constructor — set like any other Agent's model/provider.""" @@ -77,4 +77,3 @@ async def fake_run(agent_self, message, **kwargs): self.assertIn('"passed"', response.content) self.assertTrue(os.path.exists(cassette)) - self.assertEqual(len(json.loads(open(cassette).read())), 1) From e53240e2a1d6a773a29d92ccc388a668b56547ab Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 22:41:48 -0700 Subject: [PATCH 4/7] 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. --- .../src/fastapi_startkit/ai/agent.py | 20 +- .../src/fastapi_startkit/ai/ai.py | 36 ++-- .../src/fastapi_startkit/ai/runner.py | 3 + .../tests/ai/test_structured_output.py | 184 ++++++++++++++++++ 4 files changed, 226 insertions(+), 17 deletions(-) create mode 100644 fastapi_startkit/tests/ai/test_structured_output.py diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 96926185..1931cd60 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/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: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index da1dc599..327d0804 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/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,14 @@ 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 + if tools: + return chat_model.bind_tools(tools) + + schema = agent.schema() + if structured and schema is not None: + return 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) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index 49fe5a44..ba6ded73 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/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 diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py new file mode 100644 index 00000000..8e2e228c --- /dev/null +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -0,0 +1,184 @@ +"""Structured output: when an Agent declares a schema() and has no tools, the +model is built with with_structured_output() so the provider enforces the +shape, rather than relying on prompt instructions + post-hoc JSON parsing. +""" + +import unittest +from unittest import mock + +import langchain.chat_models as chat_models +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from pydantic import BaseModel + +from fastapi_startkit.ai import AIConfig +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.ai import Ai +from fastapi_startkit.ai.runner import Runner +from fastapi_startkit.application import app + + +class Movie(BaseModel): + title: str + year: int + + +class MovieAgent(Agent): + def schema(self): + return Movie + + +@tool +def noop(query: str) -> str: + """A no-op tool.""" + return query + + +class ToolMovieAgent(Agent): + def schema(self): + return Movie + + def tools(self): + return [noop] + + +def _structured_payload(parsed: Movie) -> dict: + raw = AIMessage(content="", tool_calls=[{"name": "Movie", "args": {}, "id": "1", "type": "tool_call"}]) + return {"raw": raw, "parsed": parsed, "parsing_error": None} + + +class TestBuildStructuredOutput(unittest.TestCase): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + + def tearDown(self): + Ai.reset_fakes() + + def _patch_init(self, fake): + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: fake) + patcher.start() + self.addCleanup(patcher.stop) + + def test_wraps_with_structured_output_when_schema_and_no_tools(self): + captured = {} + + class FakeModel: + def with_structured_output(self, schema, **kwargs): + captured["schema"] = schema + captured["kwargs"] = kwargs + return "STRUCTURED" + + def bind_tools(self, tools, **kwargs): + return "BOUND" + + self._patch_init(FakeModel()) + + result = Ai().build(MovieAgent()) + + self.assertEqual(result, "STRUCTURED") + self.assertIs(captured["schema"], Movie) + self.assertEqual(captured["kwargs"], {"include_raw": True}) + + def test_tools_take_precedence_over_structured_output(self): + class FakeModel: + def with_structured_output(self, schema, **kwargs): + return "STRUCTURED" + + def bind_tools(self, tools, **kwargs): + return "BOUND" + + self._patch_init(FakeModel()) + + self.assertEqual(Ai().build(ToolMovieAgent()), "BOUND") + + def test_structured_false_returns_the_plain_model(self): + class FakeModel: + def with_structured_output(self, schema, **kwargs): + return "STRUCTURED" + + def bind_tools(self, tools, **kwargs): + return "BOUND" + + fake = FakeModel() + self._patch_init(fake) + + self.assertIs(Ai().build(MovieAgent(), structured=False), fake) + + def test_no_schema_no_tools_returns_the_plain_model(self): + class FakeModel: + def with_structured_output(self, schema, **kwargs): + return "STRUCTURED" + + def bind_tools(self, tools, **kwargs): + return "BOUND" + + fake = FakeModel() + self._patch_init(fake) + + self.assertIs(Ai().build(Agent()), fake) + + +class TestStructuredResponseMapping(unittest.TestCase): + def test_unwraps_include_raw_and_sets_parsed(self): + parsed = Movie(title="Inception", year=2010) + + response = MovieAgent()._to_agent_response(_structured_payload(parsed)) + + self.assertIs(response.parsed, parsed) + self.assertEqual(response.content, parsed.model_dump_json()) + + def test_suppresses_the_synthetic_structured_output_tool_call(self): + parsed = Movie(title="Inception", year=2010) + + response = MovieAgent()._to_agent_response(_structured_payload(parsed)) + + self.assertEqual(response.tool_calls, []) + + +class TestRunnerStructuredOutput(unittest.IsolatedAsyncioTestCase): + async def test_runner_passes_structured_dict_through_without_running_tools(self): + parsed = Movie(title="Inception", year=2010) + payload = _structured_payload(parsed) + + class Model: + async def ainvoke(self, messages): + return payload + + result = await Runner(MovieAgent(), Model()).run(["hi"]) + + self.assertEqual(result, payload) + + +class TestPromptEndToEnd(unittest.IsolatedAsyncioTestCase): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + + def tearDown(self): + Ai.reset_fakes() + + async def test_prompt_populates_parsed_via_structured_output(self): + parsed = Movie(title="Inception", year=2010) + + class Structured: + async def ainvoke(self, messages): + return {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + + class FakeModel: + def with_structured_output(self, schema, **kwargs): + return Structured() + + def bind_tools(self, tools, **kwargs): + return self + + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: FakeModel()) + patcher.start() + self.addCleanup(patcher.stop) + + response = await MovieAgent().prompt("best nolan movie") + + self.assertEqual(response.parsed, parsed) + self.assertEqual(response.content, parsed.model_dump_json()) From ea9e586bde297cff0f8e2512cdfe8d484f0a3d46 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 00:46:34 -0700 Subject: [PATCH 5/7] 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. --- .../src/fastapi_startkit/ai/agent.py | 2 +- .../src/fastapi_startkit/ai/ai.py | 9 +- .../src/fastapi_startkit/ai/runner.py | 24 ++- .../tests/ai/test_structured_output.py | 188 ++++++++++++------ 4 files changed, 151 insertions(+), 72 deletions(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 1931cd60..5db24725 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -222,7 +222,7 @@ def _to_agent_response(self, result: Any) -> AgentResponse: def _apply_schema(self, response: AgentResponse) -> AgentResponse: schema = self.schema() - if schema is not None and response.parsed is None and response.content: + if schema is not None and not self.tools() and response.parsed is None and response.content: response.parsed = self._build_schema(schema, response.content) return response diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index 327d0804..d6e6f8e7 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -84,13 +84,16 @@ def build( chat_model = init_chat_model(self._resolve_model(agent, model), **kwargs) tools = list(agent.tools()) - if tools: - return chat_model.bind_tools(tools) - schema = agent.schema() + if structured and schema is not None: + if tools: + return chat_model.bind_tools([*tools, schema]) return chat_model.with_structured_output(schema, include_raw=True) + if tools: + return chat_model.bind_tools(tools) + return chat_model def _resolve_model(self, agent: "Agent", override: str | None = None) -> str: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index ba6ded73..0e1dc4ed 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -22,6 +22,8 @@ def __init__(self, agent: Agent, model: Runnable[Any, BaseMessage]) -> None: self._tools: dict[str, BaseTool] = {tool.name: tool for tool in agent.tools()} self.model: Runnable[Any, BaseMessage] = model self.max_steps = agent.max_steps + self._schema = agent.schema() + self._schema_tool = self._schema_tool_name(self._schema) if self._schema is not None else None async def run(self, messages: Sequence[Message]) -> BaseMessage: history: list[Message] = list(messages) @@ -30,10 +32,28 @@ async def run(self, messages: Sequence[Message]) -> BaseMessage: if isinstance(response, dict) and "parsed" in response: return response # type: ignore[return-value] - if not response.tool_calls: + tool_calls = list(getattr(response, "tool_calls", None) or []) + if not tool_calls: return response - return (await self._run_tools(response.tool_calls))[-1] + for call in tool_calls: + if call.get("name") == self._schema_tool: + parsed = self._parse_schema(call.get("args") or {}) + return {"raw": response, "parsed": parsed, "parsing_error": None} # type: ignore[return-value] + + return (await self._run_tools(tool_calls))[-1] + + def _parse_schema(self, args: dict[str, Any]) -> Any: + schema = self._schema + if hasattr(schema, "model_validate"): + return schema.model_validate(args) + return schema(**args) + + @staticmethod + def _schema_tool_name(schema: Any) -> str: + from langchain_core.utils.function_calling import convert_to_openai_tool # noqa: PLC0415 + + return convert_to_openai_tool(schema)["function"]["name"] async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: return [await self._resolve_tool(call["name"]).ainvoke(call) for call in tool_calls] diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py index 8e2e228c..2f9a0192 100644 --- a/fastapi_startkit/tests/ai/test_structured_output.py +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -1,6 +1,15 @@ -"""Structured output: when an Agent declares a schema() and has no tools, the -model is built with with_structured_output() so the provider enforces the -shape, rather than relying on prompt instructions + post-hoc JSON parsing. +"""Structured output. + +An Agent's schema() and tools() are passed to the model in a single payload: + +- schema + tools -> bind_tools([*tools, schema]); the model picks a real tool + call OR the schema as its structured answer. If it picks the schema, we parse + and return the structured result; otherwise we run the tool it asked for. +- schema only -> with_structured_output(); the shape is enforced. +- tools only / neither -> unchanged. + +The fake/record paths bypass build(), so they keep parsing the JSON-string +content via schema() for deterministic replay. """ import unittest @@ -30,7 +39,7 @@ def schema(self): @tool def noop(query: str) -> str: - """A no-op tool.""" + """A no-op tool that echoes its query.""" return query @@ -42,12 +51,15 @@ def tools(self): return [noop] -def _structured_payload(parsed: Movie) -> dict: - raw = AIMessage(content="", tool_calls=[{"name": "Movie", "args": {}, "id": "1", "type": "tool_call"}]) - return {"raw": raw, "parsed": parsed, "parsing_error": None} +def _schema_tool_call(**args) -> dict: + return {"name": "Movie", "args": args, "id": "1", "type": "tool_call"} -class TestBuildStructuredOutput(unittest.TestCase): +def _real_tool_call(**args) -> dict: + return {"name": "noop", "args": args, "id": "2", "type": "tool_call"} + + +class TestBuild(unittest.TestCase): def setUp(self): container = app() container.bind("ai", AIConfig()) @@ -61,86 +73,88 @@ def _patch_init(self, fake): patcher.start() self.addCleanup(patcher.stop) - def test_wraps_with_structured_output_when_schema_and_no_tools(self): - captured = {} - + def _fake_model(self): class FakeModel: - def with_structured_output(self, schema, **kwargs): - captured["schema"] = schema - captured["kwargs"] = kwargs - return "STRUCTURED" + bound = None + structured = None def bind_tools(self, tools, **kwargs): + self.bound = list(tools) return "BOUND" - self._patch_init(FakeModel()) + def with_structured_output(self, schema, **kwargs): + self.structured = (schema, kwargs) + return "STRUCTURED" - result = Ai().build(MovieAgent()) + fake = FakeModel() + self._patch_init(fake) + return fake - self.assertEqual(result, "STRUCTURED") - self.assertIs(captured["schema"], Movie) - self.assertEqual(captured["kwargs"], {"include_raw": True}) + def test_schema_and_tools_are_bound_together_in_one_payload(self): + fake = self._fake_model() - def test_tools_take_precedence_over_structured_output(self): - class FakeModel: - def with_structured_output(self, schema, **kwargs): - return "STRUCTURED" + result = Ai().build(ToolMovieAgent()) - def bind_tools(self, tools, **kwargs): - return "BOUND" + self.assertEqual(result, "BOUND") + self.assertIn(noop, fake.bound) + self.assertIn(Movie, fake.bound) - self._patch_init(FakeModel()) + def test_schema_only_uses_enforced_structured_output(self): + fake = self._fake_model() - self.assertEqual(Ai().build(ToolMovieAgent()), "BOUND") + result = Ai().build(MovieAgent()) - def test_structured_false_returns_the_plain_model(self): - class FakeModel: - def with_structured_output(self, schema, **kwargs): - return "STRUCTURED" + self.assertEqual(result, "STRUCTURED") + self.assertEqual(fake.structured, (Movie, {"include_raw": True})) - def bind_tools(self, tools, **kwargs): - return "BOUND" + def test_tools_only_binds_just_the_tools(self): + fake = self._fake_model() - fake = FakeModel() - self._patch_init(fake) + class ToolAgent(Agent): + def tools(self): + return [noop] - self.assertIs(Ai().build(MovieAgent(), structured=False), fake) + result = Ai().build(ToolAgent()) - def test_no_schema_no_tools_returns_the_plain_model(self): - class FakeModel: - def with_structured_output(self, schema, **kwargs): - return "STRUCTURED" + self.assertEqual(result, "BOUND") + self.assertEqual(fake.bound, [noop]) - def bind_tools(self, tools, **kwargs): - return "BOUND" + def test_streaming_drops_the_schema_from_the_payload(self): + fake = self._fake_model() - fake = FakeModel() - self._patch_init(fake) + result = Ai().build(ToolMovieAgent(), structured=False) - self.assertIs(Ai().build(Agent()), fake) + self.assertEqual(result, "BOUND") + self.assertEqual(fake.bound, [noop]) + def test_no_schema_no_tools_returns_the_plain_model(self): + fake = self._fake_model() -class TestStructuredResponseMapping(unittest.TestCase): - def test_unwraps_include_raw_and_sets_parsed(self): - parsed = Movie(title="Inception", year=2010) + self.assertIs(Ai().build(Agent()), fake) - response = MovieAgent()._to_agent_response(_structured_payload(parsed)) - self.assertIs(response.parsed, parsed) - self.assertEqual(response.content, parsed.model_dump_json()) +class TestRunner(unittest.IsolatedAsyncioTestCase): + async def test_returns_structured_when_model_calls_the_schema(self): + class Model: + async def ainvoke(self, messages): + return AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) - def test_suppresses_the_synthetic_structured_output_tool_call(self): - parsed = Movie(title="Inception", year=2010) + result = await Runner(ToolMovieAgent(), Model()).run(["hi"]) - response = MovieAgent()._to_agent_response(_structured_payload(parsed)) + self.assertEqual(result["parsed"], Movie(title="Inception", year=2010)) - self.assertEqual(response.tool_calls, []) + async def test_runs_the_tool_when_model_calls_a_real_tool(self): + class Model: + async def ainvoke(self, messages): + return AIMessage(content="", tool_calls=[_real_tool_call(query="hello")]) + result = await Runner(ToolMovieAgent(), Model()).run(["hi"]) -class TestRunnerStructuredOutput(unittest.IsolatedAsyncioTestCase): - async def test_runner_passes_structured_dict_through_without_running_tools(self): + self.assertEqual(result.content, "hello") + + async def test_passes_structured_output_dict_through(self): parsed = Movie(title="Inception", year=2010) - payload = _structured_payload(parsed) + payload = {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} class Model: async def ainvoke(self, messages): @@ -151,6 +165,18 @@ async def ainvoke(self, messages): self.assertEqual(result, payload) +class TestResponseMapping(unittest.TestCase): + def test_unwraps_include_raw_into_parsed_and_content(self): + parsed = Movie(title="Inception", year=2010) + raw = AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) + + response = MovieAgent()._to_agent_response({"raw": raw, "parsed": parsed, "parsing_error": None}) + + self.assertIs(response.parsed, parsed) + self.assertEqual(response.content, parsed.model_dump_json()) + self.assertEqual(response.tool_calls, []) + + class TestPromptEndToEnd(unittest.IsolatedAsyncioTestCase): def setUp(self): container = app() @@ -160,7 +186,12 @@ def setUp(self): def tearDown(self): Ai.reset_fakes() - async def test_prompt_populates_parsed_via_structured_output(self): + def _patch(self, model): + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: model) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_schema_only_populates_parsed(self): parsed = Movie(title="Inception", year=2010) class Structured: @@ -171,14 +202,39 @@ class FakeModel: def with_structured_output(self, schema, **kwargs): return Structured() - def bind_tools(self, tools, **kwargs): - return self - - patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: FakeModel()) - patcher.start() - self.addCleanup(patcher.stop) + self._patch(FakeModel()) response = await MovieAgent().prompt("best nolan movie") self.assertEqual(response.parsed, parsed) self.assertEqual(response.content, parsed.model_dump_json()) + + async def test_schema_plus_tools_returns_structured_when_model_chooses_it(self): + class FakeModel: + def bind_tools(self, tools, **kwargs): + return self + + async def ainvoke(self, messages): + return AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) + + self._patch(FakeModel()) + + response = await ToolMovieAgent().prompt("best nolan movie") + + self.assertEqual(response.parsed, Movie(title="Inception", year=2010)) + self.assertEqual(response.tool_calls, []) + + async def test_schema_plus_tools_runs_the_tool_when_model_chooses_it(self): + class FakeModel: + def bind_tools(self, tools, **kwargs): + return self + + async def ainvoke(self, messages): + return AIMessage(content="", tool_calls=[_real_tool_call(query="hello")]) + + self._patch(FakeModel()) + + response = await ToolMovieAgent().prompt("run the tool") + + self.assertIsNone(response.parsed) + self.assertEqual(response.content, "hello") From a67bd4ac3dd1df7e0769726223a0cd6510fa6961 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 00:56:51 -0700 Subject: [PATCH 6/7] 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. --- .../src/fastapi_startkit/ai/ai.py | 10 +-- .../src/fastapi_startkit/ai/runner.py | 3 - .../tests/ai/test_structured_output.py | 66 ++++++++----------- 3 files changed, 31 insertions(+), 48 deletions(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index d6e6f8e7..c331f338 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -85,16 +85,10 @@ def build( tools = list(agent.tools()) schema = agent.schema() - if structured and schema is not None: - if tools: - return chat_model.bind_tools([*tools, schema]) - return chat_model.with_structured_output(schema, include_raw=True) - - if tools: - return chat_model.bind_tools(tools) + tools.append(schema) - return chat_model + return chat_model.bind_tools(tools) if tools else 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) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index 0e1dc4ed..293814aa 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -29,9 +29,6 @@ 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] - tool_calls = list(getattr(response, "tool_calls", None) or []) if not tool_calls: return response diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py index 2f9a0192..12d95656 100644 --- a/fastapi_startkit/tests/ai/test_structured_output.py +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -1,12 +1,12 @@ """Structured output. -An Agent's schema() and tools() are passed to the model in a single payload: - -- schema + tools -> bind_tools([*tools, schema]); the model picks a real tool - call OR the schema as its structured answer. If it picks the schema, we parse - and return the structured result; otherwise we run the tool it asked for. -- schema only -> with_structured_output(); the shape is enforced. -- tools only / neither -> unchanged. +An Agent's schema() is appended to its tools() and the whole set is bound to +the model in a single payload via bind_tools([*tools, schema]). The model picks +per turn: a real tool call (which the Runner executes) or the schema as its +structured answer (which the Runner parses into response.parsed). If a +schema-only agent replies with plain JSON text instead of the tool call, +_apply_schema still parses it, so the structured result comes through either +way. The fake/record paths bypass build(), so they keep parsing the JSON-string content via schema() for deterministic replay. @@ -76,16 +76,11 @@ def _patch_init(self, fake): def _fake_model(self): class FakeModel: bound = None - structured = None def bind_tools(self, tools, **kwargs): self.bound = list(tools) return "BOUND" - def with_structured_output(self, schema, **kwargs): - self.structured = (schema, kwargs) - return "STRUCTURED" - fake = FakeModel() self._patch_init(fake) return fake @@ -96,16 +91,15 @@ def test_schema_and_tools_are_bound_together_in_one_payload(self): result = Ai().build(ToolMovieAgent()) self.assertEqual(result, "BOUND") - self.assertIn(noop, fake.bound) - self.assertIn(Movie, fake.bound) + self.assertEqual(fake.bound, [noop, Movie]) - def test_schema_only_uses_enforced_structured_output(self): + def test_schema_only_binds_the_schema_as_a_tool(self): fake = self._fake_model() result = Ai().build(MovieAgent()) - self.assertEqual(result, "STRUCTURED") - self.assertEqual(fake.structured, (Movie, {"include_raw": True})) + self.assertEqual(result, "BOUND") + self.assertEqual(fake.bound, [Movie]) def test_tools_only_binds_just_the_tools(self): fake = self._fake_model() @@ -152,18 +146,6 @@ async def ainvoke(self, messages): self.assertEqual(result.content, "hello") - async def test_passes_structured_output_dict_through(self): - parsed = Movie(title="Inception", year=2010) - payload = {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} - - class Model: - async def ainvoke(self, messages): - return payload - - result = await Runner(MovieAgent(), Model()).run(["hi"]) - - self.assertEqual(result, payload) - class TestResponseMapping(unittest.TestCase): def test_unwraps_include_raw_into_parsed_and_content(self): @@ -191,23 +173,33 @@ def _patch(self, model): patcher.start() self.addCleanup(patcher.stop) - async def test_schema_only_populates_parsed(self): - parsed = Movie(title="Inception", year=2010) + async def test_schema_only_populates_parsed_from_the_schema_tool_call(self): + class FakeModel: + def bind_tools(self, tools, **kwargs): + return self - class Structured: async def ainvoke(self, messages): - return {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + return AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) + + self._patch(FakeModel()) + + response = await MovieAgent().prompt("best nolan movie") + + self.assertEqual(response.parsed, Movie(title="Inception", year=2010)) + async def test_schema_only_parses_plain_json_text_when_model_skips_the_tool(self): class FakeModel: - def with_structured_output(self, schema, **kwargs): - return Structured() + def bind_tools(self, tools, **kwargs): + return self + + async def ainvoke(self, messages): + return AIMessage(content='{"title": "Inception", "year": 2010}') self._patch(FakeModel()) response = await MovieAgent().prompt("best nolan movie") - self.assertEqual(response.parsed, parsed) - self.assertEqual(response.content, parsed.model_dump_json()) + self.assertEqual(response.parsed, Movie(title="Inception", year=2010)) async def test_schema_plus_tools_returns_structured_when_model_chooses_it(self): class FakeModel: From 9e08be266bb622d49118f25f818f7b5e1afc3042 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 01:24:21 -0700 Subject: [PATCH 7/7] 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. --- .../src/fastapi_startkit/ai/agent.py | 2 +- .../src/fastapi_startkit/ai/ai.py | 6 +- .../src/fastapi_startkit/ai/runner.py | 27 +-- .../tests/ai/test_structured_output.py | 171 +++++++----------- 4 files changed, 79 insertions(+), 127 deletions(-) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 5db24725..1931cd60 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -222,7 +222,7 @@ def _to_agent_response(self, result: Any) -> AgentResponse: def _apply_schema(self, response: AgentResponse) -> AgentResponse: schema = self.schema() - if schema is not None and not self.tools() and response.parsed is None and response.content: + if schema is not None and response.parsed is None and response.content: response.parsed = self._build_schema(schema, response.content) return response diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index c331f338..c7433e0a 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -84,11 +84,13 @@ def build( chat_model = init_chat_model(self._resolve_model(agent, model), **kwargs) tools = list(agent.tools()) + chat_model = chat_model.bind_tools(tools) if tools else chat_model + schema = agent.schema() if structured and schema is not None: - tools.append(schema) + chat_model = chat_model.with_structured_output(schema, include_raw=True) - return chat_model.bind_tools(tools) if tools else chat_model + 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) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index 293814aa..ba6ded73 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -22,35 +22,18 @@ def __init__(self, agent: Agent, model: Runnable[Any, BaseMessage]) -> None: self._tools: dict[str, BaseTool] = {tool.name: tool for tool in agent.tools()} self.model: Runnable[Any, BaseMessage] = model self.max_steps = agent.max_steps - self._schema = agent.schema() - self._schema_tool = self._schema_tool_name(self._schema) if self._schema is not None else None async def run(self, messages: Sequence[Message]) -> BaseMessage: history: list[Message] = list(messages) response: AIMessage = await self.model.ainvoke(history) # type: ignore[assignment] - tool_calls = list(getattr(response, "tool_calls", None) or []) - if not tool_calls: - return response - - for call in tool_calls: - if call.get("name") == self._schema_tool: - parsed = self._parse_schema(call.get("args") or {}) - return {"raw": response, "parsed": parsed, "parsing_error": None} # type: ignore[return-value] - - return (await self._run_tools(tool_calls))[-1] + if isinstance(response, dict) and "parsed" in response: + return response # type: ignore[return-value] - def _parse_schema(self, args: dict[str, Any]) -> Any: - schema = self._schema - if hasattr(schema, "model_validate"): - return schema.model_validate(args) - return schema(**args) - - @staticmethod - def _schema_tool_name(schema: Any) -> str: - from langchain_core.utils.function_calling import convert_to_openai_tool # noqa: PLC0415 + if not response.tool_calls: + return response - return convert_to_openai_tool(schema)["function"]["name"] + return (await self._run_tools(response.tool_calls))[-1] async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: return [await self._resolve_tool(call["name"]).ainvoke(call) for call in tool_calls] diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py index 12d95656..bfa60a46 100644 --- a/fastapi_startkit/tests/ai/test_structured_output.py +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -1,12 +1,10 @@ """Structured output. -An Agent's schema() is appended to its tools() and the whole set is bound to -the model in a single payload via bind_tools([*tools, schema]). The model picks -per turn: a real tool call (which the Runner executes) or the schema as its -structured answer (which the Runner parses into response.parsed). If a -schema-only agent replies with plain JSON text instead of the tool call, -_apply_schema still parses it, so the structured result comes through either -way. +When an Agent declares a schema(), the built model is wrapped with +model.with_structured_output(schema, include_raw=True) so the provider returns +the parsed object. Tools are bound as usual and left untouched. The wrapped +model yields {"raw", "parsed", "parsing_error"}; the Runner passes that through +and _to_agent_response unwraps it into response.parsed. The fake/record paths bypass build(), so they keep parsing the JSON-string content via schema() for deterministic replay. @@ -51,12 +49,23 @@ def tools(self): return [noop] -def _schema_tool_call(**args) -> dict: - return {"name": "Movie", "args": args, "id": "1", "type": "tool_call"} +def _real_tool_call(**args) -> dict: + return {"name": "noop", "args": args, "id": "1", "type": "tool_call"} -def _real_tool_call(**args) -> dict: - return {"name": "noop", "args": args, "id": "2", "type": "tool_call"} +class _FakeModel: + """Records the build calls made against it.""" + + def __init__(self): + self.calls = [] + + def bind_tools(self, tools, **kwargs): + self.calls.append(("bind_tools", list(tools))) + return self + + def with_structured_output(self, schema, **kwargs): + self.calls.append(("with_structured_output", schema, kwargs)) + return "STRUCTURED" class TestBuild(unittest.TestCase): @@ -68,41 +77,35 @@ def setUp(self): def tearDown(self): Ai.reset_fakes() - def _patch_init(self, fake): + def _patch(self, fake): patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: fake) patcher.start() self.addCleanup(patcher.stop) - def _fake_model(self): - class FakeModel: - bound = None + def test_schema_wraps_model_with_structured_output(self): + fake = _FakeModel() + self._patch(fake) - def bind_tools(self, tools, **kwargs): - self.bound = list(tools) - return "BOUND" + result = Ai().build(MovieAgent()) - fake = FakeModel() - self._patch_init(fake) - return fake + self.assertEqual(result, "STRUCTURED") + self.assertEqual(fake.calls, [("with_structured_output", Movie, {"include_raw": True})]) - def test_schema_and_tools_are_bound_together_in_one_payload(self): - fake = self._fake_model() + def test_tools_are_bound_then_structured_output_is_applied(self): + fake = _FakeModel() + self._patch(fake) result = Ai().build(ToolMovieAgent()) - self.assertEqual(result, "BOUND") - self.assertEqual(fake.bound, [noop, Movie]) - - def test_schema_only_binds_the_schema_as_a_tool(self): - fake = self._fake_model() - - result = Ai().build(MovieAgent()) - - self.assertEqual(result, "BOUND") - self.assertEqual(fake.bound, [Movie]) + self.assertEqual(result, "STRUCTURED") + self.assertEqual( + fake.calls, + [("bind_tools", [noop]), ("with_structured_output", Movie, {"include_raw": True})], + ) - def test_tools_only_binds_just_the_tools(self): - fake = self._fake_model() + def test_tools_only_binds_tools_without_structured_output(self): + fake = _FakeModel() + self._patch(fake) class ToolAgent(Agent): def tools(self): @@ -110,32 +113,38 @@ def tools(self): result = Ai().build(ToolAgent()) - self.assertEqual(result, "BOUND") - self.assertEqual(fake.bound, [noop]) + self.assertIs(result, fake) + self.assertEqual(fake.calls, [("bind_tools", [noop])]) - def test_streaming_drops_the_schema_from_the_payload(self): - fake = self._fake_model() + def test_streaming_skips_structured_output(self): + fake = _FakeModel() + self._patch(fake) result = Ai().build(ToolMovieAgent(), structured=False) - self.assertEqual(result, "BOUND") - self.assertEqual(fake.bound, [noop]) + self.assertIs(result, fake) + self.assertEqual(fake.calls, [("bind_tools", [noop])]) def test_no_schema_no_tools_returns_the_plain_model(self): - fake = self._fake_model() + fake = _FakeModel() + self._patch(fake) self.assertIs(Ai().build(Agent()), fake) + self.assertEqual(fake.calls, []) class TestRunner(unittest.IsolatedAsyncioTestCase): - async def test_returns_structured_when_model_calls_the_schema(self): + async def test_passes_structured_output_dict_through(self): + parsed = Movie(title="Inception", year=2010) + payload = {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + class Model: async def ainvoke(self, messages): - return AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) + return payload - result = await Runner(ToolMovieAgent(), Model()).run(["hi"]) + result = await Runner(MovieAgent(), Model()).run(["hi"]) - self.assertEqual(result["parsed"], Movie(title="Inception", year=2010)) + self.assertEqual(result, payload) async def test_runs_the_tool_when_model_calls_a_real_tool(self): class Model: @@ -150,9 +159,10 @@ async def ainvoke(self, messages): class TestResponseMapping(unittest.TestCase): def test_unwraps_include_raw_into_parsed_and_content(self): parsed = Movie(title="Inception", year=2010) - raw = AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) - response = MovieAgent()._to_agent_response({"raw": raw, "parsed": parsed, "parsing_error": None}) + response = MovieAgent()._to_agent_response( + {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + ) self.assertIs(response.parsed, parsed) self.assertEqual(response.content, parsed.model_dump_json()) @@ -168,65 +178,22 @@ def setUp(self): def tearDown(self): Ai.reset_fakes() - def _patch(self, model): - patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: model) - patcher.start() - self.addCleanup(patcher.stop) - - async def test_schema_only_populates_parsed_from_the_schema_tool_call(self): - class FakeModel: - def bind_tools(self, tools, **kwargs): - return self + async def test_prompt_populates_parsed_via_structured_output(self): + parsed = Movie(title="Inception", year=2010) + class Structured: async def ainvoke(self, messages): - return AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) - - self._patch(FakeModel()) - - response = await MovieAgent().prompt("best nolan movie") + return {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} - self.assertEqual(response.parsed, Movie(title="Inception", year=2010)) - - async def test_schema_only_parses_plain_json_text_when_model_skips_the_tool(self): class FakeModel: - def bind_tools(self, tools, **kwargs): - return self - - async def ainvoke(self, messages): - return AIMessage(content='{"title": "Inception", "year": 2010}') + def with_structured_output(self, schema, **kwargs): + return Structured() - self._patch(FakeModel()) + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: FakeModel()) + patcher.start() + self.addCleanup(patcher.stop) response = await MovieAgent().prompt("best nolan movie") - self.assertEqual(response.parsed, Movie(title="Inception", year=2010)) - - async def test_schema_plus_tools_returns_structured_when_model_chooses_it(self): - class FakeModel: - def bind_tools(self, tools, **kwargs): - return self - - async def ainvoke(self, messages): - return AIMessage(content="", tool_calls=[_schema_tool_call(title="Inception", year=2010)]) - - self._patch(FakeModel()) - - response = await ToolMovieAgent().prompt("best nolan movie") - - self.assertEqual(response.parsed, Movie(title="Inception", year=2010)) - self.assertEqual(response.tool_calls, []) - - async def test_schema_plus_tools_runs_the_tool_when_model_chooses_it(self): - class FakeModel: - def bind_tools(self, tools, **kwargs): - return self - - async def ainvoke(self, messages): - return AIMessage(content="", tool_calls=[_real_tool_call(query="hello")]) - - self._patch(FakeModel()) - - response = await ToolMovieAgent().prompt("run the tool") - - self.assertIsNone(response.parsed) - self.assertEqual(response.content, "hello") + self.assertEqual(response.parsed, parsed) + self.assertEqual(response.content, parsed.model_dump_json())