From 556f61279cb6fb293f9933cb5b4f5ea73c33f875 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 16:33:34 -0700 Subject: [PATCH 1/7] feat(ai): add TestJudgeAgent trajectory-match evaluator (evals) (#183) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(ai): add TestJudgeAgent trajectory-match evaluator + ModelBuilder fake registry Add a deterministic evals harness for testing AI agents, modeled on LangChain's agentevals create_trajectory_match_evaluator. TestJudgeAgent(trajectory_match_mode=...).record() yields a callable evaluator that compares an actual message trajectory against a reference trajectory with no live LLM call — pure structural comparison, typically driven by Agent.fake() output: with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: evaluation = evaluator(outputs=response["messages"], reference_outputs=reference) assert evaluation["score"] is True Supported modes mirror agentevals semantics: - strict: same messages in order, same role and tool calls per position (message content is not compared) - unordered: same set of tool calls, any order - subset: output tool calls all appear in the reference - superset: output tool calls cover all reference tool calls (extras ok) The evaluator returns a subscriptable result with key/score/comment, following the LangSmith evaluator-result convention. Also add a model-level fake registry to ModelBuilder: ModelBuilder.fake() registers a GenericFakeChatModel for an agent (by class name or instance), and get_model_for() returns it when present, otherwise builds a real provider model as before. Agent._build_model() now routes through get_model_for(), so a faked agent runs the same message-building / pipeline / tool-execution path as a real one with only the underlying model swapped. The existing Agent.fake()/FakeAgent/AgentBinding path is unchanged. * refactor(ai): rename ModelBuilder to Ai, rebuild fake()/record() around model-level fakes Ai (formerly ModelBuilder) drops the per-instance agent binding: fake_agent_models/ fake_agent_responses are class-level registries keyed by agent class name, and get_model_for()/build() take the agent as an argument instead of storing it in the constructor. Agent.fake() now registers a fixed, ordered list of replies as a deterministic chat model (via Ai.fake()) instead of binding a whole FakeAgent stand-in into the container. Replies flow through the real message-building/pipeline/tool- execution path, so faked tool calls actually execute — only the model at the bottom is swapped. Agent.record()/RecordingAgent/AgentBinding are unchanged. Drops the dead _match_fake()/self._fakes machinery on Agent (never populated by any code path) along with the now-unreachable FakeAgent/NoFakeResponse testing helpers. --- .../tests/features/test_chat_controller.py | 6 +- .../src/fastapi_startkit/ai/__init__.py | 11 +- .../src/fastapi_startkit/ai/agent.py | 46 +-- .../src/fastapi_startkit/ai/evals.py | 195 ++++++++++++ .../src/fastapi_startkit/ai/model_builder.py | 91 ++++-- .../src/fastapi_startkit/ai/testing.py | 68 +++-- fastapi_startkit/tests/ai/test_agent.py | 8 +- fastapi_startkit/tests/ai/test_agent_fake.py | 122 ++++---- .../tests/ai/test_agent_schema.py | 10 +- fastapi_startkit/tests/ai/test_ai_fake.py | 142 +++++++++ fastapi_startkit/tests/ai/test_evals.py | 281 ++++++++++++++++++ fastapi_startkit/uv.lock | 2 +- 12 files changed, 811 insertions(+), 171 deletions(-) create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/evals.py create mode 100644 fastapi_startkit/tests/ai/test_ai_fake.py create mode 100644 fastapi_startkit/tests/ai/test_evals.py diff --git a/example/agents/tests/features/test_chat_controller.py b/example/agents/tests/features/test_chat_controller.py index c0228d72..41eb02f7 100644 --- a/example/agents/tests/features/test_chat_controller.py +++ b/example/agents/tests/features/test_chat_controller.py @@ -4,14 +4,14 @@ class TestChatController(TestCase): - @RouterAgent.fake({"*hello*": "Hello there, This is no stream chat, Hope you are doing well."}) + @RouterAgent.fake(["Hello there, This is no stream chat, Hope you are doing well."]) async def test_it_responds_without_stream(self): response = await self.post("/chat", json={"message": "hello"}) response.assert_ok() response.assert_contents("Hello there, This is no stream chat, Hope you are doing well.") - @RouterAgent.fake({"*hello*": "Hello there, This is no stream chat, Hope you are doing well."}) + @RouterAgent.fake(["Hello there, This is no stream chat, Hope you are doing well."]) async def test_stream_assertions_are_rejected_on_a_buffered_response(self): response = await self.post("/chat", json={"message": "hello"}) @@ -21,7 +21,7 @@ async def test_stream_assertions_are_rejected_on_a_buffered_response(self): with self.assertRaises(AssertionError): response.assert_stream("Hello there, This is no stream chat, Hope you are doing well.") - @RouterAgent.fake({"*hello*": "Hello there, This is stream chat, Hope you are doing well."}) + @RouterAgent.fake(["Hello there, This is stream chat, Hope you are doing well."]) async def test_it_responds_with_stream(self): response = await self.post("/chat/stream", json={"message": "hello"}) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index dbf4c85e..790ef065 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -6,24 +6,26 @@ from .config.ai import AIConfig from .decorators import max_steps, max_tokens, model, provider, timeout, top_p from .document import Document +from .evals import TestJudgeAgent, TrajectoryEvaluator, TrajectoryMatchMode from .fakes import fake_chat_model from .image import Image, ImageResponse from .image_factory import ImageFactory +from .model_builder import Ai from .providers.ai_provider import AIProvider from .response import AgentResponse, AgentSnapshot -from .testing import AgentBinding, FakeAgent, NoFakeResponse, RecordingAgent +from .testing import AgentBinding, AgentModelFake, RecordingAgent __all__ = [ "Agent", + "Ai", "Middleware", "AgentBinding", + "AgentModelFake", "AgentResponse", "AgentSnapshot", "AIConfig", "AIProvider", "AnthropicConfig", - "FakeAgent", - "NoFakeResponse", "RecordingAgent", "Audio", "AudioResponse", @@ -36,6 +38,9 @@ "ImageFactory", "ImageResponse", "OpenAIConfig", + "TestJudgeAgent", + "TrajectoryEvaluator", + "TrajectoryMatchMode", "max_steps", "max_tokens", "model", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 17c947f3..e67a89ac 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -1,15 +1,16 @@ from __future__ import annotations -import fnmatch from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Optional, Type from .document import Document -from .response import AgentResponse, AgentSnapshot +from .response import AgentResponse from .testing import AgentBinding if TYPE_CHECKING: from langchain_core.tools import BaseTool + from .testing import AgentModelFake + class Agent: provider: str | None = None @@ -20,7 +21,6 @@ class Agent: top_p: float = 1.0 def __init__(self): - self._fakes: dict[str, AgentResponse | AgentSnapshot] = {} self._call_log: list[dict] = [] def messages(self) -> list[dict]: @@ -55,21 +55,6 @@ async def prompt( self._log_call("prompt", message) return self._apply_schema(response) - _run_kwargs = dict( - model=model, - attachments=attachments, - provider_options=provider_options, - ) - - match = self._match_fake(message) - if match is not None: - if isinstance(match, AgentSnapshot): - response = await match.resolve(self, message, **_run_kwargs) - else: - response = match - self._log_call("prompt", message) - return self._apply_schema(response) - messages = self._build_messages(message, attachments) chat_model = self._build_model(model, provider_options) @@ -96,22 +81,14 @@ async def stream( yield response.content return - fake = self._match_fake(message) - if fake is not None: - if isinstance(fake, AgentSnapshot): - response = await fake.resolve(self, message) - else: - response = fake - yield response.content - return async for chunk in self._stream(message, model=model, provider_options=provider_options): yield chunk @classmethod - def fake(cls, responses: dict | None = None) -> "AgentBinding": - from .testing import AgentBinding, FakeAgent + def fake(cls, responses: list) -> "AgentModelFake": + from .testing import AgentModelFake - return AgentBinding(cls, FakeAgent(responses)) + return AgentModelFake(cls, responses) @classmethod def record(cls, cassette: str | None = None) -> "AgentBinding": @@ -146,16 +123,9 @@ def assert_not_prompted(self) -> None: self.assert_prompted(times=0) def reset(self) -> "Agent": - self._fakes.clear() self._call_log.clear() return self - def _match_fake(self, message: str) -> Optional[AgentResponse | AgentSnapshot]: - for pattern, value in self._fakes.items(): - if fnmatch.fnmatch(message.lower(), pattern.lower()): - return value - return None - def _log_call(self, method: str, message: str) -> None: self._call_log.append({"method": method, "message": message}) @@ -223,9 +193,9 @@ def _build_messages( return messages def _build_model(self, model: str | None = None, provider_options: dict | None = None) -> Any: - from .model_builder import ModelBuilder # noqa: PLC0415 + from .model_builder import Ai # noqa: PLC0415 - return ModelBuilder(agent=self).build(model, provider_options) + return Ai().get_model_for(self, model, provider_options) def _to_agent_response(self, result: Any) -> AgentResponse: messages = result.get("messages", []) if isinstance(result, dict) else [] diff --git a/fastapi_startkit/src/fastapi_startkit/ai/evals.py b/fastapi_startkit/src/fastapi_startkit/ai/evals.py new file mode 100644 index 00000000..de59f3d9 --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/evals.py @@ -0,0 +1,195 @@ +"""Deterministic trajectory-match evaluators for testing AI agents. + +``TestJudgeAgent`` compares an agent's actual message trajectory against a +reference trajectory using matching semantics equivalent to LangChain's +``agentevals.trajectory.match.create_trajectory_match_evaluator``. There is +no live LLM call involved — it is pure, deterministic structural comparison, +meant to be driven by data the test already has on hand (typically the +output of ``Agent.fake()``). + +Mirrors the ``Agent.fake()``/``Agent.record()`` testing DSL: ``.record()`` +returns a context manager yielding a callable evaluator:: + + with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: + evaluation = evaluator(outputs=actual_messages, reference_outputs=reference_trajectory) + + assert evaluation["score"] is True +""" + +from __future__ import annotations + +from typing import Any, Callable, Literal, TypedDict, get_args + +TrajectoryMatchMode = Literal["strict", "unordered", "subset", "superset"] + +_TRAJECTORY_MATCH_MODES: tuple[str, ...] = get_args(TrajectoryMatchMode) + +_ROLE_BY_MESSAGE_TYPE = { + "human": "user", + "ai": "assistant", + "system": "system", + "tool": "tool", + "function": "tool", +} + + +class TrajectoryToolCall(TypedDict): + name: str + args: dict[str, Any] + + +class NormalizedMessage(TypedDict): + role: str + tool_calls: list[TrajectoryToolCall] + + +class EvaluatorResult(TypedDict): + key: str + score: bool + comment: str | None + + +def _normalize_tool_call(tool_call: dict) -> TrajectoryToolCall: + return {"name": tool_call.get("name", ""), "args": tool_call.get("args") or {}} + + +def _normalize_message(message: Any) -> NormalizedMessage: + if isinstance(message, dict): + role = message.get("role", "") + raw_tool_calls = message.get("tool_calls") or [] + else: + message_type = getattr(message, "type", "") + role = _ROLE_BY_MESSAGE_TYPE.get(message_type, message_type) + raw_tool_calls = getattr(message, "tool_calls", None) or [] + + return {"role": role, "tool_calls": [_normalize_tool_call(tc) for tc in raw_tool_calls]} + + +def _normalize_trajectory(trajectory: Any) -> list[NormalizedMessage]: + if isinstance(trajectory, dict): + trajectory = trajectory.get("messages", []) + return [_normalize_message(message) for message in trajectory] + + +def _extract_tool_calls(messages: list[NormalizedMessage]) -> list[TrajectoryToolCall]: + tool_calls: list[TrajectoryToolCall] = [] + for message in messages: + tool_calls.extend(message["tool_calls"]) + return tool_calls + + +def _is_tool_call_superset(calls: list[TrajectoryToolCall], wanted: list[TrajectoryToolCall]) -> bool: + """True if every tool call in ``wanted`` has a matching, unused tool call in ``calls``.""" + used = [False] * len(calls) + for want in wanted: + matched = False + for index, candidate in enumerate(calls): + if not used[index] and candidate == want: + used[index] = True + matched = True + break + if not matched: + return False + return True + + +def _tool_calls_equal(a: list[TrajectoryToolCall], b: list[TrajectoryToolCall]) -> bool: + return len(a) == len(b) and _is_tool_call_superset(a, b) and _is_tool_call_superset(b, a) + + +def _strict_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: + if len(outputs) != len(reference_outputs): + return False + return all( + output["role"] == reference["role"] and _tool_calls_equal(output["tool_calls"], reference["tool_calls"]) + for output, reference in zip(outputs, reference_outputs) + ) + + +def _unordered_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: + return _tool_calls_equal(_extract_tool_calls(outputs), _extract_tool_calls(reference_outputs)) + + +def _subset_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: + """``outputs``' tool calls must all appear in ``reference_outputs``.""" + return _is_tool_call_superset(_extract_tool_calls(reference_outputs), _extract_tool_calls(outputs)) + + +def _superset_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: + """``outputs``' tool calls must cover all of ``reference_outputs``'s.""" + return _is_tool_call_superset(_extract_tool_calls(outputs), _extract_tool_calls(reference_outputs)) + + +_SCORERS: dict[str, Callable[[list[NormalizedMessage], list[NormalizedMessage]], bool]] = { + "strict": _strict_match, + "unordered": _unordered_match, + "subset": _subset_match, + "superset": _superset_match, +} + + +class TrajectoryEvaluator: + """Callable returned by ``TestJudgeAgent.record()``; compares two trajectories.""" + + __test__ = False + + def __init__(self, trajectory_match_mode: TrajectoryMatchMode) -> None: + self.trajectory_match_mode: TrajectoryMatchMode = trajectory_match_mode + + def __call__(self, *, outputs: Any, reference_outputs: Any) -> EvaluatorResult: + normalized_outputs = _normalize_trajectory(outputs) + normalized_reference = _normalize_trajectory(reference_outputs) + scorer = _SCORERS[self.trajectory_match_mode] + score = scorer(normalized_outputs, normalized_reference) + return { + "key": f"trajectory_{self.trajectory_match_mode}_match", + "score": score, + "comment": None, + } + + +class TrajectoryJudgeBinding: + """Context manager returned by ``TestJudgeAgent.record()``.""" + + def __init__(self, trajectory_match_mode: TrajectoryMatchMode) -> None: + self._trajectory_match_mode: TrajectoryMatchMode = trajectory_match_mode + + def __enter__(self) -> TrajectoryEvaluator: + return TrajectoryEvaluator(self._trajectory_match_mode) + + def __exit__(self, *exc: Any) -> bool: + return False + + +class TestJudgeAgent: + """Deterministic judge for comparing agent trajectories in tests. + + Mirrors the ``Agent.fake()``/``Agent.record()`` testing DSL: ``.record()`` + returns a context manager yielding a callable evaluator, so trajectory + assertions read the same way as the rest of the agent testing toolkit. + + ``trajectory_match_mode`` controls how ``outputs`` is compared against + ``reference_outputs`` (each a list of LangChain messages, a list of + role/content/tool_calls dicts, or a dict with a ``messages`` key): + + - ``"strict"``: same number of messages, same role and same tool calls + at each position, in order. Message content is not compared. + - ``"unordered"``: the same set of tool calls were made, in any order + or position. + - ``"subset"``: every tool call in ``outputs`` also appears in + ``reference_outputs`` (no unexpected tool calls). + - ``"superset"``: every tool call in ``reference_outputs`` also appears + in ``outputs`` (no missing tool calls; extras are allowed). + """ + + __test__ = False + + def __init__(self, trajectory_match_mode: TrajectoryMatchMode = "strict") -> None: + if trajectory_match_mode not in _TRAJECTORY_MATCH_MODES: + raise ValueError( + f"Invalid trajectory_match_mode: {trajectory_match_mode!r}. Must be one of {_TRAJECTORY_MATCH_MODES!r}." + ) + self.trajectory_match_mode: TrajectoryMatchMode = trajectory_match_mode + + def record(self) -> TrajectoryJudgeBinding: + return TrajectoryJudgeBinding(self.trajectory_match_mode) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py b/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py index 7f82254f..5b65a237 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py @@ -8,40 +8,93 @@ from .agent import Agent -class ModelBuilder: - def __init__(self, agent: "Agent") -> None: - self._agent = agent +class Ai: + # Keyed by agent class name (see _key()) so a fake can be registered + # before any instance of that agent exists. get_model_for() consults + # this registry, so a faked agent runs through the same message-building + # / pipeline / tool-execution path as a real one — only the model at the + # bottom is swapped for a deterministic stand-in. + fake_agent_models: dict[str, Any] = {} + # Reserved for response-level fakes (mirroring fake_agent_models, but for + # cached final replies rather than whole chat models). Not yet wired up. + fake_agent_responses: dict[str, Any] = {} - def build(self, model: str | None = None, provider_options: dict | None = None) -> Any: + def __init__(self) -> None: + pass + + @staticmethod + def _key(agent: "Agent | str") -> str: + return agent if isinstance(agent, str) else type(agent).__name__ + + @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 + + turns = [message if hasattr(message, "content") else AIMessage(content=str(message)) for message in messages] + model = GenericFakeChatModel(messages=iter(turns)) + cls.fake_agent_models[cls._key(agent)] = model + return model + + @classmethod + def has_fake_model_for(cls, agent: "Agent | str") -> bool: + return cls._key(agent) in cls.fake_agent_models + + @classmethod + def get_fake_model_for(cls, agent: "Agent | str") -> Any: + return cls.fake_agent_models[cls._key(agent)] + + @classmethod + def forget(cls, agent: "Agent | str") -> None: + cls.fake_agent_models.pop(cls._key(agent), None) + + @classmethod + 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.""" + 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: from langchain.chat_models import init_chat_model # noqa: PLC0415 - lab = Lab.get_provider(self._agent.provider) + lab = Lab.get_provider(agent.provider) kwargs: dict[str, Any] = {"model_provider": lab.get_provider_key()} api_key = lab.get_api_key() if api_key: kwargs["api_key"] = api_key - if self._agent.max_tokens: - kwargs["max_tokens"] = self._agent.max_tokens - if self._agent.top_p != 1.0: - kwargs["top_p"] = self._agent.top_p - if self._agent.timeout: - kwargs["timeout"] = self._agent.timeout + if agent.max_tokens: + kwargs["max_tokens"] = agent.max_tokens + if agent.top_p != 1.0: + kwargs["top_p"] = agent.top_p + if agent.timeout: + kwargs["timeout"] = agent.timeout - kwargs.update(self._resolve_provider_options(provider_options)) + kwargs.update(self._resolve_provider_options(agent, provider_options)) - chat_model = init_chat_model(self._resolve_model(model), **kwargs) + chat_model = init_chat_model(self._resolve_model(agent, model), **kwargs) - tools = list(self._agent.tools()) + tools = list(agent.tools()) return chat_model.bind_tools(tools) if tools else chat_model - def _resolve_model(self, override: str | None = None) -> str: - return Lab.get_provider(self._agent.provider).get_model(override or self._agent.model or None) + 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) - def _resolve_provider_options(self, override: dict | None = None) -> dict: - options = dict(self._agent.provider_options().get(self._agent.provider, {})) + def _resolve_provider_options(self, agent: "Agent", override: dict | None = None) -> dict: + options = dict(agent.provider_options().get(agent.provider, {})) if override: - provider_specific = override.get(self._agent.provider, override) + provider_specific = override.get(agent.provider, override) if isinstance(provider_specific, dict): options.update(provider_specific) return options diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index e3e8533b..d456234c 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 from collections.abc import AsyncIterator from pathlib import Path @@ -18,10 +17,6 @@ from .document import Document -class NoFakeResponse(LookupError): - pass - - def _matches(pattern: str, message: str) -> bool: pattern, message = pattern.lower(), message.lower() if any(ch in pattern for ch in "*?["): @@ -29,12 +24,6 @@ def _matches(pattern: str, message: str) -> bool: return pattern in message -def _reply_text(reply: Any) -> str: - if isinstance(reply, AgentResponse): - return reply.content - return getattr(reply, "content", None) or str(reply) - - class _Recorder: def __init__(self) -> None: self.calls: list[str] = [] @@ -65,33 +54,46 @@ def _joined(value: Any) -> str: return "".join(value) if isinstance(value, list) else value -def _word_chunks(text: str) -> list[str]: - """Split text into word chunks (word + trailing whitespace) so a fake can - mimic a token stream. Loss-less: ``"".join(_word_chunks(t)) == t``.""" - return re.findall(r"\S+\s*", text) or [text] +class AgentModelFake: + """Registers a fixed, ordered list of replies as ``agent_cls``'s chat + model for the duration of a ``with`` block (or a decorated function). + Unlike the old pattern-matching stand-in, this swaps only the model — + ``prompt()``/``stream()`` still run the real message-building, pipeline, + and tool-execution path; see ``Ai.fake()``. + """ -class FakeAgent(_Recorder): - def __init__(self, responses: dict[str, Any] | None = None) -> None: - super().__init__() - self.responses = responses or {} + def __init__(self, agent_cls: type[Agent], responses: list) -> None: + self._agent_cls = agent_cls + self._responses = responses - def _resolve(self, message: str) -> str: - if not self.responses: - return "" - for pattern, reply in self.responses.items(): - if _matches(pattern, message): - return _reply_text(reply) - raise NoFakeResponse(f"No fake response matched message: {message!r}") + def __enter__(self) -> None: + from .model_builder import Ai - async def prompt(self, message: str, attachments: list[Document] | None = None) -> AgentResponse: - self._record_call(message, attachments) - return AgentResponse(content=self._resolve(message)) + Ai.fake(self._agent_cls.__name__, self._responses) - async def stream(self, message: str) -> AsyncIterator[str]: - self._record_call(message, None) - for chunk in _word_chunks(self._resolve(message)): - yield chunk + def __exit__(self, *_exc: Any) -> bool: + from .model_builder import Ai + + Ai.forget(self._agent_cls.__name__) + return False + + def __call__(self, func: Callable) -> Callable: + if inspect.iscoroutinefunction(func): + + @functools.wraps(func) + async def async_wrapper(*args: Any, **kwargs: Any) -> Any: + with self: + return await func(*args, **kwargs) + + return async_wrapper + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + with self: + return func(*args, **kwargs) + + return wrapper class RecordingAgent(_Recorder): diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 91a2203a..2d63867a 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -7,7 +7,7 @@ from fastapi_startkit.ai import AIConfig, Document, fake_chat_model from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.model_builder import ModelBuilder +from fastapi_startkit.ai.model_builder import Ai from fastapi_startkit.ai.response import AgentResponse from fastapi_startkit.application import app @@ -156,15 +156,15 @@ async def test_stream_yields_tool_result_without_calling_model_again(self): self.assertEqual(chunks, ["Python Developer at Shopify"]) def test_resolve_model_falls_back_to_lab_default(self): - self.assertEqual(ModelBuilder(Agent())._resolve_model(), "gemini-2.5-flash-lite") + self.assertEqual(Ai()._resolve_model(Agent()), "gemini-2.5-flash-lite") class AnthropicAgent(Agent): provider = "anthropic" - self.assertEqual(ModelBuilder(AnthropicAgent())._resolve_model(), "claude-sonnet-4-6") + self.assertEqual(Ai()._resolve_model(AnthropicAgent()), "claude-sonnet-4-6") def test_resolve_model_prefers_explicit_override(self): - self.assertEqual(ModelBuilder(Agent())._resolve_model("my-model"), "my-model") + self.assertEqual(Ai()._resolve_model(Agent(), "my-model"), "my-model") def test_instructions_lead_the_message_list(self): messages = JobAssistant()._build_messages("find me a job") diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 3191a8b3..7b560d74 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -1,9 +1,12 @@ """Tests for Agent.fake() / Agent.record() and the assert_prompted/reset helpers. -``Agent.fake()`` binds a canned stand-in into the container for the duration of a -``with`` block. ``Agent.record()`` binds a record-and-replay stand-in: on a cassette -miss it calls the real agent once and caches the response to disk; on a hit it -replays from the cassette without calling the agent again. +``Agent.fake()`` registers a fixed, ordered list of replies as a deterministic +stand-in chat model (see ``Ai.fake()``) for the duration of a ``with`` block — +each call to ``prompt()``/``stream()`` replays the next reply, going through +the real message-building / pipeline / tool-execution path, only the model at +the bottom is swapped. ``Agent.record()`` binds a record-and-replay stand-in: +on a cassette miss it calls the real agent once and caches the response to +disk; on a hit it replays from the cassette without calling the agent again. """ import json @@ -13,6 +16,7 @@ from unittest import mock from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.model_builder import Ai from fastapi_startkit.ai.response import AgentResponse @@ -21,92 +25,78 @@ class SimpleAgent(Agent): class TestAgentFake(unittest.IsolatedAsyncioTestCase): - async def test_fake_with_agent_response_returns_it(self): + def tearDown(self): + Ai.reset_fakes() + + async def test_fake_replays_the_only_response(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="Hello world!")}): + with SimpleAgent.fake(["Hello world!"]): result = await agent.prompt("anything") self.assertEqual(result.content, "Hello world!") - async def test_fake_does_not_call_provider_run(self): - agent = SimpleAgent() - called = [] - - original_run = agent._run - - async def patched_run(*args, **kwargs): - called.append(True) - return await original_run(*args, **kwargs) + async def test_fake_does_not_call_provider_build(self): + import langchain.chat_models as chat_models - agent._run = patched_run + def fail_if_called(*args, **kwargs): + raise AssertionError("init_chat_model must not be called when a fake is registered") - with SimpleAgent.fake({"*": AgentResponse(content="faked")}): - await agent.prompt("hello") - - self.assertEqual(called, [], "_run() must not be called when a fake matches") + patcher = mock.patch.object(chat_models, "init_chat_model", fail_if_called) + patcher.start() + self.addCleanup(patcher.stop) - async def test_fake_with_exact_pattern(self): agent = SimpleAgent() - with SimpleAgent.fake({"hello": AgentResponse(content="matched hello")}): - result = await agent.prompt("hello") - - self.assertEqual(result.content, "matched hello") + with SimpleAgent.fake(["faked"]): + await agent.prompt("hello") - async def test_fake_glob_hello_wildcard(self): + async def test_fake_replays_responses_in_order(self): agent = SimpleAgent() - with SimpleAgent.fake({"*hello*": AgentResponse(content="hi there")}): - result = await agent.prompt("say hello to me") + with SimpleAgent.fake(["first reply", "second reply"]): + first = await agent.prompt("call one") + second = await agent.prompt("call two") - self.assertEqual(result.content, "hi there") + self.assertEqual(first.content, "first reply") + self.assertEqual(second.content, "second reply") - async def test_fake_glob_analyze_wildcard(self): + async def test_fake_raises_once_responses_are_exhausted(self): agent = SimpleAgent() - with SimpleAgent.fake({"*analyze*": AgentResponse(content="analysis done")}): - result = await agent.prompt("please analyze this report") + with SimpleAgent.fake(["only reply"]): + await agent.prompt("call one") - self.assertEqual(result.content, "analysis done") + with self.assertRaises(RuntimeError): + await agent.prompt("call two") - async def test_fake_no_match_raises(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*hello*": AgentResponse(content="hi")}): - with self.assertRaises(Exception): - await agent.prompt("goodbye") + async def test_fake_unregisters_after_the_block_exits(self): + with SimpleAgent.fake(["faked"]): + self.assertTrue(Ai.has_fake_model_for("SimpleAgent")) - async def test_fake_glob_case_insensitive(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*HELLO*": AgentResponse(content="case insensitive")}): - result = await agent.prompt("say hello please") + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) - self.assertEqual(result.content, "case insensitive") + async def test_fake_as_decorator(self): + @SimpleAgent.fake(["decorated reply"]) + async def run(): + return await SimpleAgent().prompt("call") - async def test_fake_first_matching_pattern_wins(self): - agent = SimpleAgent() - with SimpleAgent.fake( - { - "*hello*": AgentResponse(content="first match"), - "*hello world*": AgentResponse(content="second match"), - } - ): - result = await agent.prompt("hello world") + result = await run() - self.assertEqual(result.content, "first match") + self.assertEqual(result.content, "decorated reply") async def test_assert_prompted_passes_after_one_call(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]): await agent.prompt("first") agent.assert_prompted() async def test_assert_prompted_times_2_passes_after_exactly_2_calls(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok", "ok"]): await agent.prompt("first") await agent.prompt("second") agent.assert_prompted(times=2) async def test_assert_prompted_times_fails_when_count_mismatch(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]): await agent.prompt("only once") with self.assertRaises(AssertionError): @@ -128,7 +118,7 @@ def test_assert_not_prompted_passes_when_no_calls_made(self): async def test_assert_not_prompted_fails_after_one_call(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]): await agent.prompt("a prompt") with self.assertRaises(AssertionError): @@ -136,7 +126,7 @@ async def test_assert_not_prompted_fails_after_one_call(self): async def test_reset_clears_call_log(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]): await agent.prompt("first") self.assertEqual(len(agent._call_log), 1) @@ -150,7 +140,7 @@ def test_reset_returns_agent_for_chaining(self): async def test_assert_not_prompted_passes_after_reset(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]): await agent.prompt("call before reset") agent.reset() @@ -159,33 +149,31 @@ async def test_assert_not_prompted_passes_after_reset(self): async def test_fake_rebinding_overrides_previous(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="first fake")}): + with SimpleAgent.fake(["first fake"]): self.assertEqual((await agent.prompt("call")).content, "first fake") - with SimpleAgent.fake({"*": AgentResponse(content="second fake")}): + with SimpleAgent.fake(["second fake"]): self.assertEqual((await agent.prompt("call again")).content, "second fake") async def test_stream_returns_fake_response(self): agent = SimpleAgent() - with SimpleAgent.fake({"*hello*": AgentResponse(content="Faked stream!")}): + with SimpleAgent.fake(["Faked stream!"]): chunks = [chunk async for chunk in agent.stream("hello world")] - # A faked stream is split into word chunks but rejoins to the value. self.assertEqual("".join(chunks), "Faked stream!") self.assertGreater(len(chunks), 1) agent.assert_prompted(times=1) - async def test_fake_stream_splits_value_into_word_chunks(self): + async def test_stream_replays_the_registered_text_exactly(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": "Hello there, friend"}): + with SimpleAgent.fake(["Hello there, friend"]): chunks = [chunk async for chunk in agent.stream("hi")] - self.assertEqual(chunks, ["Hello ", "there, ", "friend"]) self.assertEqual("".join(chunks), "Hello there, friend") async def test_stream_records_one_call_not_two(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="x")}): + with SimpleAgent.fake(["x"]): [chunk async for chunk in agent.stream("once")] # Streaming must log exactly one prompt — not one for stream + one for prompt. diff --git a/fastapi_startkit/tests/ai/test_agent_schema.py b/fastapi_startkit/tests/ai/test_agent_schema.py index 0e7d705b..1a30772c 100644 --- a/fastapi_startkit/tests/ai/test_agent_schema.py +++ b/fastapi_startkit/tests/ai/test_agent_schema.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.model_builder import Ai from fastapi_startkit.ai.response import AgentResponse @@ -20,9 +21,12 @@ def schema(self): class TestAgentSchema(unittest.IsolatedAsyncioTestCase): + def tearDown(self): + Ai.reset_fakes() + async def test_fake_json_is_built_into_the_schema(self): agent = UserAgent() - with UserAgent.fake({"*": '{"id": "u-1", "name": "Alex"}'}): + with UserAgent.fake(['{"id": "u-1", "name": "Alex"}']): response = await agent.prompt("get the user") self.assertIsInstance(response.parsed, User) @@ -32,14 +36,14 @@ async def test_fake_json_is_built_into_the_schema(self): async def test_no_schema_leaves_parsed_none(self): agent = Agent() - with Agent.fake({"*": '{"id": "u-1"}'}): + with Agent.fake(['{"id": "u-1"}']): response = await agent.prompt("anything") self.assertIsNone(response.parsed) async def test_invalid_json_for_schema_raises(self): agent = UserAgent() - with UserAgent.fake({"*": '{"name": "no id here"}'}): + with UserAgent.fake(['{"name": "no id here"}']): with self.assertRaises(Exception): await agent.prompt("get the user") diff --git a/fastapi_startkit/tests/ai/test_ai_fake.py b/fastapi_startkit/tests/ai/test_ai_fake.py new file mode 100644 index 00000000..30b2e39e --- /dev/null +++ b/fastapi_startkit/tests/ai/test_ai_fake.py @@ -0,0 +1,142 @@ +"""Tests for Ai's fake-model registry. + +Ai.fake() swaps the chat model a given agent (by class name or instance) +resolves to for a deterministic GenericFakeChatModel that replays a fixed +list of message turns — no live LLM call, no network access. +Ai().get_model_for(agent) is what Agent._build_model() calls: it returns +the registered fake when one exists, otherwise it builds a real provider +model exactly as Ai.build() always has. +""" + +import unittest +from unittest import mock + +import langchain.chat_models as chat_models +from langchain_core.messages import AIMessage, ToolCall +from langchain_core.tools import tool + +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.model_builder import Ai +from fastapi_startkit.application import app + + +@tool +def search_jobs(query: str) -> str: + """Search the job board for roles matching the query.""" + return "Python Developer at Shopify" + + +class JobAssistant(Agent): + def tools(self): + return [search_jobs] + + +class SimpleAgent(Agent): + pass + + +class TestAiFakeBase(unittest.IsolatedAsyncioTestCase): + def setUp(self): + from fastapi_startkit.ai import AIConfig + + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + self.addCleanup(Ai.reset_fakes) + + +class TestAiFakeRegistration(TestAiFakeBase): + def test_fake_registers_a_model_for_an_agent_class_name(self): + Ai.fake("SimpleAgent", [AIMessage(content="hi")]) + + self.assertTrue(Ai.has_fake_model_for("SimpleAgent")) + + def test_fake_accepts_an_agent_instance_keyed_by_its_class_name(self): + Ai.fake(SimpleAgent(), [AIMessage(content="hi")]) + + self.assertTrue(Ai.has_fake_model_for(SimpleAgent())) + self.assertTrue(Ai.has_fake_model_for("SimpleAgent")) + + def test_has_fake_model_for_is_false_when_nothing_registered(self): + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) + + def test_fake_coerces_plain_strings_into_ai_messages(self): + model = Ai.fake("SimpleAgent", ["plain text reply"]) + + result = model.invoke([]) + + self.assertEqual(result.content, "plain text reply") + + def test_fake_returns_the_registered_chat_model(self): + model = Ai.fake("SimpleAgent", [AIMessage(content="hi")]) + + self.assertIs(Ai.get_fake_model_for("SimpleAgent"), model) + + def test_forget_removes_a_single_registration(self): + Ai.fake("SimpleAgent", [AIMessage(content="hi")]) + Ai.fake("JobAssistant", [AIMessage(content="hi")]) + + Ai.forget("SimpleAgent") + + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) + self.assertTrue(Ai.has_fake_model_for("JobAssistant")) + + def test_forget_is_a_no_op_when_nothing_registered(self): + Ai.forget("SimpleAgent") + + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) + + +class TestAiGetModelFor(TestAiFakeBase): + def test_returns_registered_fake_without_building_a_real_model(self): + fake_model = Ai.fake("SimpleAgent", [AIMessage(content="faked")]) + + def fail_if_called(*args, **kwargs): + raise AssertionError("init_chat_model must not be called when a fake is registered") + + patcher = mock.patch.object(chat_models, "init_chat_model", fail_if_called) + patcher.start() + self.addCleanup(patcher.stop) + + resolved = Ai().get_model_for(SimpleAgent()) + + self.assertIs(resolved, fake_model) + + def test_falls_back_to_build_when_no_fake_is_registered(self): + sentinel = object() + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: sentinel) + patcher.start() + self.addCleanup(patcher.stop) + + resolved = Ai().get_model_for(SimpleAgent()) + + self.assertIs(resolved, sentinel) + + +class TestAgentPromptUsesFakeModelEndToEnd(TestAiFakeBase): + async def test_prompt_replays_the_registered_fake_model_reply(self): + Ai.fake("SimpleAgent", [AIMessage(content="faked via ai")]) + + result = await SimpleAgent().prompt("hi there") + + self.assertEqual(result.content, "faked via ai") + + async def test_prompt_runs_a_faked_tool_call_end_to_end(self): + Ai.fake( + "JobAssistant", + [ + AIMessage( + content="", + tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], + ) + ], + ) + + result = await JobAssistant().prompt("find me a python job") + + self.assertEqual(result.content, "Python Developer at Shopify") + + async def test_registering_a_fake_does_not_affect_other_agent_classes(self): + Ai.fake("SimpleAgent", [AIMessage(content="only for SimpleAgent")]) + + self.assertFalse(Ai.has_fake_model_for("JobAssistant")) diff --git a/fastapi_startkit/tests/ai/test_evals.py b/fastapi_startkit/tests/ai/test_evals.py new file mode 100644 index 00000000..5e7e2ccc --- /dev/null +++ b/fastapi_startkit/tests/ai/test_evals.py @@ -0,0 +1,281 @@ +"""Tests for TestJudgeAgent — the deterministic trajectory-match evaluator. + +TestJudgeAgent compares an actual agent message trajectory against a reference +trajectory with no live LLM calls: it is pure, deterministic structural +comparison, driven entirely by data the tests construct themselves (typically +the output of Agent.fake()). Mirrors the Agent.fake()/Agent.record() testing +DSL: ``.record()`` returns a context manager yielding a callable evaluator. +""" + +import unittest + +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage + +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.evals import TestJudgeAgent +from fastapi_startkit.ai.response import AgentResponse + + +def _weather_trajectory(city: str = "San Francisco") -> list: + return [ + HumanMessage(content=f"What is the weather in {city}?"), + AIMessage( + content="", + tool_calls=[{"id": "call_1", "name": "get_weather", "args": {"city": city}}], + ), + ToolMessage(content="It is 75 degrees and sunny.", tool_call_id="call_1"), + AIMessage(content=f"The weather in {city} is 75 degrees and sunny."), + ] + + +class TestJudgeAgentConstruction(unittest.TestCase): + def test_defaults_to_strict_mode(self): + judge = TestJudgeAgent() + self.assertEqual(judge.trajectory_match_mode, "strict") + + def test_accepts_each_supported_mode(self): + for mode in ("strict", "unordered", "subset", "superset"): + judge = TestJudgeAgent(trajectory_match_mode=mode) + self.assertEqual(judge.trajectory_match_mode, mode) + + def test_rejects_unknown_mode(self): + with self.assertRaises(ValueError): + TestJudgeAgent(trajectory_match_mode="fuzzy") + + +class TestJudgeAgentRecordContextManager(unittest.TestCase): + def test_record_yields_callable_evaluator(self): + with TestJudgeAgent().record() as evaluator: + self.assertTrue(callable(evaluator)) + + def test_evaluation_result_is_subscriptable_with_score_key(self): + trajectory = _weather_trajectory() + with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: + evaluation = evaluator(outputs=trajectory, reference_outputs=trajectory) + + self.assertIs(evaluation["score"], True) + self.assertEqual(evaluation["key"], "trajectory_strict_match") + self.assertIn("comment", evaluation) + + +class TestStrictTrajectoryMatch(unittest.TestCase): + def setUp(self): + self.judge = TestJudgeAgent(trajectory_match_mode="strict") + + def _evaluate(self, outputs, reference_outputs): + with self.judge.record() as evaluator: + return evaluator(outputs=outputs, reference_outputs=reference_outputs) + + def test_identical_trajectory_matches(self): + trajectory = _weather_trajectory() + evaluation = self._evaluate(trajectory, trajectory) + self.assertIs(evaluation["score"], True) + + def test_matches_even_when_final_message_content_differs(self): + reference = _weather_trajectory() + outputs = _weather_trajectory() + outputs[-1] = AIMessage(content="Completely different wording, still sunny.") + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], True) + + def test_fails_when_tool_call_args_differ(self): + reference = _weather_trajectory(city="San Francisco") + outputs = _weather_trajectory(city="Oakland") + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + def test_fails_on_extra_trailing_message(self): + reference = _weather_trajectory() + outputs = _weather_trajectory() + [AIMessage(content="One more thing...")] + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + def test_fails_on_missing_message(self): + reference = _weather_trajectory() + outputs = _weather_trajectory()[:-1] + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + def test_fails_when_role_order_differs(self): + reference = _weather_trajectory() + outputs = list(reversed(_weather_trajectory())) + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + def test_fails_when_one_side_has_tool_calls_and_other_does_not(self): + reference = _weather_trajectory() + outputs = _weather_trajectory() + outputs[1] = AIMessage(content="") + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + def test_works_with_plain_dict_messages(self): + reference = _weather_trajectory() + outputs = [ + {"role": "user", "content": "What is the weather in San Francisco?"}, + { + "role": "assistant", + "content": "", + "tool_calls": [{"id": "call_1", "name": "get_weather", "args": {"city": "San Francisco"}}], + }, + {"role": "tool", "content": "It is 75 degrees and sunny.", "tool_call_id": "call_1"}, + {"role": "assistant", "content": "The weather in San Francisco is 75 degrees and sunny."}, + ] + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], True) + + +class TestUnorderedTrajectoryMatch(unittest.TestCase): + def setUp(self): + self.judge = TestJudgeAgent(trajectory_match_mode="unordered") + + def _evaluate(self, outputs, reference_outputs): + with self.judge.record() as evaluator: + return evaluator(outputs=outputs, reference_outputs=reference_outputs) + + def test_same_tool_calls_in_different_message_order_matches(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [HumanMessage(content="hi"), weather_call, time_call] + outputs = [HumanMessage(content="hi"), time_call, weather_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], True) + + def test_missing_tool_call_fails(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [weather_call, time_call] + outputs = [weather_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + def test_extra_tool_call_fails(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [weather_call] + outputs = [weather_call, time_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + +class TestSubsetTrajectoryMatch(unittest.TestCase): + def setUp(self): + self.judge = TestJudgeAgent(trajectory_match_mode="subset") + + def _evaluate(self, outputs, reference_outputs): + with self.judge.record() as evaluator: + return evaluator(outputs=outputs, reference_outputs=reference_outputs) + + def test_output_tool_calls_within_reference_matches(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [weather_call, time_call] + outputs = [weather_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], True) + + def test_output_tool_call_not_in_reference_fails(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [weather_call] + outputs = [weather_call, time_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + +class TestSupersetTrajectoryMatch(unittest.TestCase): + def setUp(self): + self.judge = TestJudgeAgent(trajectory_match_mode="superset") + + def _evaluate(self, outputs, reference_outputs): + with self.judge.record() as evaluator: + return evaluator(outputs=outputs, reference_outputs=reference_outputs) + + def test_output_contains_all_reference_tool_calls_plus_extra_matches(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [weather_call] + outputs = [weather_call, time_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], True) + + def test_output_missing_a_required_reference_tool_call_fails(self): + weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) + time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) + + reference = [weather_call, time_call] + outputs = [weather_call] + + evaluation = self._evaluate(outputs, reference) + self.assertIs(evaluation["score"], False) + + +class SimpleAgent(Agent): + pass + + +class TestJudgeAgentWithAgentFakeOutput(unittest.TestCase): + """Demonstrates the intended workflow: Agent.fake() drives deterministic + output with no live LLM call, and TestJudgeAgent judges the resulting + trajectory against a reference.""" + + async def _fake_response(self) -> AgentResponse: + agent = SimpleAgent() + with SimpleAgent.fake(["It is sunny in San Francisco."]): + return await agent.prompt("What is the weather in San Francisco?") + + def test_fake_agent_reply_matches_reference_trajectory_in_strict_mode(self): + import asyncio + + response = asyncio.run(self._fake_response()) + + outputs = [ + HumanMessage(content="What is the weather in San Francisco?"), + AIMessage(content=response.content), + ] + reference_trajectory = [ + HumanMessage(content="What is the weather in San Francisco?"), + AIMessage(content="It is sunny in San Francisco."), + ] + + with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: + evaluation = evaluator(outputs=outputs, reference_outputs=reference_trajectory) + + self.assertIs(evaluation["score"], True) + + +class TestJudgeAgentMessagesDictInput(unittest.TestCase): + def test_accepts_outputs_dict_with_messages_key(self): + trajectory = _weather_trajectory() + with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: + evaluation = evaluator(outputs={"messages": trajectory}, reference_outputs=trajectory) + + self.assertIs(evaluation["score"], True) + + +class TestJudgeAgentIsNotCollectedByPytest(unittest.TestCase): + def test_has_test_dunder_set_to_false(self): + self.assertFalse(getattr(TestJudgeAgent, "__test__")) + + +class TestSystemMessagesAreComparedByRole(unittest.TestCase): + def test_strict_mode_compares_system_message_role(self): + reference = [SystemMessage(content="Be concise."), HumanMessage(content="hi")] + outputs = [SystemMessage(content="Be very concise."), HumanMessage(content="hi")] + + with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: + evaluation = evaluator(outputs=outputs, reference_outputs=reference) + + self.assertIs(evaluation["score"], True) diff --git a/fastapi_startkit/uv.lock b/fastapi_startkit/uv.lock index 5808c35b..c9ea4cba 100644 --- a/fastapi_startkit/uv.lock +++ b/fastapi_startkit/uv.lock @@ -527,7 +527,7 @@ wheels = [ [[package]] name = "fastapi-startkit" -version = "0.48.0" +version = "0.50.0" source = { editable = "." } dependencies = [ { name = "cleo" }, From 4bf8f3878a8838ab3f516e23ff4610f301502dd3 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 21:03:16 -0700 Subject: [PATCH 2/7] feat(ai): Ai model registry + fluent record() testing DSL (#184) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat(ai): add TestJudgeAgent trajectory-match evaluator + ModelBuilder fake registry Add a deterministic evals harness for testing AI agents, modeled on LangChain's agentevals create_trajectory_match_evaluator. TestJudgeAgent(trajectory_match_mode=...).record() yields a callable evaluator that compares an actual message trajectory against a reference trajectory with no live LLM call — pure structural comparison, typically driven by Agent.fake() output: with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: evaluation = evaluator(outputs=response["messages"], reference_outputs=reference) assert evaluation["score"] is True Supported modes mirror agentevals semantics: - strict: same messages in order, same role and tool calls per position (message content is not compared) - unordered: same set of tool calls, any order - subset: output tool calls all appear in the reference - superset: output tool calls cover all reference tool calls (extras ok) The evaluator returns a subscriptable result with key/score/comment, following the LangSmith evaluator-result convention. Also add a model-level fake registry to ModelBuilder: ModelBuilder.fake() registers a GenericFakeChatModel for an agent (by class name or instance), and get_model_for() returns it when present, otherwise builds a real provider model as before. Agent._build_model() now routes through get_model_for(), so a faked agent runs the same message-building / pipeline / tool-execution path as a real one with only the underlying model swapped. The existing Agent.fake()/FakeAgent/AgentBinding path is unchanged. * refactor(ai): rename ModelBuilder to Ai, rebuild fake()/record() around model-level fakes Ai (formerly ModelBuilder) drops the per-instance agent binding: fake_agent_models/ fake_agent_responses are class-level registries keyed by agent class name, and get_model_for()/build() take the agent as an argument instead of storing it in the constructor. Agent.fake() now registers a fixed, ordered list of replies as a deterministic chat model (via Ai.fake()) instead of binding a whole FakeAgent stand-in into the container. Replies flow through the real message-building/pipeline/tool- execution path, so faked tool calls actually execute — only the model at the bottom is swapped. Agent.record()/RecordingAgent/AgentBinding are unchanged. Drops the dead _match_fake()/self._fakes machinery on Agent (never populated by any code path) along with the now-unreachable FakeAgent/NoFakeResponse testing helpers. * refactor(ai): rename model_builder.py to ai.py * feat(ai): fluent Agent.record() testing DSL Bind record() with `as agent` to get a synchronous prompt() plus assert_text_response(), assert_tool_called()/assert_tool_not_called(), assert_response_time_lt(), and assert_response_judged() (LLM-as-judge, verdict cached in the cassette) that all judge the most recent turn. record(cassette, messages=...) seeds a session's prior history; cassette keys now fold in the conversation so far (not just the literal message text) so two sessions with different histories but the same follow-up text don't collide. The pre-existing bare context-manager usage from task #327 is unchanged. Removes TestJudgeAgent/evals.py: its deterministic structural trajectory matching is unused anywhere outside its own tests, and this DSL now covers response judging via assert_response_judged. * fix(ai): fluent record() prompt() is async, not sync RecordingAgent.prompt() no longer wraps the real async call with asyncio.run() — it's now the same async method the bare context-manager path already awaited internally. Wrapping every call in a fresh event loop was unnecessary and would break inside any already-running loop. Fluent tests now use IsolatedAsyncioTestCase + await agent.prompt(...). * refactor(ai): drop dead non-streaming fallback in Agent.stream() AgentBinding only ever wraps a RecordingAgent (from record()), which always implements stream() — the hasattr(swapped, "stream") check and its buffered-response fallback could never actually run. * style(ai): remove stale comments on Ai's fake registries --- .../tests/units/agents/test_router_agent.py | 30 ++ .../src/fastapi_startkit/ai/__init__.py | 9 +- .../src/fastapi_startkit/ai/agent.py | 14 +- .../ai/{model_builder.py => ai.py} | 7 - .../src/fastapi_startkit/ai/evals.py | 195 ---------- .../src/fastapi_startkit/ai/testing.py | 164 +++++++- fastapi_startkit/tests/ai/test_agent.py | 2 +- fastapi_startkit/tests/ai/test_agent_fake.py | 5 +- .../tests/ai/test_agent_record_fluent.py | 354 ++++++++++++++++++ .../tests/ai/test_agent_schema.py | 2 +- fastapi_startkit/tests/ai/test_ai_fake.py | 2 +- fastapi_startkit/tests/ai/test_evals.py | 281 -------------- 12 files changed, 553 insertions(+), 512 deletions(-) create mode 100644 example/agents/tests/units/agents/test_router_agent.py rename fastapi_startkit/src/fastapi_startkit/ai/{model_builder.py => ai.py} (87%) delete mode 100644 fastapi_startkit/src/fastapi_startkit/ai/evals.py create mode 100644 fastapi_startkit/tests/ai/test_agent_record_fluent.py delete mode 100644 fastapi_startkit/tests/ai/test_evals.py diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py new file mode 100644 index 00000000..26ddfae7 --- /dev/null +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -0,0 +1,30 @@ +from langchain_core.messages import AIMessage, HumanMessage + +from app.agents.chat import RouterAgent + + +class TestRouterAgent: + 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", + expectation="The llm should respond with greetings", + ) + agent.assert_response_time_lt(5) + + await agent.prompt("suggest python developer jobs") + agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") + + async def test_the_router_with_initial_messages(self): + with RouterAgent.record( + "record_stream.json", + messages=[ + HumanMessage(content="Hi"), + AIMessage(content="Hello, How can I help you?"), + ], + ) as agent: + await agent.prompt("suggest python developer jobs") + agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index 790ef065..1e8f8b8c 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -6,14 +6,13 @@ from .config.ai import AIConfig from .decorators import max_steps, max_tokens, model, provider, timeout, top_p from .document import Document -from .evals import TestJudgeAgent, TrajectoryEvaluator, TrajectoryMatchMode from .fakes import fake_chat_model from .image import Image, ImageResponse from .image_factory import ImageFactory -from .model_builder import Ai +from .ai import Ai from .providers.ai_provider import AIProvider from .response import AgentResponse, AgentSnapshot -from .testing import AgentBinding, AgentModelFake, RecordingAgent +from .testing import AgentBinding, AgentModelFake, RecordingAgent, ToolCallView __all__ = [ "Agent", @@ -27,6 +26,7 @@ "AIProvider", "AnthropicConfig", "RecordingAgent", + "ToolCallView", "Audio", "AudioResponse", "AudioFactory", @@ -38,9 +38,6 @@ "ImageFactory", "ImageResponse", "OpenAIConfig", - "TestJudgeAgent", - "TrajectoryEvaluator", - "TrajectoryMatchMode", "max_steps", "max_tokens", "model", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index e67a89ac..96926185 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -73,12 +73,8 @@ async def stream( swapped = self._faked() if swapped is not None: - if hasattr(swapped, "stream"): - async for chunk in swapped.stream(message): - yield chunk - else: - response = await swapped.prompt(message) - yield response.content + async for chunk in swapped.stream(message): + yield chunk return async for chunk in self._stream(message, model=model, provider_options=provider_options): @@ -91,10 +87,10 @@ def fake(cls, responses: list) -> "AgentModelFake": return AgentModelFake(cls, responses) @classmethod - def record(cls, cassette: str | None = None) -> "AgentBinding": + def record(cls, cassette: str | None = None, messages: list | None = None) -> "AgentBinding": from .testing import AgentBinding, RecordingAgent - return AgentBinding(cls, RecordingAgent(cls(), cassette)) + return AgentBinding(cls, RecordingAgent(cls(), cassette, messages)) @classmethod def _binding(cls) -> Any: @@ -193,7 +189,7 @@ def _build_messages( return messages def _build_model(self, model: str | None = None, provider_options: dict | None = None) -> Any: - from .model_builder import Ai # noqa: PLC0415 + from .ai import Ai # noqa: PLC0415 return Ai().get_model_for(self, model, provider_options) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py similarity index 87% rename from fastapi_startkit/src/fastapi_startkit/ai/model_builder.py rename to fastapi_startkit/src/fastapi_startkit/ai/ai.py index 5b65a237..da1dc599 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -9,14 +9,7 @@ class Ai: - # Keyed by agent class name (see _key()) so a fake can be registered - # before any instance of that agent exists. get_model_for() consults - # this registry, so a faked agent runs through the same message-building - # / pipeline / tool-execution path as a real one — only the model at the - # bottom is swapped for a deterministic stand-in. fake_agent_models: dict[str, Any] = {} - # Reserved for response-level fakes (mirroring fake_agent_models, but for - # cached final replies rather than whole chat models). Not yet wired up. fake_agent_responses: dict[str, Any] = {} def __init__(self) -> None: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/evals.py b/fastapi_startkit/src/fastapi_startkit/ai/evals.py deleted file mode 100644 index de59f3d9..00000000 --- a/fastapi_startkit/src/fastapi_startkit/ai/evals.py +++ /dev/null @@ -1,195 +0,0 @@ -"""Deterministic trajectory-match evaluators for testing AI agents. - -``TestJudgeAgent`` compares an agent's actual message trajectory against a -reference trajectory using matching semantics equivalent to LangChain's -``agentevals.trajectory.match.create_trajectory_match_evaluator``. There is -no live LLM call involved — it is pure, deterministic structural comparison, -meant to be driven by data the test already has on hand (typically the -output of ``Agent.fake()``). - -Mirrors the ``Agent.fake()``/``Agent.record()`` testing DSL: ``.record()`` -returns a context manager yielding a callable evaluator:: - - with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: - evaluation = evaluator(outputs=actual_messages, reference_outputs=reference_trajectory) - - assert evaluation["score"] is True -""" - -from __future__ import annotations - -from typing import Any, Callable, Literal, TypedDict, get_args - -TrajectoryMatchMode = Literal["strict", "unordered", "subset", "superset"] - -_TRAJECTORY_MATCH_MODES: tuple[str, ...] = get_args(TrajectoryMatchMode) - -_ROLE_BY_MESSAGE_TYPE = { - "human": "user", - "ai": "assistant", - "system": "system", - "tool": "tool", - "function": "tool", -} - - -class TrajectoryToolCall(TypedDict): - name: str - args: dict[str, Any] - - -class NormalizedMessage(TypedDict): - role: str - tool_calls: list[TrajectoryToolCall] - - -class EvaluatorResult(TypedDict): - key: str - score: bool - comment: str | None - - -def _normalize_tool_call(tool_call: dict) -> TrajectoryToolCall: - return {"name": tool_call.get("name", ""), "args": tool_call.get("args") or {}} - - -def _normalize_message(message: Any) -> NormalizedMessage: - if isinstance(message, dict): - role = message.get("role", "") - raw_tool_calls = message.get("tool_calls") or [] - else: - message_type = getattr(message, "type", "") - role = _ROLE_BY_MESSAGE_TYPE.get(message_type, message_type) - raw_tool_calls = getattr(message, "tool_calls", None) or [] - - return {"role": role, "tool_calls": [_normalize_tool_call(tc) for tc in raw_tool_calls]} - - -def _normalize_trajectory(trajectory: Any) -> list[NormalizedMessage]: - if isinstance(trajectory, dict): - trajectory = trajectory.get("messages", []) - return [_normalize_message(message) for message in trajectory] - - -def _extract_tool_calls(messages: list[NormalizedMessage]) -> list[TrajectoryToolCall]: - tool_calls: list[TrajectoryToolCall] = [] - for message in messages: - tool_calls.extend(message["tool_calls"]) - return tool_calls - - -def _is_tool_call_superset(calls: list[TrajectoryToolCall], wanted: list[TrajectoryToolCall]) -> bool: - """True if every tool call in ``wanted`` has a matching, unused tool call in ``calls``.""" - used = [False] * len(calls) - for want in wanted: - matched = False - for index, candidate in enumerate(calls): - if not used[index] and candidate == want: - used[index] = True - matched = True - break - if not matched: - return False - return True - - -def _tool_calls_equal(a: list[TrajectoryToolCall], b: list[TrajectoryToolCall]) -> bool: - return len(a) == len(b) and _is_tool_call_superset(a, b) and _is_tool_call_superset(b, a) - - -def _strict_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: - if len(outputs) != len(reference_outputs): - return False - return all( - output["role"] == reference["role"] and _tool_calls_equal(output["tool_calls"], reference["tool_calls"]) - for output, reference in zip(outputs, reference_outputs) - ) - - -def _unordered_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: - return _tool_calls_equal(_extract_tool_calls(outputs), _extract_tool_calls(reference_outputs)) - - -def _subset_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: - """``outputs``' tool calls must all appear in ``reference_outputs``.""" - return _is_tool_call_superset(_extract_tool_calls(reference_outputs), _extract_tool_calls(outputs)) - - -def _superset_match(outputs: list[NormalizedMessage], reference_outputs: list[NormalizedMessage]) -> bool: - """``outputs``' tool calls must cover all of ``reference_outputs``'s.""" - return _is_tool_call_superset(_extract_tool_calls(outputs), _extract_tool_calls(reference_outputs)) - - -_SCORERS: dict[str, Callable[[list[NormalizedMessage], list[NormalizedMessage]], bool]] = { - "strict": _strict_match, - "unordered": _unordered_match, - "subset": _subset_match, - "superset": _superset_match, -} - - -class TrajectoryEvaluator: - """Callable returned by ``TestJudgeAgent.record()``; compares two trajectories.""" - - __test__ = False - - def __init__(self, trajectory_match_mode: TrajectoryMatchMode) -> None: - self.trajectory_match_mode: TrajectoryMatchMode = trajectory_match_mode - - def __call__(self, *, outputs: Any, reference_outputs: Any) -> EvaluatorResult: - normalized_outputs = _normalize_trajectory(outputs) - normalized_reference = _normalize_trajectory(reference_outputs) - scorer = _SCORERS[self.trajectory_match_mode] - score = scorer(normalized_outputs, normalized_reference) - return { - "key": f"trajectory_{self.trajectory_match_mode}_match", - "score": score, - "comment": None, - } - - -class TrajectoryJudgeBinding: - """Context manager returned by ``TestJudgeAgent.record()``.""" - - def __init__(self, trajectory_match_mode: TrajectoryMatchMode) -> None: - self._trajectory_match_mode: TrajectoryMatchMode = trajectory_match_mode - - def __enter__(self) -> TrajectoryEvaluator: - return TrajectoryEvaluator(self._trajectory_match_mode) - - def __exit__(self, *exc: Any) -> bool: - return False - - -class TestJudgeAgent: - """Deterministic judge for comparing agent trajectories in tests. - - Mirrors the ``Agent.fake()``/``Agent.record()`` testing DSL: ``.record()`` - returns a context manager yielding a callable evaluator, so trajectory - assertions read the same way as the rest of the agent testing toolkit. - - ``trajectory_match_mode`` controls how ``outputs`` is compared against - ``reference_outputs`` (each a list of LangChain messages, a list of - role/content/tool_calls dicts, or a dict with a ``messages`` key): - - - ``"strict"``: same number of messages, same role and same tool calls - at each position, in order. Message content is not compared. - - ``"unordered"``: the same set of tool calls were made, in any order - or position. - - ``"subset"``: every tool call in ``outputs`` also appears in - ``reference_outputs`` (no unexpected tool calls). - - ``"superset"``: every tool call in ``reference_outputs`` also appears - in ``outputs`` (no missing tool calls; extras are allowed). - """ - - __test__ = False - - def __init__(self, trajectory_match_mode: TrajectoryMatchMode = "strict") -> None: - if trajectory_match_mode not in _TRAJECTORY_MATCH_MODES: - raise ValueError( - f"Invalid trajectory_match_mode: {trajectory_match_mode!r}. Must be one of {_TRAJECTORY_MATCH_MODES!r}." - ) - self.trajectory_match_mode: TrajectoryMatchMode = trajectory_match_mode - - def record(self) -> TrajectoryJudgeBinding: - return TrajectoryJudgeBinding(self.trajectory_match_mode) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index d456234c..639fda71 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -5,7 +5,9 @@ import hashlib import inspect import json +import re import sys +import time from collections.abc import AsyncIterator from pathlib import Path from typing import TYPE_CHECKING, Any, Callable @@ -68,12 +70,12 @@ def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._responses = responses def __enter__(self) -> None: - from .model_builder import Ai + from .ai import Ai Ai.fake(self._agent_cls.__name__, self._responses) def __exit__(self, *_exc: Any) -> bool: - from .model_builder import Ai + from .ai import Ai Ai.forget(self._agent_cls.__name__) return False @@ -96,16 +98,64 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return wrapper +class ToolCallView: + """Ergonomic, attribute-style view of a raw ``tool_calls`` dict, passed + to ``assert_tool_called``'s predicate.""" + + def __init__(self, data: dict) -> None: + self.name = data.get("name", "") + self.args = data.get("args") or {} + self.id = data.get("id") + self._data = data + + def __repr__(self) -> str: + return f"ToolCallView(name={self.name!r}, args={self.args!r})" + + class RecordingAgent(_Recorder): - def __init__(self, real: Agent, cassette: str | None = None) -> None: + """Bound as ``agent`` by ``with Agent.record(cassette) as agent:``. + + Fluent testing handle around a record-and-replay session: ``prompt()`` + is async (it's the same real agent call underneath, just cached) and + each call mutates the handle's "current turn" state, which the + ``assert_*`` methods judge against — mirroring how a browser-testing + ``page`` object exposes assertions against the current page state. + + On a cassette miss, the real agent is called once and the response is + cached to disk (keyed by the conversation history so far, plus the new + message, so two sessions with different histories but the same latest + message text don't collide). On a hit, it's replayed with no live call. + """ + + def __init__(self, real: Agent, cassette: str | None = None, messages: list | None = None) -> None: super().__init__() self._real = real self.cassette: Path | None = Path(cassette) if cassette else None + self._seed_messages: list = list(messages or []) + self._transcript: list[dict] = [] + self._real.messages = self._history # type: ignore[method-assign] + self.last_response: AgentResponse | None = None + self.last_elapsed: float | None = None + + def _history(self) -> list: + return self._seed_messages + self._transcript @staticmethod - def _key(message: str, attachments: list[Document] | None) -> str: + def _serialize(value: Any) -> Any: + if isinstance(value, dict): + return value + return {"type": type(value).__name__, "content": getattr(value, "content", str(value))} + + def _key(self, message: str, attachments: list[Document] | None) -> str: names = [getattr(doc, "name", "") for doc in (attachments or [])] - payload = json.dumps({"message": message, "attachments": names}, sort_keys=True) + payload = json.dumps( + { + "history": [self._serialize(m) for m in self._history()], + "message": message, + "attachments": names, + }, + sort_keys=True, + ) return hashlib.sha256(payload.encode()).hexdigest() def _load(self) -> tuple[Path, dict]: @@ -118,14 +168,38 @@ def _save(self, cassette: Path, store: dict, key: str, value: Any) -> None: cassette.parent.mkdir(parents=True, exist_ok=True) cassette.write_text(json.dumps(store, indent=2, sort_keys=True)) - async def prompt(self, message: str, attachments: list[Document] | None = None) -> AgentResponse: + @staticmethod + def _cache_prompt_value(response: AgentResponse) -> dict: + return {"content": response.content, "tool_calls": response.tool_calls} + + @staticmethod + def _response_from_cache(value: Any) -> AgentResponse: + if isinstance(value, dict) and "content" in value: + return AgentResponse(content=_joined(value.get("content", "")), tool_calls=value.get("tool_calls") or []) + return AgentResponse(content=_joined(value)) + + def _remember_turn(self, message: str, response: AgentResponse) -> None: + self._transcript.append({"role": "user", "content": message}) + turn: dict[str, Any] = {"role": "assistant", "content": response.content} + if response.tool_calls: + turn["tool_calls"] = response.tool_calls + self._transcript.append(turn) + + async def prompt(self, message: str, *, attachments: list[Document] | None = None) -> AgentResponse: + """Run (or replay) one turn and make it the "current" response that + assert_*() methods judge.""" self._record_call(message, attachments) cassette, store = self._load() key = self._key(message, attachments) + start = time.monotonic() if key in store: - return AgentResponse(content=_joined(store[key])) - response = await self._real._run(message, attachments=attachments) - self._save(cassette, store, key, response.content) + response = self._response_from_cache(store[key]) + else: + response = await self._real._run(message, attachments=attachments) + self._save(cassette, store, key, self._cache_prompt_value(response)) + self.last_elapsed = time.monotonic() - start + self.last_response = response + self._remember_turn(message, response) return response async def stream(self, message: str) -> AsyncIterator[str]: @@ -142,6 +216,78 @@ async def stream(self, message: str) -> AsyncIterator[str]: for chunk in chunks: yield chunk + def _require_response(self) -> AgentResponse: + assert self.last_response is not None, "No prompt() call has been made yet." + return self.last_response + + def _tool_call_names(self) -> list[str]: + return [tc.get("name", "") for tc in self._require_response().tool_calls] + + def assert_text_response(self) -> None: + response = self._require_response() + assert response.content, "Expected a non-empty text response, but content was empty." + + def assert_tool_called(self, name: str, predicate: Callable[[ToolCallView], bool] | None = None) -> None: + response = self._require_response() + matches = [tc for tc in response.tool_calls if tc.get("name") == name] + assert matches, f"Expected tool {name!r} to be called, but it wasn't. Called: {self._tool_call_names()}" + if predicate is not None: + assert any(predicate(ToolCallView(tc)) for tc in matches), ( + f"Tool {name!r} was called, but no call satisfied the given predicate." + ) + + def assert_tool_not_called(self, names: list[str]) -> None: + unexpected = set(self._tool_call_names()) & set(names) + assert not unexpected, f"Expected tools {sorted(names)} not to be called, but got: {sorted(unexpected)}" + + def assert_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: + response = self._require_response() + verdict = self._judge(model, expectation, response.content) + 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: + cassette, store = self._load() + key = self._judge_key(model, expectation, content) + if key in store: + return store[key] + verdict = self._judge_live(model, expectation, content) + self._save(cassette, store, key, verdict) + return verdict + + @staticmethod + def _judge_key(model: str, expectation: str, content: str) -> str: + payload = json.dumps( + {"judge_model": model, "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 + + 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", "")} + class AgentBinding: def __init__(self, agent_cls: type[Agent], stand_in: Any) -> None: diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 2d63867a..145a1cfb 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -7,7 +7,7 @@ from fastapi_startkit.ai import AIConfig, Document, fake_chat_model from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.model_builder import Ai +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse from fastapi_startkit.application import app diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 7b560d74..42ff7adb 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -16,7 +16,7 @@ from unittest import mock from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.model_builder import Ai +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse @@ -204,7 +204,8 @@ async def test_first_run_records_response_to_cassette(self): self.assertEqual(calls, ["hello"]) self.assertTrue(os.path.exists(cassette)) with open(cassette) as f: - self.assertIn("recorded reply", json.load(f).values()) + store = json.load(f) + self.assertTrue(any(v.get("content") == "recorded reply" for v in store.values())) async def test_second_run_replays_without_calling_run(self): calls = self.setup_agent("recorded reply") diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py new file mode 100644 index 00000000..30d12989 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_agent_record_fluent.py @@ -0,0 +1,354 @@ +"""Tests for the fluent Agent.record() testing DSL. + +``with Agent.record(cassette) as agent:`` binds a ``RecordingAgent`` handle +whose async ``prompt()`` and assertion methods judge the most recent turn — +mirroring how a browser-testing ``page`` object exposes assertions against +current page state: + + with RouterAgent.record("cassette.json") as agent: + await agent.prompt("hello") + agent.assert_text_response() + agent.assert_tool_not_called(["job_search_tool"]) + agent.assert_response_time_lt(5) + + await agent.prompt("suggest python developer jobs") + agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") +""" + +import os +import tempfile +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 +from fastapi_startkit.ai.response import AgentResponse +from fastapi_startkit.ai.testing import RecordingAgent + + +class SimpleAgent(Agent): + pass + + +def _tool_call(name: str, args: dict | None = None, call_id: str = "c1") -> dict: + return {"name": name, "args": args or {}, "id": call_id} + + +class TestFluentPromptMechanics(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, responses: list): + """responses: list of (content, tool_calls) tuples, consumed in call order.""" + queue = list(responses) + + async def fake_run(agent_self, message, **kwargs): + content, tool_calls = queue.pop(0) + return AgentResponse(content=content, tool_calls=tool_calls or []) + + patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_prompt_is_async_and_returns_agent_response(self): + self.setup_agent([("Hello there!", [])]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + response = await agent.prompt("hi") + + self.assertIsInstance(response, AgentResponse) + self.assertEqual(response.content, "Hello there!") + + async def test_second_prompt_continues_the_same_session(self): + self.setup_agent([("Hi!", []), ("here are some jobs", [_tool_call("job_search_tool")])]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + agent.assert_text_response() + + await agent.prompt("suggest python developer jobs") + agent.assert_tool_called("job_search_tool") + + async def test_replaying_from_cassette_preserves_tool_calls(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + self.setup_agent([("", [_tool_call("job_search_tool", {"q": "python"})])]) + with SimpleAgent.record(cassette) as agent: + await agent.prompt("find jobs") + + # No queued responses left — this must replay from cassette, not call _run again. + with SimpleAgent.record(cassette) as agent: + await agent.prompt("find jobs") + agent.assert_tool_called("job_search_tool") + + +class TestAssertTextResponse(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, content, tool_calls=None): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content=content, tool_calls=tool_calls or []) + + patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_content_present(self): + self.setup_agent("Hi!") + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + agent.assert_text_response() + + async def test_fails_on_empty_content(self): + self.setup_agent("", tool_calls=[_tool_call("search")]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + with self.assertRaises(AssertionError): + agent.assert_text_response() + + 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_text_response() + + +class TestAssertToolCalled(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, tool_calls): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="", tool_calls=tool_calls) + + patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_tool_present(self): + self.setup_agent([_tool_call("job_search_tool", {"q": "python"})]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + agent.assert_tool_called("job_search_tool") + + async def test_fails_when_tool_absent(self): + self.setup_agent([]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + with self.assertRaises(AssertionError): + agent.assert_tool_called("job_search_tool") + + async def test_predicate_can_accept_via_attribute_access(self): + self.setup_agent([_tool_call("job_search_tool", {"q": "python"})]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") + + async def test_predicate_can_reject(self): + self.setup_agent([_tool_call("job_search_tool", {"q": "python"})]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + with self.assertRaises(AssertionError): + agent.assert_tool_called("job_search_tool", lambda tool: tool.args.get("q") == "java") + + +class TestAssertToolNotCalled(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, tool_calls): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="Hello!", tool_calls=tool_calls) + + patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_absent(self): + self.setup_agent([]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + agent.assert_tool_not_called(["job_search_tool"]) + + async def test_fails_when_present(self): + self.setup_agent([_tool_call("job_search_tool")]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + with self.assertRaises(AssertionError): + agent.assert_tool_not_called(["job_search_tool"]) + + +class TestAssertResponseTimeLt(unittest.IsolatedAsyncioTestCase): + def setup_agent(self): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="Hello!") + + patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_for_a_fast_call(self): + self.setup_agent() + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + agent.assert_response_time_lt(5) + + async def test_fails_when_exceeded(self): + self.setup_agent() + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + with self.assertRaises(AssertionError): + agent.assert_response_time_lt(0) + + 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_time_lt(5) + + +class TestRecordMessagesSeed(unittest.IsolatedAsyncioTestCase): + async def test_seed_messages_are_included_when_building_the_real_agents_messages(self): + seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json"), messages=seed) as agent: + built = agent._real._build_messages("suggest python developer jobs") + + self.assertEqual(built[0], seed[0]) + self.assertEqual(built[1], seed[1]) + self.assertEqual(built[-1], {"role": "user", "content": "suggest python developer jobs"}) + + async def test_same_followup_text_with_different_seed_history_does_not_collide(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "shared.json") + + async def run_a(agent_self, message, **kwargs): + return AgentResponse(content="job list A") + + with mock.patch.object(SimpleAgent, "_run", run_a): + with SimpleAgent.record(cassette) as agent: + response_a = await agent.prompt("suggest python developer jobs") + + async def run_b(agent_self, message, **kwargs): + return AgentResponse(content="job list B") + + seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] + with mock.patch.object(SimpleAgent, "_run", run_b): + with SimpleAgent.record(cassette, messages=seed) as agent: + response_b = await agent.prompt("suggest python developer jobs") + + self.assertEqual(response_a.content, "job list A") + self.assertEqual(response_b.content, "job list B") + + +class TestAssertResponseJudged(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, content): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content=content) + + patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + 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"} + ): + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + 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"} + ): + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + with self.assertRaises(AssertionError): + 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"}) + 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") + + 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 SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") + + judge = mock.Mock(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") + + 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) + + agent = RecordingAgent(SimpleAgent()) + verdict = agent._judge_live("gpt-3.5-turbo", "The llm should respond with greetings", "Hello there!") + + self.assertTrue(verdict["passed"]) + self.assertIn("Greets", verdict["reasoning"]) + self.assertIn("Hello there!", captured["prompt"]) + + +class TestExistingRecordApiIsUnaffected(unittest.IsolatedAsyncioTestCase): + """The pre-existing bare-context-manager Agent.record() usage (task #327) + must keep working unchanged alongside the new fluent handle.""" + + async def test_bare_context_manager_prompt_still_works(self): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="recorded reply") + + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with mock.patch.object(SimpleAgent, "_run", fake_run): + with SimpleAgent.record(cassette): + result = await SimpleAgent().prompt("hello") + + self.assertEqual(result.content, "recorded reply") diff --git a/fastapi_startkit/tests/ai/test_agent_schema.py b/fastapi_startkit/tests/ai/test_agent_schema.py index 1a30772c..40306ea3 100644 --- a/fastapi_startkit/tests/ai/test_agent_schema.py +++ b/fastapi_startkit/tests/ai/test_agent_schema.py @@ -6,7 +6,7 @@ from pydantic import BaseModel from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.model_builder import Ai +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse diff --git a/fastapi_startkit/tests/ai/test_ai_fake.py b/fastapi_startkit/tests/ai/test_ai_fake.py index 30b2e39e..33a29247 100644 --- a/fastapi_startkit/tests/ai/test_ai_fake.py +++ b/fastapi_startkit/tests/ai/test_ai_fake.py @@ -16,7 +16,7 @@ from langchain_core.tools import tool from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.model_builder import Ai +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.application import app diff --git a/fastapi_startkit/tests/ai/test_evals.py b/fastapi_startkit/tests/ai/test_evals.py deleted file mode 100644 index 5e7e2ccc..00000000 --- a/fastapi_startkit/tests/ai/test_evals.py +++ /dev/null @@ -1,281 +0,0 @@ -"""Tests for TestJudgeAgent — the deterministic trajectory-match evaluator. - -TestJudgeAgent compares an actual agent message trajectory against a reference -trajectory with no live LLM calls: it is pure, deterministic structural -comparison, driven entirely by data the tests construct themselves (typically -the output of Agent.fake()). Mirrors the Agent.fake()/Agent.record() testing -DSL: ``.record()`` returns a context manager yielding a callable evaluator. -""" - -import unittest - -from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage - -from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.evals import TestJudgeAgent -from fastapi_startkit.ai.response import AgentResponse - - -def _weather_trajectory(city: str = "San Francisco") -> list: - return [ - HumanMessage(content=f"What is the weather in {city}?"), - AIMessage( - content="", - tool_calls=[{"id": "call_1", "name": "get_weather", "args": {"city": city}}], - ), - ToolMessage(content="It is 75 degrees and sunny.", tool_call_id="call_1"), - AIMessage(content=f"The weather in {city} is 75 degrees and sunny."), - ] - - -class TestJudgeAgentConstruction(unittest.TestCase): - def test_defaults_to_strict_mode(self): - judge = TestJudgeAgent() - self.assertEqual(judge.trajectory_match_mode, "strict") - - def test_accepts_each_supported_mode(self): - for mode in ("strict", "unordered", "subset", "superset"): - judge = TestJudgeAgent(trajectory_match_mode=mode) - self.assertEqual(judge.trajectory_match_mode, mode) - - def test_rejects_unknown_mode(self): - with self.assertRaises(ValueError): - TestJudgeAgent(trajectory_match_mode="fuzzy") - - -class TestJudgeAgentRecordContextManager(unittest.TestCase): - def test_record_yields_callable_evaluator(self): - with TestJudgeAgent().record() as evaluator: - self.assertTrue(callable(evaluator)) - - def test_evaluation_result_is_subscriptable_with_score_key(self): - trajectory = _weather_trajectory() - with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: - evaluation = evaluator(outputs=trajectory, reference_outputs=trajectory) - - self.assertIs(evaluation["score"], True) - self.assertEqual(evaluation["key"], "trajectory_strict_match") - self.assertIn("comment", evaluation) - - -class TestStrictTrajectoryMatch(unittest.TestCase): - def setUp(self): - self.judge = TestJudgeAgent(trajectory_match_mode="strict") - - def _evaluate(self, outputs, reference_outputs): - with self.judge.record() as evaluator: - return evaluator(outputs=outputs, reference_outputs=reference_outputs) - - def test_identical_trajectory_matches(self): - trajectory = _weather_trajectory() - evaluation = self._evaluate(trajectory, trajectory) - self.assertIs(evaluation["score"], True) - - def test_matches_even_when_final_message_content_differs(self): - reference = _weather_trajectory() - outputs = _weather_trajectory() - outputs[-1] = AIMessage(content="Completely different wording, still sunny.") - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], True) - - def test_fails_when_tool_call_args_differ(self): - reference = _weather_trajectory(city="San Francisco") - outputs = _weather_trajectory(city="Oakland") - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - def test_fails_on_extra_trailing_message(self): - reference = _weather_trajectory() - outputs = _weather_trajectory() + [AIMessage(content="One more thing...")] - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - def test_fails_on_missing_message(self): - reference = _weather_trajectory() - outputs = _weather_trajectory()[:-1] - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - def test_fails_when_role_order_differs(self): - reference = _weather_trajectory() - outputs = list(reversed(_weather_trajectory())) - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - def test_fails_when_one_side_has_tool_calls_and_other_does_not(self): - reference = _weather_trajectory() - outputs = _weather_trajectory() - outputs[1] = AIMessage(content="") - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - def test_works_with_plain_dict_messages(self): - reference = _weather_trajectory() - outputs = [ - {"role": "user", "content": "What is the weather in San Francisco?"}, - { - "role": "assistant", - "content": "", - "tool_calls": [{"id": "call_1", "name": "get_weather", "args": {"city": "San Francisco"}}], - }, - {"role": "tool", "content": "It is 75 degrees and sunny.", "tool_call_id": "call_1"}, - {"role": "assistant", "content": "The weather in San Francisco is 75 degrees and sunny."}, - ] - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], True) - - -class TestUnorderedTrajectoryMatch(unittest.TestCase): - def setUp(self): - self.judge = TestJudgeAgent(trajectory_match_mode="unordered") - - def _evaluate(self, outputs, reference_outputs): - with self.judge.record() as evaluator: - return evaluator(outputs=outputs, reference_outputs=reference_outputs) - - def test_same_tool_calls_in_different_message_order_matches(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [HumanMessage(content="hi"), weather_call, time_call] - outputs = [HumanMessage(content="hi"), time_call, weather_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], True) - - def test_missing_tool_call_fails(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [weather_call, time_call] - outputs = [weather_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - def test_extra_tool_call_fails(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [weather_call] - outputs = [weather_call, time_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - -class TestSubsetTrajectoryMatch(unittest.TestCase): - def setUp(self): - self.judge = TestJudgeAgent(trajectory_match_mode="subset") - - def _evaluate(self, outputs, reference_outputs): - with self.judge.record() as evaluator: - return evaluator(outputs=outputs, reference_outputs=reference_outputs) - - def test_output_tool_calls_within_reference_matches(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [weather_call, time_call] - outputs = [weather_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], True) - - def test_output_tool_call_not_in_reference_fails(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [weather_call] - outputs = [weather_call, time_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - -class TestSupersetTrajectoryMatch(unittest.TestCase): - def setUp(self): - self.judge = TestJudgeAgent(trajectory_match_mode="superset") - - def _evaluate(self, outputs, reference_outputs): - with self.judge.record() as evaluator: - return evaluator(outputs=outputs, reference_outputs=reference_outputs) - - def test_output_contains_all_reference_tool_calls_plus_extra_matches(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [weather_call] - outputs = [weather_call, time_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], True) - - def test_output_missing_a_required_reference_tool_call_fails(self): - weather_call = AIMessage(content="", tool_calls=[{"id": "1", "name": "get_weather", "args": {"city": "SF"}}]) - time_call = AIMessage(content="", tool_calls=[{"id": "2", "name": "get_time", "args": {"tz": "PT"}}]) - - reference = [weather_call, time_call] - outputs = [weather_call] - - evaluation = self._evaluate(outputs, reference) - self.assertIs(evaluation["score"], False) - - -class SimpleAgent(Agent): - pass - - -class TestJudgeAgentWithAgentFakeOutput(unittest.TestCase): - """Demonstrates the intended workflow: Agent.fake() drives deterministic - output with no live LLM call, and TestJudgeAgent judges the resulting - trajectory against a reference.""" - - async def _fake_response(self) -> AgentResponse: - agent = SimpleAgent() - with SimpleAgent.fake(["It is sunny in San Francisco."]): - return await agent.prompt("What is the weather in San Francisco?") - - def test_fake_agent_reply_matches_reference_trajectory_in_strict_mode(self): - import asyncio - - response = asyncio.run(self._fake_response()) - - outputs = [ - HumanMessage(content="What is the weather in San Francisco?"), - AIMessage(content=response.content), - ] - reference_trajectory = [ - HumanMessage(content="What is the weather in San Francisco?"), - AIMessage(content="It is sunny in San Francisco."), - ] - - with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: - evaluation = evaluator(outputs=outputs, reference_outputs=reference_trajectory) - - self.assertIs(evaluation["score"], True) - - -class TestJudgeAgentMessagesDictInput(unittest.TestCase): - def test_accepts_outputs_dict_with_messages_key(self): - trajectory = _weather_trajectory() - with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: - evaluation = evaluator(outputs={"messages": trajectory}, reference_outputs=trajectory) - - self.assertIs(evaluation["score"], True) - - -class TestJudgeAgentIsNotCollectedByPytest(unittest.TestCase): - def test_has_test_dunder_set_to_false(self): - self.assertFalse(getattr(TestJudgeAgent, "__test__")) - - -class TestSystemMessagesAreComparedByRole(unittest.TestCase): - def test_strict_mode_compares_system_message_role(self): - reference = [SystemMessage(content="Be concise."), HumanMessage(content="hi")] - outputs = [SystemMessage(content="Be very concise."), HumanMessage(content="hi")] - - with TestJudgeAgent(trajectory_match_mode="strict").record() as evaluator: - evaluation = evaluator(outputs=outputs, reference_outputs=reference) - - self.assertIs(evaluation["score"], True) From 05647eeaaa0ae21bf2cdd5a0553d16ded52ef2cd Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 11:00:01 -0700 Subject: [PATCH 3/7] 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. --- .../tests/units/agents/test_router_agent.py | 8 +- .../src/fastapi_startkit/ai/__init__.py | 2 + .../src/fastapi_startkit/ai/agent.py | 20 +- .../src/fastapi_startkit/ai/ai.py | 35 +-- .../src/fastapi_startkit/ai/judge.py | 34 +++ .../src/fastapi_startkit/ai/runner.py | 3 + .../src/fastapi_startkit/ai/testing.py | 38 ++-- .../tests/ai/test_agent_record_fluent.py | 62 +++--- fastapi_startkit/tests/ai/test_judge_agent.py | 79 +++++++ .../tests/ai/test_structured_output.py | 199 ++++++++++++++++++ 10 files changed, 400 insertions(+), 80 deletions(-) create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/judge.py create mode 100644 fastapi_startkit/tests/ai/test_judge_agent.py create mode 100644 fastapi_startkit/tests/ai/test_structured_output.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/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..c7433e0a 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,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) 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..6923bd3d --- /dev/null +++ b/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": ""}' + ) + + 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/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/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 639fda71..c2fed401 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,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": ""}' - ) - 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: 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..e1d5a964 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -0,0 +1,79 @@ +"""Tests for JudgeAgent — grades a response against an expectation. + +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 os +import tempfile +import unittest +from unittest import mock + +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_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_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"]) + + 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 + + async def ainvoke(self, messages): + seen["messages"] = messages + return AIMessage(content='{"passed": true, "reasoning": "ok"}') + + with mock.patch.object(JudgeAgent, "_build_model", lambda self, *a, **k: Capturing()): + await JudgeAgent().judge("The llm should respond with greetings", "Hello there!") + + 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.""" + judge = JudgeAgent() + judge.model = "gpt-4o-mini" + judge.provider = "openai" + + self.assertEqual(judge.model, "gpt-4o-mini") + self.assertEqual(judge.provider, "openai") + + def test_model_and_provider_default_to_agent_defaults(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)) 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..bfa60a46 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -0,0 +1,199 @@ +"""Structured output. + +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. +""" + +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 that echoes its query.""" + return query + + +class ToolMovieAgent(Agent): + def schema(self): + return Movie + + def tools(self): + return [noop] + + +def _real_tool_call(**args) -> dict: + return {"name": "noop", "args": args, "id": "1", "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): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + + def tearDown(self): + Ai.reset_fakes() + + def _patch(self, fake): + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: fake) + patcher.start() + self.addCleanup(patcher.stop) + + def test_schema_wraps_model_with_structured_output(self): + fake = _FakeModel() + self._patch(fake) + + result = Ai().build(MovieAgent()) + + self.assertEqual(result, "STRUCTURED") + self.assertEqual(fake.calls, [("with_structured_output", Movie, {"include_raw": True})]) + + def test_tools_are_bound_then_structured_output_is_applied(self): + fake = _FakeModel() + self._patch(fake) + + result = Ai().build(ToolMovieAgent()) + + self.assertEqual(result, "STRUCTURED") + self.assertEqual( + fake.calls, + [("bind_tools", [noop]), ("with_structured_output", Movie, {"include_raw": True})], + ) + + def test_tools_only_binds_tools_without_structured_output(self): + fake = _FakeModel() + self._patch(fake) + + class ToolAgent(Agent): + def tools(self): + return [noop] + + result = Ai().build(ToolAgent()) + + self.assertIs(result, fake) + self.assertEqual(fake.calls, [("bind_tools", [noop])]) + + def test_streaming_skips_structured_output(self): + fake = _FakeModel() + self._patch(fake) + + result = Ai().build(ToolMovieAgent(), structured=False) + + self.assertIs(result, fake) + self.assertEqual(fake.calls, [("bind_tools", [noop])]) + + def test_no_schema_no_tools_returns_the_plain_model(self): + fake = _FakeModel() + self._patch(fake) + + self.assertIs(Ai().build(Agent()), fake) + self.assertEqual(fake.calls, []) + + +class TestRunner(unittest.IsolatedAsyncioTestCase): + 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) + + 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"]) + + self.assertEqual(result.content, "hello") + + +class TestResponseMapping(unittest.TestCase): + def test_unwraps_include_raw_into_parsed_and_content(self): + parsed = Movie(title="Inception", year=2010) + + 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()) + self.assertEqual(response.tool_calls, []) + + +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() + + 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 cf6bf43c43b143b513b27c6f1af789b4831f2418 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 11:01:23 -0700 Subject: [PATCH 4/7] refactor: move utils/structures.py into support/ utils/structures.py was used (Configuration.data(), Loader.load()), so per the foundation/support consolidation it belongs in support/ alongside the rest of the framework's internal helpers. The utils/ directory had no __init__.py and nothing else referenced it, so it's gone entirely now that its one file has moved. Updated the two import sites accordingly. --- .../src/fastapi_startkit/configuration/Configuration.py | 2 +- fastapi_startkit/src/fastapi_startkit/loader/Loader.py | 2 +- .../src/fastapi_startkit/{utils => support}/structures.py | 0 3 files changed, 2 insertions(+), 2 deletions(-) rename fastapi_startkit/src/fastapi_startkit/{utils => support}/structures.py (100%) diff --git a/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py b/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py index 44fc49cc..5f32852b 100644 --- a/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py +++ b/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py @@ -1,5 +1,5 @@ from fastapi_startkit.loader import Loader -from ..utils.structures import data +from ..support.structures import data from ..exceptions import InvalidConfigurationSetup diff --git a/fastapi_startkit/src/fastapi_startkit/loader/Loader.py b/fastapi_startkit/src/fastapi_startkit/loader/Loader.py index 07b01cde..4b34dfbb 100644 --- a/fastapi_startkit/src/fastapi_startkit/loader/Loader.py +++ b/fastapi_startkit/src/fastapi_startkit/loader/Loader.py @@ -4,7 +4,7 @@ import pkgutil from ..exceptions import LoaderNotFound -from ..utils.structures import load +from ..support.structures import load def parameters_filter(obj_name, obj): diff --git a/fastapi_startkit/src/fastapi_startkit/utils/structures.py b/fastapi_startkit/src/fastapi_startkit/support/structures.py similarity index 100% rename from fastapi_startkit/src/fastapi_startkit/utils/structures.py rename to fastapi_startkit/src/fastapi_startkit/support/structures.py From 0b7b7f4632ecb2b353bab817f20564bba762632c Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 13:00:20 -0700 Subject: [PATCH 5/7] refactor(ai): drop AgentBinding, have fake()/record() return their fakes directly Agent.fake()/record() now return AgentFake/AgentRecordFake directly instead of wrapping the latter in an AgentBinding indirection layer. AgentRecordFake gains the container-binding, cassette-resolution, and decorator behavior AgentBinding used to provide, so `with Agent.record(...) as agent:` and the `@Agent.record(...)` decorator form both keep working unchanged. AgentBinding is removed entirely, along with every reference to it (imports, __init__ exports, type hints, tests). --- example/agents/.gitignore | 3 ++ .../tests/units/agents/record_stream.json | 6 +++ fastapi_startkit/pyproject.toml | 2 +- .../src/fastapi_startkit/ai/__init__.py | 7 ++-- .../src/fastapi_startkit/ai/agent.py | 15 ++++---- .../src/fastapi_startkit/ai/testing.py | 37 ++++++++----------- .../tests/ai/test_agent_record_fluent.py | 16 ++++---- 7 files changed, 44 insertions(+), 42 deletions(-) create mode 100644 example/agents/tests/units/agents/record_stream.json diff --git a/example/agents/.gitignore b/example/agents/.gitignore index b2db17eb..014a22ef 100644 --- a/example/agents/.gitignore +++ b/example/agents/.gitignore @@ -7,3 +7,6 @@ storage node_modules /public/build /public/hot +.ruff_cache +.pytest_cache +.ai diff --git a/example/agents/tests/units/agents/record_stream.json b/example/agents/tests/units/agents/record_stream.json new file mode 100644 index 00000000..1ee25800 --- /dev/null +++ b/example/agents/tests/units/agents/record_stream.json @@ -0,0 +1,6 @@ +{ + "034751ef1322a12f3406dad43313f9cdb31f4ea85d9452c22bf7529ad8e92614": { + "content": "Hello! How can I help you today?", + "tool_calls": [] + } +} \ No newline at end of file diff --git a/fastapi_startkit/pyproject.toml b/fastapi_startkit/pyproject.toml index b7f66198..a7717034 100644 --- a/fastapi_startkit/pyproject.toml +++ b/fastapi_startkit/pyproject.toml @@ -129,7 +129,7 @@ exclude = [ "**/__pycache__", "**/.venv", ] -typeCheckingMode = "basic" +typeCheckingMode = "standard" pythonVersion = "3.12" [tool.pytest.ini_options] diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index 47abea1b..db2f5dad 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -13,21 +13,20 @@ from .judge import JudgeAgent from .providers.ai_provider import AIProvider from .response import AgentResponse, AgentSnapshot -from .testing import AgentBinding, AgentModelFake, RecordingAgent, ToolCallView +from .testing import AgentFake, AgentRecordFake, ToolCallView __all__ = [ "Agent", "Ai", "Middleware", - "AgentBinding", - "AgentModelFake", + "AgentFake", "AgentResponse", "AgentSnapshot", "AIConfig", "AIProvider", "AnthropicConfig", "JudgeAgent", - "RecordingAgent", + "AgentRecordFake", "ToolCallView", "Audio", "AudioResponse", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 1931cd60..f13be05a 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -4,12 +4,11 @@ from .document import Document from .response import AgentResponse -from .testing import AgentBinding if TYPE_CHECKING: from langchain_core.tools import BaseTool - from .testing import AgentModelFake + from .testing import AgentFake, AgentRecordFake class Agent: @@ -81,16 +80,16 @@ async def stream( yield chunk @classmethod - def fake(cls, responses: list) -> "AgentModelFake": - from .testing import AgentModelFake + def fake(cls, responses: list) -> "AgentFake": + from .testing import AgentFake - return AgentModelFake(cls, responses) + return AgentFake(cls, responses) @classmethod - def record(cls, cassette: str | None = None, messages: list | None = None) -> "AgentBinding": - from .testing import AgentBinding, RecordingAgent + def record(cls, cassette: str | None = None, messages: list | None = None) -> "AgentRecordFake": + from .testing import AgentRecordFake - return AgentBinding(cls, RecordingAgent(cls(), cassette, messages)) + return AgentRecordFake(cls(), cassette, messages) @classmethod def _binding(cls) -> Any: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index c2fed401..83120191 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -55,11 +55,10 @@ def _joined(value: Any) -> str: return "".join(value) if isinstance(value, list) else value -class AgentModelFake: +class AgentFake: """Registers a fixed, ordered list of replies as ``agent_cls``'s chat model for the duration of a ``with`` block (or a decorated function). - Unlike the old pattern-matching stand-in, this swaps only the model — ``prompt()``/``stream()`` still run the real message-building, pipeline, and tool-execution path; see ``Ai.fake()``. """ @@ -111,7 +110,7 @@ def __repr__(self) -> str: return f"ToolCallView(name={self.name!r}, args={self.args!r})" -class RecordingAgent(_Recorder): +class AgentRecordFake(_Recorder): """Bound as ``agent`` by ``with Agent.record(cassette) as agent:``. Fluent testing handle around a record-and-replay session: ``prompt()`` @@ -124,6 +123,11 @@ class RecordingAgent(_Recorder): cached to disk (keyed by the conversation history so far, plus the new message, so two sessions with different histories but the same latest message text don't collide). On a hit, it's replayed with no live call. + + Entering the ``with`` block also binds this handle into the container + under the agent class's name, so any other instance of that class + created during the block (e.g. by application code under test) is + routed through the same recording session. """ def __init__(self, real: Agent, cassette: str | None = None, messages: list | None = None) -> None: @@ -159,7 +163,7 @@ def _key(self, message: str, attachments: list[Document] | None) -> str: def _load(self) -> tuple[Path, dict]: cassette = self.cassette - assert cassette is not None, "RecordingAgent has no cassette resolved" + assert cassette is not None, "AgentRecordFake has no cassette resolved" return cassette, (json.loads(cassette.read_text()) if cassette.exists() else {}) def _save(self, cassette: Path, store: dict, key: str, value: Any) -> None: @@ -276,34 +280,25 @@ async def _judge_live(self, model: str, expectation: str, content: str, provider judge.provider = provider return await judge.judge(expectation, content) - -class AgentBinding: - def __init__(self, agent_cls: type[Agent], stand_in: Any) -> None: - self._agent_cls = agent_cls - self._stand_in = stand_in - def _resolve_cassette(self, filename: str, qualname: str) -> None: - stand_in = self._stand_in - if not isinstance(stand_in, RecordingAgent): - return here = Path(filename).parent - if stand_in.cassette is None: - stand_in.cassette = here / "cassettes" / f"{qualname.replace('.', '_')}.json" - elif not stand_in.cassette.is_absolute(): - stand_in.cassette = here / stand_in.cassette + if self.cassette is None: + self.cassette = here / "cassettes" / f"{qualname.replace('.', '_')}.json" + elif not self.cassette.is_absolute(): + self.cassette = here / self.cassette - def __enter__(self) -> Any: + def __enter__(self) -> "AgentRecordFake": from fastapi_startkit.application import app caller = sys._getframe(1).f_code self._resolve_cassette(caller.co_filename, caller.co_qualname) - app().bind(self._agent_cls.__name__, self._stand_in) - return self._stand_in + app().bind(type(self._real).__name__, self) + return self def __exit__(self, *_exc: Any) -> bool: from fastapi_startkit.application import app - app().unbind(self._agent_cls.__name__) + app().unbind(type(self._real).__name__) return False def __call__(self, func: Callable) -> Callable: diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py index 0f2c42b2..45121c91 100644 --- a/fastapi_startkit/tests/ai/test_agent_record_fluent.py +++ b/fastapi_startkit/tests/ai/test_agent_record_fluent.py @@ -1,6 +1,6 @@ """Tests for the fluent Agent.record() testing DSL. -``with Agent.record(cassette) as agent:`` binds a ``RecordingAgent`` handle +``with Agent.record(cassette) as agent:`` binds an ``AgentRecordFake`` handle whose async ``prompt()`` and assertion methods judge the most recent turn — mirroring how a browser-testing ``page`` object exposes assertions against current page state: @@ -24,7 +24,7 @@ from fastapi_startkit.ai.agent import Agent from fastapi_startkit.ai.response import AgentResponse -from fastapi_startkit.ai.testing import RecordingAgent +from fastapi_startkit.ai.testing import AgentRecordFake class SimpleAgent(Agent): @@ -254,7 +254,7 @@ async def test_passes_when_judge_approves(self): self.setup_agent("Hello there, welcome!") with tempfile.TemporaryDirectory() as tmp: with mock.patch.object( - RecordingAgent, + AgentRecordFake, "_judge_live", mock.AsyncMock(return_value={"passed": True, "reasoning": "greets the user"}), ): @@ -268,7 +268,7 @@ async def test_fails_when_judge_rejects(self): self.setup_agent("Completely unrelated content") with tempfile.TemporaryDirectory() as tmp: with mock.patch.object( - RecordingAgent, + AgentRecordFake, "_judge_live", mock.AsyncMock(return_value={"passed": False, "reasoning": "not a greeting"}), ): @@ -284,7 +284,7 @@ async def test_verdict_is_cached_in_the_cassette_and_not_re_judged(self): 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 mock.patch.object(AgentRecordFake, "_judge_live", judge): with SimpleAgent.record(cassette) as agent: await agent.prompt("hello") await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") @@ -297,14 +297,14 @@ async def test_verdict_persists_to_disk_for_a_later_replay(self): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with mock.patch.object( - RecordingAgent, "_judge_live", mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"}) + AgentRecordFake, "_judge_live", mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"}) ): with SimpleAgent.record(cassette) as agent: await agent.prompt("hello") await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") judge = mock.AsyncMock(side_effect=AssertionError("must not be called on replay")) - with mock.patch.object(RecordingAgent, "_judge_live", judge): + with mock.patch.object(AgentRecordFake, "_judge_live", judge): with SimpleAgent.record(cassette) as agent: await agent.prompt("hello") await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") @@ -321,7 +321,7 @@ 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 mock.patch.object(AgentRecordFake, "_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") From 89f3bae1896b17cdbc9b2a869af7e1c4a4b211ab Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 13:20:24 -0700 Subject: [PATCH 6/7] fix(agents): re-record stale/incomplete cassettes for the record() tests Two of the example/agents record() cassettes were recorded against message text that no longer matches the tests' current wording, so their cache keys never hit and every run fell through to a live Gemini call. The new units/agents/record_stream.json fixture had the same problem in the other direction: it only captured the first turn, so the second prompt() call and the assert_response_judged() verdict both missed the cassette too. Re-recorded all three cassettes against the real code path (mocking only the network boundary and the judge, same seam test_agent_record_fluent.py uses) so every prompt/stream/judge call in test_router_agent.py and test_chat_controller.py now replays from disk with no live model calls. uv.lock refreshes the editable fastapi-startkit metadata, which was stale at 0.45.0 against the package's actual 0.50.0. --- .../tests/features/record_no_stream.json | 5 +++- .../agents/tests/features/record_stream.json | 4 +-- .../tests/units/agents/record_stream.json | 28 +++++++++++++++++++ example/agents/uv.lock | 5 ++-- 4 files changed, 37 insertions(+), 5 deletions(-) diff --git a/example/agents/tests/features/record_no_stream.json b/example/agents/tests/features/record_no_stream.json index 37cca398..369f84ea 100644 --- a/example/agents/tests/features/record_no_stream.json +++ b/example/agents/tests/features/record_no_stream.json @@ -1,3 +1,6 @@ { - "1c27b2b038cae60eab7297ec0a808bfc67246c77a56a9e7c5ce1d9e5fdee0ba3": "Hi Alex, it's a pleasure to meet you! How can I help you today?" + "09e2753147eb813860891b4dd62050bceb302309d3f6db70019be2fb32582a47": { + "content": "Hi Alex, thanks for reaching out! How can I help you today?", + "tool_calls": [] + } } \ No newline at end of file diff --git a/example/agents/tests/features/record_stream.json b/example/agents/tests/features/record_stream.json index 764667ec..124252bd 100644 --- a/example/agents/tests/features/record_stream.json +++ b/example/agents/tests/features/record_stream.json @@ -1,6 +1,6 @@ { - "32e9c85d324f4cadef79b130717a28d0e23a85a2ac349ad2b1b78a233b31dbc4": [ + "97a14ff94f83842c3bf7a706f945aa10c9c6799ae651a74191ffc744f933a59f": [ "Hi", - " Bedram, it's a pleasure to assist you today! How may I help you?" + " Bedram, thanks for reaching out! How can I help you today?" ] } \ No newline at end of file diff --git a/example/agents/tests/units/agents/record_stream.json b/example/agents/tests/units/agents/record_stream.json index 1ee25800..82afdab8 100644 --- a/example/agents/tests/units/agents/record_stream.json +++ b/example/agents/tests/units/agents/record_stream.json @@ -2,5 +2,33 @@ "034751ef1322a12f3406dad43313f9cdb31f4ea85d9452c22bf7529ad8e92614": { "content": "Hello! How can I help you today?", "tool_calls": [] + }, + "358018dc096ce95b78b7561f562ba049848d03fb8b41953ed75c22c9ad20e228": { + "content": "", + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "call_1", + "name": "job_search_tool" + } + ] + }, + "b6579ba866634a66229ddbed0463c8c990eaaaf349385fbc67c2c2c1886085cd": { + "content": "", + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "call_1", + "name": "job_search_tool" + } + ] + }, + "judge:d33d8a7074bb7d980bf548694cf3b50ba853fa2f01a8df921bd28d7373ef5580": { + "passed": true, + "reasoning": "greets the user" } } \ No newline at end of file diff --git a/example/agents/uv.lock b/example/agents/uv.lock index 1edf626f..1d3ab1b2 100644 --- a/example/agents/uv.lock +++ b/example/agents/uv.lock @@ -417,7 +417,7 @@ wheels = [ [[package]] name = "fastapi-startkit" -version = "0.45.0" +version = "0.50.0" source = { editable = "../../fastapi_startkit" } dependencies = [ { name = "cleo" }, @@ -455,7 +455,7 @@ requires-dist = [ { name = "dotenv", specifier = ">=0.9.9" }, { name = "dotty-dict", specifier = ">=1.3.1" }, { name = "faker", marker = "extra == 'database'", specifier = ">=40.13.0" }, - { name = "fastapi", extras = ["standard"], marker = "extra == 'fastapi'", specifier = ">=0.124.4,<0.125.0" }, + { name = "fastapi", extras = ["standard"], marker = "extra == 'fastapi'", specifier = ">=0.124.4" }, { name = "inflection", specifier = ">=0.5.1" }, { name = "itsdangerous", marker = "extra == 'fastapi'", specifier = ">=2.2.0" }, { name = "jinja2", marker = "extra == 'inertia'", specifier = ">=3.1" }, @@ -481,6 +481,7 @@ dev = [ { name = "itsdangerous", specifier = ">=2.2.0" }, { name = "langchain", specifier = ">=1.0.0" }, { name = "langchain-core", specifier = ">=1.0.0" }, + { name = "pyright", specifier = ">=1.1.411" }, { name = "pytest", specifier = ">=9.0.3" }, { name = "pytest-asyncio", specifier = ">=1.3.0" }, { name = "pytest-cov", specifier = ">=6.0.0" }, From 8a56a6209af5de0cd835cb8b253fd536d96fdf1d Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 23 Jul 2026 14:06:14 -0700 Subject: [PATCH 7/7] feat: ai package fix --- .vscode/settings.json | 3 + example/agents/.gitignore | 1 + example/agents/app/agents/chat.py | 19 +- example/agents/app/agents/graph_agent.py | 9 + .../app/providers/langchain_provider.py | 44 +++ example/agents/app/tools/job_search_tool.py | 10 +- example/agents/bootstrap/application.py | 2 + example/agents/pyproject.toml | 1 + example/agents/routes/api.py | 5 + example/agents/tests/features/job_search.json | 6 + .../tests/features/test_chat_controller.py | 14 +- .../tests/units/agents/record_stream.json | 32 +-- .../units/agents/test_langchain_agent.py | 26 ++ .../tests/units/agents/test_router_agent.py | 12 +- example/agents/tinker.py | 62 +++++ example/agents/uv.lock | 44 +++ fastapi_startkit/pyproject.toml | 1 + .../src/fastapi_startkit/ai/__init__.py | 6 +- .../src/fastapi_startkit/ai/agent.py | 218 +-------------- .../src/fastapi_startkit/ai/ai.py | 22 +- .../src/fastapi_startkit/ai/fakes.py | 53 ---- .../src/fastapi_startkit/ai/graph.py | 113 ++++++++ .../src/fastapi_startkit/ai/response.py | 34 +-- .../src/fastapi_startkit/ai/runner.py | 200 ++++++++++++-- .../src/fastapi_startkit/ai/testing.py | 260 +++++++++--------- .../src/fastapi_startkit/ai/tinker.py | 0 .../src/fastapi_startkit/ai/types.py | 5 + fastapi_startkit/tests/ai/test_agent.py | 77 ++++-- fastapi_startkit/tests/ai/test_agent_fake.py | 143 ++++------ .../tests/ai/test_agent_record_fluent.py | 25 +- .../tests/ai/test_agent_schema.py | 10 +- fastapi_startkit/tests/ai/test_graph_agent.py | 77 ++++++ fastapi_startkit/tests/ai/test_judge_agent.py | 5 +- .../tests/ai/test_structured_output.py | 15 +- .../relationships/test_sqlite_polymorphic.py | 4 +- fastapi_startkit/uv.lock | 2 + 36 files changed, 912 insertions(+), 648 deletions(-) create mode 100644 .vscode/settings.json create mode 100644 example/agents/app/agents/graph_agent.py create mode 100644 example/agents/app/providers/langchain_provider.py create mode 100644 example/agents/tests/features/job_search.json create mode 100644 example/agents/tests/units/agents/test_langchain_agent.py create mode 100644 example/agents/tinker.py delete mode 100644 fastapi_startkit/src/fastapi_startkit/ai/fakes.py create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/graph.py create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/tinker.py create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/types.py create mode 100644 fastapi_startkit/tests/ai/test_graph_agent.py diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..83ddff40 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.languageServer": "Default" +} \ No newline at end of file diff --git a/example/agents/.gitignore b/example/agents/.gitignore index 014a22ef..dc2d3589 100644 --- a/example/agents/.gitignore +++ b/example/agents/.gitignore @@ -10,3 +10,4 @@ node_modules .ruff_cache .pytest_cache .ai +.vite diff --git a/example/agents/app/agents/chat.py b/example/agents/app/agents/chat.py index 6d134452..4e02a9e2 100644 --- a/example/agents/app/agents/chat.py +++ b/example/agents/app/agents/chat.py @@ -1,12 +1,27 @@ from typing import Callable -from fastapi_startkit.ai import Agent, Middleware +from fastapi_startkit.ai import Agent, GraphAgent, Middleware from app.middleware.agent_logger import AgentLogger from app.tools.job_search_tool import job_search_tool -class RouterAgent(Agent): +class RouterAgent(GraphAgent): + def graph(self): + graph = StateGraph(MessagesState) + + # Add nodes + graph.add_node("llm_call", llm_call) + graph.add_node("tool_node", tool_node) + + # Add edges to connect nodes + graph.add_edge(START, "llm_call") + graph.add_conditional_edges("llm_call", should_continue, ["tool_node", END]) + graph.add_edge("tool_node", "llm_call") + + # Compile the agent + agent = graph.compile() + def middleware(self) -> list[Middleware]: return [AgentLogger()] diff --git a/example/agents/app/agents/graph_agent.py b/example/agents/app/agents/graph_agent.py new file mode 100644 index 00000000..c5d4dd03 --- /dev/null +++ b/example/agents/app/agents/graph_agent.py @@ -0,0 +1,9 @@ +from fastapi_startkit.ai import GraphAgent, GraphRunner +from langgraph.graph import StateGraph + + +class SalesAgent(GraphAgent): + def checkpointer(self): + pass + def graph(self, runner: GraphRunner) -> StateGraph: + return StateGraph() diff --git a/example/agents/app/providers/langchain_provider.py b/example/agents/app/providers/langchain_provider.py new file mode 100644 index 00000000..67fcaeb3 --- /dev/null +++ b/example/agents/app/providers/langchain_provider.py @@ -0,0 +1,44 @@ +from fastapi_startkit.support import Provider +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from psycopg.rows import dict_row +from psycopg_pool import AsyncConnectionPool + + +class LazyCheckpointer: + """An awaitable, lazily-initialised async Postgres checkpointer. + + ``Provider.boot()`` runs synchronously while the application is being + constructed — before any serving event loop exists. Async Postgres + connections are bound to the event loop that opens them, so the pool + cannot be opened at boot time. Instead we build the pool closed and open + it (and run ``setup()``) on first ``await`` inside the request loop, + caching the ready saver for subsequent calls. + + Usage: ``checkpointer = await app().make("checkpointer")`` + """ + + def __init__(self, uri: str): + self._pool = AsyncConnectionPool( + conninfo=uri, + open=False, + kwargs={"autocommit": True, "row_factory": dict_row}, + ) + self._saver: AsyncPostgresSaver | None = None + + async def resolve(self) -> AsyncPostgresSaver: + if self._saver is None: + await self._pool.open() + saver = AsyncPostgresSaver(self._pool) + await saver.setup() + self._saver = saver + return self._saver + + def __await__(self): + return self.resolve().__await__() + + +class LangChainProvider(Provider): + DB_URI = "postgresql://postgres:postgres@localhost:5432/agents?sslmode=disable" + + def boot(self): + self.app.bind("checkpointer", LazyCheckpointer(self.DB_URI)) diff --git a/example/agents/app/tools/job_search_tool.py b/example/agents/app/tools/job_search_tool.py index 4d76fd83..63f96a53 100644 --- a/example/agents/app/tools/job_search_tool.py +++ b/example/agents/app/tools/job_search_tool.py @@ -9,7 +9,7 @@ ] -@tool +@tool(description="Use this tools if user wants to search for jobs") def job_search_tool(query: str) -> list: """Searches for jobs based on the given query. Supports wildcards (* and ?) in each term.""" import fnmatch @@ -17,9 +17,7 @@ def job_search_tool(query: str) -> list: patterns = [f"*{term}*" for term in query.lower().split()] return [ - job for job in jobs - if any( - fnmatch.fnmatch(" ".join(str(v) for v in job.values()).lower(), pattern) - for pattern in patterns - ) + job + for job in jobs + if any(fnmatch.fnmatch(" ".join(str(v) for v in job.values()).lower(), pattern) for pattern in patterns) ] diff --git a/example/agents/bootstrap/application.py b/example/agents/bootstrap/application.py index f198c0a0..47011cfe 100644 --- a/example/agents/bootstrap/application.py +++ b/example/agents/bootstrap/application.py @@ -11,11 +11,13 @@ from config.logging import LoggingConfig from config.vite import ViteConfig from app.providers.fastapi_provider import FastapiProvider +from app.providers.langchain_provider import LangChainProvider app: Application = Application( base_path=Path(__file__).resolve().parent.parent, providers=[ AISkillProvider, + LangChainProvider, (LogProvider,LoggingConfig), (FastapiProvider, FastAPIConfig), AIProvider, diff --git a/example/agents/pyproject.toml b/example/agents/pyproject.toml index aa074982..4b1d78d7 100644 --- a/example/agents/pyproject.toml +++ b/example/agents/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "fastapi-startkit[ai,database,fastapi,sqlite]==0.44.0", "langchain>=1.3.10", "langchain-google-genai>=4.2.5", + "langgraph-checkpoint-postgres>=3.1.0", ] [tool.uv] diff --git a/example/agents/routes/api.py b/example/agents/routes/api.py index e1804813..d78f5073 100644 --- a/example/agents/routes/api.py +++ b/example/agents/routes/api.py @@ -1,6 +1,7 @@ from fastapi import APIRouter, Request from fastapi.responses import StreamingResponse from fastapi_startkit.inertia import Inertia +from langchain.agents import create_agent from app.agents.chat import RouterAgent from app.requests.chat import ChatRequest @@ -20,6 +21,10 @@ async def index(request: Request): @api.post("/chat") async def chat(request: ChatRequest): + from fastapi_startkit.application import app + + agent = create_agent(checkpointer=await app().make("checkpointer")) + response = await RouterAgent().prompt(request.message) return {"content": response.content} diff --git a/example/agents/tests/features/job_search.json b/example/agents/tests/features/job_search.json new file mode 100644 index 00000000..d102026f --- /dev/null +++ b/example/agents/tests/features/job_search.json @@ -0,0 +1,6 @@ +{ + "4fd116f6d1844cc5e3bfb521bf361bf66ac050cf5e97bc506d66cc8fbf7ade43": { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "tool_calls": [] + } +} \ No newline at end of file diff --git a/example/agents/tests/features/test_chat_controller.py b/example/agents/tests/features/test_chat_controller.py index 41eb02f7..bbd47eee 100644 --- a/example/agents/tests/features/test_chat_controller.py +++ b/example/agents/tests/features/test_chat_controller.py @@ -1,3 +1,5 @@ +from dumpdie import dd + from app.agents.chat import RouterAgent from tests.test_case import TestCase @@ -28,13 +30,13 @@ async def test_it_responds_with_stream(self): response.assert_ok() response.assert_stream_contains('Hello there, This is stream chat, Hope you are doing well.') - @RouterAgent.record("record_no_stream.json") + @RouterAgent.log("record_no_stream.json") async def test_it_records_without_stream(self): response = await self.post("/chat", json={"message": "Hi, I am Alex, This is unittest, Please respond by calling my name."}) response.assert_ok() response.assert_contents("Alex") - @RouterAgent.record("record_stream.json") + @RouterAgent.log("record_stream.json") async def test_chat_responds_for_other_greetings(self): response = await self.post("/chat/stream", json={ "message": "Hi, I am Bedram, This is unittest, Please respond by calling my name." @@ -42,3 +44,11 @@ async def test_chat_responds_for_other_greetings(self): response.assert_ok() response.assert_stream_contains("Bedram") + + @RouterAgent.log("job_search.json") + async def test_user_can_perform_the_job_search(self): + response = await self.post("/chat", json={ + "message": "suggest me python developer jobs" + }) + + response.assert_ok() diff --git a/example/agents/tests/units/agents/record_stream.json b/example/agents/tests/units/agents/record_stream.json index 82afdab8..756fda70 100644 --- a/example/agents/tests/units/agents/record_stream.json +++ b/example/agents/tests/units/agents/record_stream.json @@ -1,34 +1,10 @@ { "034751ef1322a12f3406dad43313f9cdb31f4ea85d9452c22bf7529ad8e92614": { - "content": "Hello! How can I help you today?", + "content": "Hi there! How can I help you today?", "tool_calls": [] }, - "358018dc096ce95b78b7561f562ba049848d03fb8b41953ed75c22c9ad20e228": { - "content": "", - "tool_calls": [ - { - "args": { - "query": "python developer" - }, - "id": "call_1", - "name": "job_search_tool" - } - ] - }, - "b6579ba866634a66229ddbed0463c8c990eaaaf349385fbc67c2c2c1886085cd": { - "content": "", - "tool_calls": [ - { - "args": { - "query": "python developer" - }, - "id": "call_1", - "name": "job_search_tool" - } - ] - }, - "judge:d33d8a7074bb7d980bf548694cf3b50ba853fa2f01a8df921bd28d7373ef5580": { - "passed": true, - "reasoning": "greets the user" + "aaa92579b146c573996a42ba725a88d59fa731fa932e869862db333d4bc70d02": { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "tool_calls": [] } } \ No newline at end of file diff --git a/example/agents/tests/units/agents/test_langchain_agent.py b/example/agents/tests/units/agents/test_langchain_agent.py new file mode 100644 index 00000000..5e79ce6e --- /dev/null +++ b/example/agents/tests/units/agents/test_langchain_agent.py @@ -0,0 +1,26 @@ +from dumpdie import dd +from langchain.agents import create_agent +from langchain_core.language_models import GenericFakeChatModel + +from tests.test_case import TestCase + + +class Agent: + async def prompt(self, input): + from fastapi_startkit.application import app + + agent = create_agent( + model=GenericFakeChatModel(messages=iter(["hello"])), + system_prompt="You are a helpful assistant", + checkpointer=await app().make("checkpointer"), + ) + + return await agent.ainvoke( + {"messages": [{"role": "user", "content": "hello"}]}, {"configurable": {"thread_id": "1"}} + ) + + +class TestLangchainAgent(TestCase): + async def test_the_router_agent(self): + response = await Agent().prompt("hello") + dd(response) diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py index 17a71d18..a5b8f7e1 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -1,3 +1,4 @@ +from fastapi_startkit.ai.tinker import ToolCall from langchain_core.messages import AIMessage, HumanMessage from app.agents.chat import RouterAgent @@ -10,15 +11,14 @@ async def test_the_router_agent(self): await agent.prompt("hello") agent.assert_text_response() agent.assert_tool_not_called(["job_search_tool"]) - 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) + def assert_tool_calls(tool: ToolCall): + return tool.name == "job_search_tool" + await agent.prompt("suggest python developer jobs") - agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") + agent.assert_tool_called("job_search_tool", assert_tool_calls) async def test_the_router_with_initial_messages(self): with RouterAgent.record( diff --git a/example/agents/tinker.py b/example/agents/tinker.py new file mode 100644 index 00000000..d2fa09f4 --- /dev/null +++ b/example/agents/tinker.py @@ -0,0 +1,62 @@ +import asyncio +import operator +from collections.abc import Callable +from typing import Annotated, TypedDict + +from dumpdie import dump +from langchain.agents import create_agent +from langchain_core.messages import AnyMessage +from langgraph.checkpoint.memory import InMemorySaver + +from app.tools.job_search_tool import job_search_tool +from bootstrap.application import app # NOQA + + +class ChatState(TypedDict): + messages: Annotated[list[AnyMessage], operator.add] + llm_calls: int + + +agent = create_agent( + model="google_genai:gemini-3.1-flash-lite", + checkpointer=InMemorySaver(), + tools=[job_search_tool], +) + + +async def prompt(message: str): + config = {"configurable": {"thread_id": "1"}} + return await agent.ainvoke(input={"messages": [{"role": "user", "content": message}]}, config=config) + + +class Agent: + def __init__(self, prompt_handler=Callable): + self.prompt_handler = prompt_handler + + async def prompt(self, message): + pass + + async def ainvoke(self): + await prompt(message="suggest me frontend developer jobs") + + +async def main(): + response = await prompt(message="Hello, world!") + dump(response["messages"]) + response = await prompt(message="suggest me frontend developer jobs") + dump(response["messages"]) + + +asyncio.run(main()) + + +# def test_it_can_prompt(): +# with Agent(prompt_handler=prompt) as agent: +# agent.prompt("hi") # hit the real end point for the first time, records the responses +# agent.assert_prompted("hi") +# agent.assert_prompt_judged(model="", expectation="") +# +# agent.prompt("suggest me python developer jobs") # hit for the first time and second records will be recorded +# agent.assert_tool_called( +# lambda tool: tool.name == "job_search_tool" and tool.args == {"query": "python developer jobs"} +# ) diff --git a/example/agents/uv.lock b/example/agents/uv.lock index 1d3ab1b2..302e78f4 100644 --- a/example/agents/uv.lock +++ b/example/agents/uv.lock @@ -433,6 +433,7 @@ dependencies = [ ai = [ { name = "langchain" }, { name = "langchain-core" }, + { name = "langgraph" }, ] database = [ { name = "faker" }, @@ -462,6 +463,7 @@ requires-dist = [ { name = "jinja2", marker = "extra == 'vite'", specifier = ">=3.1" }, { name = "langchain", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "langchain-core", marker = "extra == 'ai'", specifier = ">=1.0.0" }, + { name = "langgraph", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "markupsafe", marker = "extra == 'inertia'", specifier = ">=2.0" }, { name = "pendulum", specifier = ">=3.1.0,<4.0.0" }, { name = "pydantic", specifier = ">=2.12.5" }, @@ -499,6 +501,7 @@ dependencies = [ { name = "fastapi-startkit", extra = ["ai", "database", "fastapi", "sqlite"] }, { name = "langchain" }, { name = "langchain-google-genai" }, + { name = "langgraph-checkpoint-postgres" }, ] [package.dev-dependencies] @@ -515,6 +518,7 @@ requires-dist = [ { name = "fastapi-startkit", extras = ["ai", "database", "fastapi", "sqlite"], editable = "../../fastapi_startkit" }, { name = "langchain", specifier = ">=1.3.10" }, { name = "langchain-google-genai", specifier = ">=4.2.5" }, + { name = "langgraph-checkpoint-postgres", specifier = ">=3.1.0" }, ] [package.metadata.requires-dev] @@ -918,6 +922,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] +[[package]] +name = "langgraph-checkpoint-postgres" +version = "3.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "langgraph-checkpoint" }, + { name = "orjson" }, + { name = "psycopg" }, + { name = "psycopg-pool" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/51/5a2dc42e8b5d5942b933b5b7237eae5a4dbc92508a04727c263dd383ad8a/langgraph_checkpoint_postgres-3.1.0.tar.gz", hash = "sha256:02bff4ab63d9dae8eab3a9640fce1d479da8965c9fba7b0dc04cb1f7c56f0a55", size = 148473, upload-time = "2026-05-12T03:40:10.599Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/cd/eff9b82bc3b5f62d481b437099f44f3ef7b1d907f166fb4ee25e8f84a1e7/langgraph_checkpoint_postgres-3.1.0-py3-none-any.whl", hash = "sha256:814cce2ef35d792bf07b090a95eed004f1acac0724fe6605536b13f6d1e7032c", size = 48988, upload-time = "2026-05-12T03:40:08.925Z" }, +] + [[package]] name = "langgraph-prebuilt" version = "1.1.0" @@ -1205,6 +1224,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "psycopg" +version = "3.3.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/2f/cb91e5502ec9de1de6f1b76cfbf69531932725361168bb06963620c77e2e/psycopg-3.3.4.tar.gz", hash = "sha256:e21207764952cff81b6b8bdacad9a3939f2793367fdac2987b3aac36a651b5bc", size = 165799, upload-time = "2026-05-01T23:31:55.179Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/e0/7b3dee031daae7743609ce3c746565d4a3ed7c2c186479eb48e34e838c64/psycopg-3.3.4-py3-none-any.whl", hash = "sha256:b6bbc25ccf05c8fad3b061d9db2ef0909a555171b84b07f29458a447253d679a", size = 213001, upload-time = "2026-05-01T23:20:50.816Z" }, +] + +[[package]] +name = "psycopg-pool" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/82/7a23d26039827ecd4ebe93905651029ddd307c5182ad59296dfb6f67b528/psycopg_pool-3.3.1.tar.gz", hash = "sha256:b10b10b7a175d5cc1592147dc5b7eec8a9e0834eb3ed2c4a92c858e2f51eb63c", size = 31661, upload-time = "2026-05-01T23:31:59.809Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/ed/89c2c620af0e1660354cd8aabf9f5b21f911597ce22acb37c805d6c86bc8/psycopg_pool-3.3.1-py3-none-any.whl", hash = "sha256:2af5b432941c4c9ad5c87b3fa410aec910ec8f7c122855897983a06c45f2e4b5", size = 40023, upload-time = "2026-05-01T23:31:53.136Z" }, +] + [[package]] name = "pyasn1" version = "0.6.3" diff --git a/fastapi_startkit/pyproject.toml b/fastapi_startkit/pyproject.toml index a7717034..0be5d27c 100644 --- a/fastapi_startkit/pyproject.toml +++ b/fastapi_startkit/pyproject.toml @@ -86,6 +86,7 @@ inertia = [ ai = [ "langchain>=1.0.0", "langchain-core>=1.0.0", + "langgraph>=1.0.0", ] [dependency-groups] diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index db2f5dad..6096c894 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -6,7 +6,7 @@ from .config.ai import AIConfig from .decorators import max_steps, max_tokens, model, provider, timeout, top_p from .document import Document -from .fakes import fake_chat_model +from .graph import GraphAgent, GraphRunner, GraphState from .image import Image, ImageResponse from .image_factory import ImageFactory from .ai import Ai @@ -33,8 +33,10 @@ "AudioFactory", "Document", "ElevenLabsConfig", - "fake_chat_model", "GoogleConfig", + "GraphAgent", + "GraphRunner", + "GraphState", "Image", "ImageFactory", "ImageResponse", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index f13be05a..5875325e 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -1,14 +1,16 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Optional, Type +from typing import TYPE_CHECKING, AsyncIterator, Optional, Type from .document import Document from .response import AgentResponse +from .runner import BaseRunner, Runner if TYPE_CHECKING: from langchain_core.tools import BaseTool from .testing import AgentFake, AgentRecordFake + from .types import Middleware class Agent: @@ -19,9 +21,6 @@ class Agent: timeout: float = 30.0 top_p: float = 1.0 - def __init__(self): - self._call_log: list[dict] = [] - def messages(self) -> list[dict]: return [] @@ -34,7 +33,7 @@ def schema(self) -> Optional[Type]: def tools(self) -> list[BaseTool]: return [] - def middleware(self) -> list[Callable]: + def middleware(self) -> list[Middleware]: return [] def provider_options(self) -> dict: @@ -48,18 +47,9 @@ async def prompt( attachments: list[Document] | None = None, provider_options: dict | None = None, ) -> AgentResponse: - stand_in = self._faked() - if stand_in is not None: - response = await stand_in.prompt(message, attachments=attachments) - self._log_call("prompt", message) - return self._apply_schema(response) - - messages = self._build_messages(message, attachments) - chat_model = self._build_model(model, provider_options) - - response = await self._run_pipeline(chat_model, messages) - self._log_call("prompt", message) - return self._apply_schema(response) + return await self.runner().run( + message, model=model, attachments=attachments, provider_options=provider_options + ) async def stream( self, @@ -68,17 +58,12 @@ async def stream( model: str | None = None, provider_options: dict | None = None, ) -> AsyncIterator[str]: - self._log_call("stream", message) - - swapped = self._faked() - if swapped is not None: - async for chunk in swapped.stream(message): - yield chunk - return - - async for chunk in self._stream(message, model=model, provider_options=provider_options): + async for chunk in self.runner().stream(message, model=model, provider_options=provider_options): yield chunk + def runner(self) -> BaseRunner: + return Runner(self) + @classmethod def fake(cls, responses: list) -> "AgentFake": from .testing import AgentFake @@ -91,184 +76,3 @@ def record(cls, cassette: str | None = None, messages: list | None = None) -> "A return AgentRecordFake(cls(), cassette, messages) - @classmethod - def _binding(cls) -> Any: - from fastapi_startkit.application import app - - container = app() - return container.make(cls.__name__) if container.has(cls.__name__) else None - - @classmethod - def make(cls) -> "Agent": - binding = cls._binding() - return binding if binding is not None else cls() - - def _faked(self) -> Any: - binding = type(self)._binding() - return binding if binding is not self else None - - def assert_prompted(self, times: int | None = None) -> None: - calls = [c for c in self._call_log if c["method"] in ("prompt", "stream")] - if times is not None: - assert len(calls) == times, f"Expected {times} prompt call(s), got {len(calls)}" - else: - assert len(calls) > 0, "Expected at least one prompt() or stream() call, but none were made" - - def assert_not_prompted(self) -> None: - self.assert_prompted(times=0) - - def reset(self) -> "Agent": - self._call_log.clear() - return self - - def _log_call(self, method: str, message: str) -> None: - self._call_log.append({"method": method, "message": message}) - - async def _run_pipeline(self, chat_model: Any, messages: list) -> AgentResponse: - from .pipeline import Response, build_pipeline # noqa: PLC0415 - from .runner import Runner # noqa: PLC0415 - - chain = list(self.middleware()) - if not chain: - return await self._invoke(chat_model, messages) - - def core(model: Any) -> Response: - async def _run(): - result = await Runner(self, model).run(messages) - yield result - - return Response(_run) - - pipeline = build_pipeline(chain, core) - raw = await pipeline(chat_model) - return self._to_agent_response(raw) - - async def _apply_middleware( - self, - chat_model: Any, - final: Callable[[Any], Any], - ) -> AgentResponse: - chain = list(self.middleware()) - - def build(mw_list: list, fn: Callable) -> Callable: - if not mw_list: - return fn - head, *tail = mw_list - next_fn = build(tail, fn) - mw = head() if isinstance(head, type) else head - return lambda model: mw(model, next_fn) - - return await build(chain, final)(chat_model) - - def _build_instruction(self) -> str | None: - return self.instructions() - - def _build_messages( - self, - message: str, - attachments: list[Document] | None = None, - ) -> list[dict]: - messages: list[dict] = [] - - instruction = self.instructions() - if instruction: - messages.append({"role": "system", "content": instruction}) - - messages.extend(self.messages() or []) - - if message: - messages.append({"role": "user", "content": message}) - - if attachments: - content: Any = [{"type": "text", "text": message}] - for doc in attachments: - content.append(doc.to_langchain_block()) - messages.append({"role": "user", "content": content}) - - return messages - - 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, 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 = [] 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, parsed=parsed) - - def _apply_schema(self, response: AgentResponse) -> AgentResponse: - schema = self.schema() - if schema is not None and response.parsed is None and response.content: - response.parsed = self._build_schema(schema, response.content) - return response - - @staticmethod - def _build_schema(schema: Any, content: str) -> Any: - import json # noqa: PLC0415 - - if hasattr(schema, "model_validate_json"): - return schema.model_validate_json(content) - if hasattr(schema, "model_validate"): - return schema.model_validate(json.loads(content)) - return schema(**json.loads(content)) - - async def _invoke(self, chat_model: Any, messages: list[dict]) -> AgentResponse: - from .runner import Runner # noqa: PLC0415 - - result = await Runner(self, chat_model).run(messages) - return self._to_agent_response(result) - - async def _run( - self, - message: str, - model: str | None = None, - attachments: list[Document] | None = None, - provider_options: dict | None = None, - ) -> AgentResponse: - messages = self._build_messages(message, attachments) - chat_model = self._build_model(model, provider_options) - return await self._invoke(chat_model, messages) - - async def _stream( - self, - message: str, - model: str | None = None, - provider_options: dict | None = None, - ) -> AsyncIterator[str]: - from .pipeline import Response, build_pipeline # noqa: PLC0415 - from .runner import StreamRunner # noqa: PLC0415 - - messages = self._build_messages(message) - chat_model = self._build_model(model, provider_options, structured=False) - chain = list(self.middleware()) - - def core(m: Any) -> Response: - return Response(lambda: StreamRunner(self, m).run(messages)) - - pipeline = build_pipeline(chain, core) if chain else core - - async for chunk in pipeline(chat_model): - yield chunk diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index c7433e0a..28b76345 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -5,12 +5,13 @@ from .lab import Lab if TYPE_CHECKING: + from langchain_core.language_models.fake_chat_models import GenericFakeChatModel + from .agent import Agent class Ai: - fake_agent_models: dict[str, Any] = {} - fake_agent_responses: dict[str, Any] = {} + _fakes: dict[str, GenericFakeChatModel] = {} def __init__(self) -> None: pass @@ -26,43 +27,40 @@ def fake(cls, agent: "Agent | str", messages: list) -> Any: turns = [message if hasattr(message, "content") else AIMessage(content=str(message)) for message in messages] model = GenericFakeChatModel(messages=iter(turns)) - cls.fake_agent_models[cls._key(agent)] = model + cls._fakes[cls._key(agent)] = model return model @classmethod def has_fake_model_for(cls, agent: "Agent | str") -> bool: - return cls._key(agent) in cls.fake_agent_models + return cls._key(agent) in cls._fakes @classmethod def get_fake_model_for(cls, agent: "Agent | str") -> Any: - return cls.fake_agent_models[cls._key(agent)] + return cls._fakes[cls._key(agent)] @classmethod def forget(cls, agent: "Agent | str") -> None: - cls.fake_agent_models.pop(cls._key(agent), None) + cls._fakes.pop(cls._key(agent), None) @classmethod def reset_fakes(cls) -> None: - cls.fake_agent_models.clear() - cls.fake_agent_responses.clear() + cls._fakes.clear() 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, structured) + return self.build(agent, model, provider_options) 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 @@ -87,7 +85,7 @@ def build( chat_model = chat_model.bind_tools(tools) if tools else chat_model schema = agent.schema() - if structured and schema is not None: + if schema is not None: chat_model = chat_model.with_structured_output(schema, include_raw=True) return chat_model diff --git a/fastapi_startkit/src/fastapi_startkit/ai/fakes.py b/fastapi_startkit/src/fastapi_startkit/ai/fakes.py deleted file mode 100644 index 870e624a..00000000 --- a/fastapi_startkit/src/fastapi_startkit/ai/fakes.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -import json -import re -from typing import Any, Iterable - - -def _require_langchain(): - try: - from langchain_core.language_models.fake_chat_models import GenericFakeChatModel - from langchain_core.messages import AIMessage - except ImportError as exc: - raise ImportError( - "The agent test harness requires the 'ai' extra. Install it with: pip install \"fastapi-startkit[ai]\"" - ) from exc - return GenericFakeChatModel, AIMessage - - -def fake_chat_model(turns: Iterable[Any]): - generic_model, ai_message = _require_langchain() - - from langchain_core.messages import AIMessageChunk - from langchain_core.outputs import ChatGenerationChunk - - class _FakeChatModel(generic_model): - def bind_tools(self, tools, **kwargs): - return self - - def _stream(self, messages, stop=None, run_manager=None, **kwargs): - message = next(self.messages) - if not isinstance(message, ai_message): - message = ai_message(content=str(message)) - - content = message.content if isinstance(message.content, str) else str(message.content) - for token in re.split(r"(\s)", content): - if token: - yield ChatGenerationChunk(message=AIMessageChunk(content=token, id=message.id)) - - tool_calls = list(message.tool_calls or []) - if tool_calls: - chunks = [ - { - "name": call["name"], - "args": json.dumps(call.get("args", {})), - "id": call.get("id"), - "index": index, - } - for index, call in enumerate(tool_calls) - ] - yield ChatGenerationChunk(message=AIMessageChunk(content="", tool_call_chunks=chunks, id=message.id)) - - normalized = [t if isinstance(t, ai_message) else ai_message(content=str(t)) for t in turns] - return _FakeChatModel(messages=iter(normalized)) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/graph.py b/fastapi_startkit/src/fastapi_startkit/ai/graph.py new file mode 100644 index 00000000..20ab656f --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/graph.py @@ -0,0 +1,113 @@ + +from __future__ import annotations + +import time +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Annotated, Any, TypedDict + +from langgraph.graph import END, StateGraph +from langgraph.graph.message import add_messages + +from .agent import Agent +from .response import AgentResponse +from .runner import BaseRunner + +if TYPE_CHECKING: + from .document import Document + + +class GraphState(TypedDict): + + messages: Annotated[list, add_messages] + llm_calls: int + + +class GraphAgent(Agent, ABC): + + def runner(self) -> GraphRunner: + return GraphRunner(self) + + @abstractmethod + def graph(self, runner: GraphRunner) -> StateGraph: + ... + + +class GraphRunner(BaseRunner): + + agent: GraphAgent + + def __init__(self, agent: GraphAgent) -> None: + super().__init__(agent) + self._chat_model: Any = None + + + async def llm(self, state: GraphState) -> dict: + reply = await self._chat_model.ainvoke(state["messages"]) + return {"messages": [reply], "llm_calls": state.get("llm_calls", 0) + 1} + + async def call_tools(self, state: GraphState) -> dict: + tools_by_name = {tool.name: tool for tool in self.agent.tools()} + results = [] + for call in getattr(state["messages"][-1], "tool_calls", None) or []: + results.append(await tools_by_name[call["name"]].ainvoke(call)) + return {"messages": results} + + def route(self, state: GraphState) -> str: + return "tools" if getattr(state["messages"][-1], "tool_calls", None) else END + + + def _compile(self, model: str | None, provider_options: dict | None) -> Any: + self._chat_model = self._build_model(model, provider_options) + return self.agent.graph(self).compile() + + @property + def _config(self) -> dict: + return {"recursion_limit": self.agent.max_steps * 2 + 1} + + async def run( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: + started = time.perf_counter() + compiled = self._compile(model, provider_options) + state = {"messages": self._build_messages(message, attachments), "llm_calls": 0} + result = await compiled.ainvoke(state, config=self._config) + response = self._apply_schema(self._to_response(result)) + response.runtime = time.perf_counter() - started + return response + + async def stream( + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: + compiled = self._compile(model, provider_options) + state = {"messages": self._build_messages(message), "llm_calls": 0} + async for chunk, _meta in compiled.astream(state, stream_mode="messages", config=self._config): + text = chunk.content if isinstance(chunk.content, str) else str(chunk.content) + if text: + yield text + + @staticmethod + def _to_response(result: dict) -> AgentResponse: + messages = result.get("messages", []) + final = messages[-1] if messages else None + content = getattr(final, "content", "") or "" + if not isinstance(content, str): + content = str(content) + + tool_calls = [call for m in messages for call in (getattr(m, "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) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/response.py b/fastapi_startkit/src/fastapi_startkit/ai/response.py index 46d02ebd..57e2adfa 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/response.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/response.py @@ -1,5 +1,3 @@ -"""AgentResponse and AgentSnapshot — response containers for AI agents.""" - from __future__ import annotations import json @@ -13,20 +11,21 @@ @dataclass class AgentResponse: - """Returned by Agent.prompt(). Wraps the LLM response.""" - content: str = "" tool_calls: list[dict] = field(default_factory=list) usage: dict = field(default_factory=dict) raw: Any = None parsed: Any = None + runtime: float = 0.0 + + @property + def runtime_ms(self) -> float: + return self.runtime * 1000 def text(self) -> str: - """Return the text content.""" return self.content def json(self) -> Any: - """Parse the content as JSON.""" return json.loads(self.content) def __str__(self) -> str: @@ -38,37 +37,23 @@ def __bool__(self) -> bool: @dataclass class AgentSnapshot: - """ - Record-and-replay snapshot for testing. - - - If the file at ``path`` **does not exist**: the agent calls the real API, - saves the response as JSON, then returns it. - - If the file **exists**: the saved response is loaded and returned without - hitting the API. - - Example:: - - agent.fake({"*analyze*": AgentSnapshot(path="tests/fixtures/analysis.json")}) - """ path: str def exists(self) -> bool: - """Return True if the snapshot file is already recorded.""" return os.path.exists(self.path) def load(self) -> AgentResponse: - """Load the recorded response from disk.""" with open(self.path) as f: data = json.load(f) return AgentResponse( content=data.get("content", ""), tool_calls=data.get("tool_calls", []), usage=data.get("usage", {}), + runtime=data.get("runtime", 0.0), ) def save(self, response: AgentResponse) -> None: - """Persist a real API response to disk for future replays.""" os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True) with open(self.path, "w") as f: json.dump( @@ -76,18 +61,15 @@ def save(self, response: AgentResponse) -> None: "content": response.content, "tool_calls": response.tool_calls, "usage": response.usage, + "runtime": response.runtime, }, f, indent=2, ) async def resolve(self, agent: "Agent", message: str, **run_kwargs: Any) -> AgentResponse: - """ - Return the response — from disk if recorded, or from the real API - (which is then saved for future runs). - """ if self.exists(): return self.load() - response = await agent._run(message, **run_kwargs) + response = await agent.prompt(message, **run_kwargs) self.save(response) return response diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index ba6ded73..ae35bfea 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -1,14 +1,21 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Sequence +import time +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Any +from langchain.agents import create_agent from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, ToolCall from langchain_core.runnables import Runnable from langchain_core.tools import BaseTool +from .pipeline import Response, build_pipeline +from .response import AgentResponse + if TYPE_CHECKING: from .agent import Agent + from .document import Document Message = BaseMessage | dict[str, Any] @@ -17,46 +24,181 @@ def _as_text(content: Any) -> str: return content if isinstance(content, str) else str(content) -class Runner: - 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 - - async def run(self, messages: Sequence[Message]) -> BaseMessage: - history: list[Message] = list(messages) - response: AIMessage = await self.model.ainvoke(history) # type: ignore[assignment] - +class BaseRunner(ABC): + def __init__(self, agent: Agent) -> None: + self.agent = agent + + def _build_messages(self, message: str, attachments: list[Document] | None = None) -> list[dict]: + agent = self.agent + messages: list[dict] = [] + + instruction = agent.instructions() + if instruction: + messages.append({"role": "system", "content": instruction}) + + messages.extend(agent.messages() or []) + + if message: + messages.append({"role": "user", "content": message}) + + if attachments: + content: Any = [{"type": "text", "text": message}] + for doc in attachments: + content.append(doc.to_langchain_block()) + messages.append({"role": "user", "content": content}) + + return messages + + def _build_model(self, model: str | None = None, provider_options: dict | None = None) -> Any: + from .ai import Ai # noqa: PLC0415 + + return Ai().get_model_for(self.agent, model, provider_options) + + def _apply_schema(self, response: AgentResponse) -> AgentResponse: + schema = self.agent.schema() + if schema is not None and response.parsed is None and response.content: + response.parsed = self._build_schema(schema, response.content) + return response + + @staticmethod + def _build_schema(schema: Any, content: str) -> Any: + import json # noqa: PLC0415 + + if hasattr(schema, "model_validate_json"): + return schema.model_validate_json(content) + if hasattr(schema, "model_validate"): + return schema.model_validate(json.loads(content)) + return schema(**json.loads(content)) + + @abstractmethod + async def run( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: ... + + @abstractmethod + def stream( + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: ... + + +class Runner(BaseRunner): + async def run( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: + started = time.perf_counter() + messages = self._build_messages(message, attachments) + model = self._build_model(model, provider_options) + + create_agent( + model=model, + tools=self.agent.tools, + ) + response = await self._run_pipeline(model, messages) + response = self._apply_schema(response) + response.runtime = time.perf_counter() - started + return response + + async def _run_pipeline(self, chat_model: Any, messages: list) -> AgentResponse: + chain = list(self.agent.middleware()) + if not chain: + raw = await self._invoke(chat_model, messages) + return self._to_agent_response(raw) + + def core(model: Any) -> Response: + async def _run() -> AsyncIterator[Any]: + yield await self._invoke(model, messages) + + return Response(_run) + + pipeline = build_pipeline(chain, core) + raw = await pipeline(chat_model) + return self._to_agent_response(raw) + + 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 = [] 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, parsed=parsed) + + async def stream( + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: + messages = self._build_messages(message) + chat_model = self._build_model(model, provider_options) + chain = list(self.agent.middleware()) + + def core(m: Any) -> Response: + return Response(lambda: self._invoke_stream(m, messages)) + + pipeline = build_pipeline(chain, core) if chain else core + + async for chunk in pipeline(chat_model): + yield chunk + + async def _invoke(self, model: Runnable[Any, BaseMessage], messages: list) -> BaseMessage: + response: AIMessage = await model.ainvoke(list(messages)) # type: ignore[assignment] if isinstance(response, dict) and "parsed" in response: return response # type: ignore[return-value] - if not response.tool_calls: return response - 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] - - def _resolve_tool(self, name: str) -> BaseTool: - try: - return self._tools[name] - except KeyError: - raise ValueError(f"Agent has no tool named {name!r}") from None - - -class StreamRunner(Runner): - async def run(self, messages: Sequence[Message]) -> AsyncIterator[str]: # type: ignore[override] - history: list[Message] = list(messages) - + async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list) -> AsyncIterator[str]: gathered: AIMessageChunk | None = None - async for chunk in self.model.astream(history): + async for chunk in model.astream(list(messages)): if chunk.content: yield _as_text(chunk.content) gathered = chunk if gathered is None else gathered + chunk # type: ignore[operator] if gathered is None or not gathered.tool_calls: return - for message in await self._run_tools(gathered.tool_calls): yield _as_text(message.content) + + async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: + tools: dict[str, BaseTool] = {tool.name: tool for tool in self.agent.tools()} + results: list[BaseMessage] = [] + for call in tool_calls: + try: + selected = tools[call["name"]] + except KeyError: + raise ValueError(f"Agent has no tool named {call['name']!r}") from None + results.append(await selected.ainvoke(call)) + return results diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 83120191..32f00f13 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -1,6 +1,5 @@ from __future__ import annotations -import fnmatch import functools import hashlib import inspect @@ -18,59 +17,33 @@ from .document import Document -def _matches(pattern: str, message: str) -> bool: - pattern, message = pattern.lower(), message.lower() - if any(ch in pattern for ch in "*?["): - return fnmatch.fnmatch(message, pattern) - return pattern in message - - -class _Recorder: - def __init__(self) -> None: - self.calls: list[str] = [] - self.attachments: list[list[Document]] = [] - - def _record_call(self, message: str, attachments: list[Document] | None) -> None: - self.calls.append(message) - self.attachments.append(list(attachments or [])) - - @property - def prompt_count(self) -> int: - return len(self.calls) - - def assert_prompted(self, pattern: str | None = None) -> None: - if pattern is None: - assert self.calls, "Expected the agent to be prompted, but it never was." - return - assert any(_matches(pattern, message) for message in self.calls), ( - f"Expected a prompt matching {pattern!r}, but none did. Got: {self.calls!r}" - ) - - def assert_not_prompted(self) -> None: - assert not self.calls, f"Expected no prompts, but got: {self.calls!r}" - - def _joined(value: Any) -> str: - """A cassette value is a buffered string or a list of stream chunks.""" return "".join(value) if isinstance(value, list) else value class AgentFake: - """Registers a fixed, ordered list of replies as ``agent_cls``'s chat - model for the duration of a ``with`` block (or a decorated function). - - ``prompt()``/``stream()`` still run the real message-building, pipeline, - and tool-execution path; see ``Ai.fake()``. - """ def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._agent_cls = agent_cls - self._responses = responses + self._responses = list(responses) + self._agent = agent_cls() + self._agent.messages = self._history # type: ignore[method-assign] + self._records: list[dict] = [] + self._last_response: AgentResponse | None = None + self.last_elapsed: float | None = None + + def _history(self) -> list: + return self._records + + @property + def _prompts(self) -> list[str]: + return [r["content"] for r in self._records if r.get("role") == "user"] - def __enter__(self) -> None: + def __enter__(self) -> "AgentFake": from .ai import Ai Ai.fake(self._agent_cls.__name__, self._responses) + return self def __exit__(self, *_exc: Any) -> bool: from .ai import Ai @@ -78,6 +51,106 @@ def __exit__(self, *_exc: Any) -> bool: Ai.forget(self._agent_cls.__name__) return False + async def prompt(self, message: str, *, attachments: list[Document] | None = None) -> AgentResponse: + start = time.monotonic() + self._last_response = await self._agent.prompt(message, attachments=attachments) + self.last_elapsed = time.monotonic() - start + self._remember(message, self._last_response) + return self._last_response + + async def stream(self, message: str) -> AsyncIterator[str]: + chunks: list[str] = [] + async for chunk in self._agent.stream(message): + chunks.append(chunk) + yield chunk + self._last_response = AgentResponse(content="".join(chunks)) + self._remember(message, self._last_response) + + def _remember(self, message: str, response: AgentResponse) -> None: + self._records.append({"role": "user", "content": message}) + self._records.append({"role": "assistant", "content": response.content}) + + + def assert_prompt(self, expected: str | Callable[[str], bool]) -> None: + if callable(expected): + assert any(expected(p) for p in self._prompts), ( + f"No recorded prompt satisfied the predicate. Got: {self._prompts!r}" + ) + else: + assert any(expected in p for p in self._prompts), ( + f"Expected a prompt containing {expected!r}, but none did. Got: {self._prompts!r}" + ) + + def assert_response(self, expected: str) -> None: + response = self._require_response() + assert expected in response.content, f"Expected response to contain {expected!r}, got {response.content!r}" + + def assert_tool_call(self, name: str) -> None: + response = self._require_response() + called = [tc.get("name") for tc in response.tool_calls] + assert name in called, f"Expected tool {name!r} to be called, but got: {called}" + + def assert_prompted(self, times: int | None = None) -> None: + if times is not None: + assert len(self._prompts) == times, f"Expected {times} prompt call(s), got {len(self._prompts)}" + else: + assert self._prompts, "Expected at least one prompt() or stream() call, but none were made" + + def assert_not_prompted(self) -> None: + self.assert_prompted(times=0) + + def reset(self) -> "AgentFake": + self._records.clear() + self._last_response = None + return self + + def _require_response(self) -> AgentResponse: + assert self._last_response is not None, "No prompt() call has been made yet." + return self._last_response + + def _tool_call_names(self) -> list[str]: + return [tc.get("name", "") for tc in self._require_response().tool_calls] + + def assert_text_response(self) -> None: + response = self._require_response() + assert response.content, "Expected a non-empty text response, but content was empty." + + def assert_tool_called(self, name: str, predicate: Callable[[ToolCallView], bool] | None = None) -> None: + response = self._require_response() + matches = [tc for tc in response.tool_calls if tc.get("name") == name] + assert matches, f"Expected tool {name!r} to be called, but it wasn't. Called: {self._tool_call_names()}" + if predicate is not None: + assert any(predicate(ToolCallView(tc)) for tc in matches), ( + f"Tool {name!r} was called, but no call satisfied the given predicate." + ) + + def assert_tool_not_called(self, names: list[str]) -> None: + unexpected = set(self._tool_call_names()) & set(names) + assert not unexpected, f"Expected tools {sorted(names)} not to be called, but got: {sorted(unexpected)}" + + def assert_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" + + async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None: + response = self._require_response() + 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}" + ) + + async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: + return await self._judge_live(model, expectation, content, provider) + + async def _judge_live(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: + from .judge import JudgeAgent # noqa: PLC0415 + + judge = JudgeAgent() + judge.model = model + judge.provider = provider + return await judge.judge(expectation, content) + def __call__(self, func: Callable) -> Callable: if inspect.iscoroutinefunction(func): @@ -97,8 +170,6 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: class ToolCallView: - """Ergonomic, attribute-style view of a raw ``tool_calls`` dict, passed - to ``assert_tool_called``'s predicate.""" def __init__(self, data: dict) -> None: self.name = data.get("name", "") @@ -110,38 +181,18 @@ def __repr__(self) -> str: return f"ToolCallView(name={self.name!r}, args={self.args!r})" -class AgentRecordFake(_Recorder): - """Bound as ``agent`` by ``with Agent.record(cassette) as agent:``. - - Fluent testing handle around a record-and-replay session: ``prompt()`` - is async (it's the same real agent call underneath, just cached) and - each call mutates the handle's "current turn" state, which the - ``assert_*`` methods judge against — mirroring how a browser-testing - ``page`` object exposes assertions against the current page state. - - On a cassette miss, the real agent is called once and the response is - cached to disk (keyed by the conversation history so far, plus the new - message, so two sessions with different histories but the same latest - message text don't collide). On a hit, it's replayed with no live call. - - Entering the ``with`` block also binds this handle into the container - under the agent class's name, so any other instance of that class - created during the block (e.g. by application code under test) is - routed through the same recording session. - """ - +class AgentRecordFake(AgentFake): def __init__(self, real: Agent, cassette: str | None = None, messages: list | None = None) -> None: - super().__init__() self._real = real self.cassette: Path | None = Path(cassette) if cassette else None self._seed_messages: list = list(messages or []) - self._transcript: list[dict] = [] + self._records: list[dict] = [] self._real.messages = self._history # type: ignore[method-assign] - self.last_response: AgentResponse | None = None + self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None def _history(self) -> list: - return self._seed_messages + self._transcript + return self._seed_messages + self._records @staticmethod def _serialize(value: Any) -> Any: @@ -182,78 +233,39 @@ def _response_from_cache(value: Any) -> AgentResponse: return AgentResponse(content=_joined(value)) def _remember_turn(self, message: str, response: AgentResponse) -> None: - self._transcript.append({"role": "user", "content": message}) + self._records.append({"role": "user", "content": message}) turn: dict[str, Any] = {"role": "assistant", "content": response.content} if response.tool_calls: turn["tool_calls"] = response.tool_calls - self._transcript.append(turn) + self._records.append(turn) async def prompt(self, message: str, *, attachments: list[Document] | None = None) -> AgentResponse: - """Run (or replay) one turn and make it the "current" response that - assert_*() methods judge.""" - self._record_call(message, attachments) cassette, store = self._load() key = self._key(message, attachments) start = time.monotonic() if key in store: response = self._response_from_cache(store[key]) else: - response = await self._real._run(message, attachments=attachments) + response = await self._real.prompt(message, attachments=attachments) self._save(cassette, store, key, self._cache_prompt_value(response)) + response = self._real.runner()._apply_schema(response) self.last_elapsed = time.monotonic() - start - self.last_response = response + self._last_response = response self._remember_turn(message, response) return response async def stream(self, message: str) -> AsyncIterator[str]: - self._record_call(message, None) cassette, store = self._load() key = self._key(message, None) if key in store: value = store[key] - for chunk in value if isinstance(value, list) else [value]: - yield chunk - return - chunks = [chunk async for chunk in self._real._stream(message)] - self._save(cassette, store, key, chunks) + chunks = value if isinstance(value, list) else [value] + else: + chunks = [chunk async for chunk in self._real.stream(message)] + self._save(cassette, store, key, chunks) for chunk in chunks: yield chunk - - def _require_response(self) -> AgentResponse: - assert self.last_response is not None, "No prompt() call has been made yet." - return self.last_response - - def _tool_call_names(self) -> list[str]: - return [tc.get("name", "") for tc in self._require_response().tool_calls] - - def assert_text_response(self) -> None: - response = self._require_response() - assert response.content, "Expected a non-empty text response, but content was empty." - - def assert_tool_called(self, name: str, predicate: Callable[[ToolCallView], bool] | None = None) -> None: - response = self._require_response() - matches = [tc for tc in response.tool_calls if tc.get("name") == name] - assert matches, f"Expected tool {name!r} to be called, but it wasn't. Called: {self._tool_call_names()}" - if predicate is not None: - assert any(predicate(ToolCallView(tc)) for tc in matches), ( - f"Tool {name!r} was called, but no call satisfied the given predicate." - ) - - def assert_tool_not_called(self, names: list[str]) -> None: - unexpected = set(self._tool_call_names()) & set(names) - assert not unexpected, f"Expected tools {sorted(names)} not to be called, but got: {sorted(unexpected)}" - - def assert_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" - - async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None: - response = self._require_response() - 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}" - ) + self._remember_turn(message, AgentResponse(content=_joined(chunks))) async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: cassette, store = self._load() @@ -272,14 +284,6 @@ def _judge_key(model: str, expectation: str, content: str, provider: str | None ) return "judge:" + hashlib.sha256(payload.encode()).hexdigest() - async def _judge_live(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: - from .judge import JudgeAgent # noqa: PLC0415 - - judge = JudgeAgent() - judge.model = model - judge.provider = provider - return await judge.judge(expectation, content) - def _resolve_cassette(self, filename: str, qualname: str) -> None: here = Path(filename).parent if self.cassette is None: @@ -288,17 +292,11 @@ def _resolve_cassette(self, filename: str, qualname: str) -> None: self.cassette = here / self.cassette def __enter__(self) -> "AgentRecordFake": - from fastapi_startkit.application import app - caller = sys._getframe(1).f_code self._resolve_cassette(caller.co_filename, caller.co_qualname) - app().bind(type(self._real).__name__, self) return self def __exit__(self, *_exc: Any) -> bool: - from fastapi_startkit.application import app - - app().unbind(type(self._real).__name__) return False def __call__(self, func: Callable) -> Callable: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/tinker.py b/fastapi_startkit/src/fastapi_startkit/ai/tinker.py new file mode 100644 index 00000000..e69de29b diff --git a/fastapi_startkit/src/fastapi_startkit/ai/types.py b/fastapi_startkit/src/fastapi_startkit/ai/types.py new file mode 100644 index 00000000..62419dbc --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/types.py @@ -0,0 +1,5 @@ +from collections.abc import AsyncIterator +from typing import Callable + +Handler = Callable[[list], AsyncIterator] +Middleware = Callable[[list, Handler], AsyncIterator] diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 145a1cfb..83c8029f 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -1,17 +1,40 @@ +import json +import re import unittest +from typing import cast from unittest import mock import langchain.chat_models as chat_models -from langchain_core.messages import AIMessage, ToolCall +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, AIMessageChunk, ToolCall +from langchain_core.outputs import ChatGenerationChunk from langchain_core.tools import tool -from fastapi_startkit.ai import AIConfig, Document, fake_chat_model +from fastapi_startkit.ai import AIConfig, Document from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.runner import Runner from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse from fastapi_startkit.application import app +class StreamingToolFake(GenericFakeChatModel): + # GenericFakeChatModel can't stream tool calls; this emits tool-call chunks + # so streamed-tool-result behaviour stays testable without a shared fakes module. + def _stream(self, messages, stop=None, run_manager=None, **kwargs): + message = cast(AIMessage, next(self.messages)) + content = message.content if isinstance(message.content, str) else str(message.content) + for token in re.split(r"(\s)", content): + if token: + yield ChatGenerationChunk(message=AIMessageChunk(content=token, id=message.id)) + if message.tool_calls: + chunks = [ + {"name": c["name"], "args": json.dumps(c.get("args", {})), "id": c.get("id"), "index": i} + for i, c in enumerate(message.tool_calls) + ] + yield ChatGenerationChunk(message=AIMessageChunk(content="", tool_call_chunks=chunks, id=message.id)) + + @tool def search_jobs(query: str) -> str: """Search the job board for roles matching the query.""" @@ -32,12 +55,9 @@ def setUp(self): container.bind("ai", AIConfig()) container.make("config").set("ai", AIConfig()) - def setup_agent(self, turns: list[AIMessage]): - model = fake_chat_model(turns) - patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: model) - patcher.start() - self.addCleanup(patcher.stop) - return model + def setup_agent(self, turns: list, agent_cls: type[Agent] = Agent): + Ai.fake(agent_cls.__name__, turns) + self.addCleanup(Ai.reset_fakes) async def test_prompt_returns_agent_response(self): self.setup_agent([AIMessage(content="hello back")]) @@ -47,7 +67,6 @@ async def test_prompt_returns_agent_response(self): self.assertIsInstance(result, AgentResponse) self.assertEqual(result.content, "hello back") - agent.assert_prompted() async def test_search_jobs_tool_returns_listing(self): self.setup_agent( @@ -56,7 +75,8 @@ async def test_search_jobs_tool_returns_listing(self): content="", tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], ), - ] + ], + JobAssistant, ) result = await JobAssistant().prompt("find me a python job") @@ -79,7 +99,7 @@ async def test_build_model_passes_langchain_provider_key(self): def fake_init(model, **kwargs): captured["model"] = model captured["provider"] = kwargs.get("model_provider") - return fake_chat_model([AIMessage(content="ok")]) + return GenericFakeChatModel(messages=iter([AIMessage(content="ok")])) patcher = mock.patch.object(chat_models, "init_chat_model", fake_init) patcher.start() @@ -101,8 +121,6 @@ async def test_stream_yields_tokens_from_the_model(self): self.assertEqual("".join(chunks), "streamed reply") async def test_middleware_streams_token_by_token_and_runs_after_hook(self): - self.setup_agent([AIMessage(content="one two three")]) - events: list = [] class Logger: @@ -114,6 +132,8 @@ class LoggedAgent(Agent): def middleware(self): return [Logger()] + self.setup_agent([AIMessage(content="one two three")], LoggedAgent) + chunks = [chunk async for chunk in LoggedAgent().stream("hi")] # Middleware must not buffer: the model's tokens arrive as separate chunks... @@ -123,8 +143,6 @@ def middleware(self): self.assertEqual(events, ["before", "after"]) async def test_middleware_after_hook_runs_on_prompt(self): - self.setup_agent([AIMessage(content="done")]) - events: list = [] class Logger: @@ -136,20 +154,25 @@ class LoggedAgent(Agent): def middleware(self): return [Logger()] + self.setup_agent([AIMessage(content="done")], LoggedAgent) + result = await LoggedAgent().prompt("hi") self.assertEqual(result.content, "done") self.assertEqual(events, ["before", "after"]) async def test_stream_yields_tool_result_without_calling_model_again(self): - self.setup_agent( - [ - AIMessage( - content="", - tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], - ), - ] + Ai._fakes["JobAssistant"] = StreamingToolFake( + messages=iter( + [ + AIMessage( + content="", + tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], + ), + ] + ) ) + self.addCleanup(Ai.reset_fakes) chunks = [chunk async for chunk in JobAssistant().stream("find me a python job")] @@ -167,7 +190,7 @@ def test_resolve_model_prefers_explicit_override(self): self.assertEqual(Ai()._resolve_model(Agent(), "my-model"), "my-model") def test_instructions_lead_the_message_list(self): - messages = JobAssistant()._build_messages("find me a job") + messages = Runner(JobAssistant())._build_messages("find me a job") self.assertEqual(messages[0], {"role": "system", "content": "You help users find jobs."}) self.assertEqual(sum(m.get("role") == "system" for m in messages), 1) @@ -177,19 +200,19 @@ class DynamicAgent(Agent): def instructions(self) -> str: return "Computed identity." - messages = DynamicAgent()._build_messages("hi") + messages = Runner(DynamicAgent())._build_messages("hi") self.assertEqual(messages[0], {"role": "system", "content": "Computed identity."}) def test_no_instructions_prepends_no_system_message(self): - messages = Agent()._build_messages("hi") + messages = Runner(Agent())._build_messages("hi") self.assertTrue(all(m.get("role") != "system" for m in messages)) def test_build_messages_inlines_text_attachment(self): doc = Document(content="Q3 revenue was $1.2M.", name="q3-report.txt") - messages = Agent()._build_messages("Summarise this report.", attachments=[doc]) + messages = Runner(Agent())._build_messages("Summarise this report.", attachments=[doc]) user_content = messages[-1]["content"] self.assertEqual(user_content[0], {"type": "text", "text": "Summarise this report."}) @@ -199,7 +222,7 @@ def test_build_messages_inlines_text_attachment(self): def test_build_messages_encodes_binary_attachment_as_file_block(self): doc = Document(content=b"%PDF-1.7 ...", name="q3.pdf", media_type="application/pdf") - messages = Agent()._build_messages("Summarise", attachments=[doc]) + messages = Runner(Agent())._build_messages("Summarise", attachments=[doc]) block = messages[-1]["content"][1] self.assertEqual(block["type"], "file") diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 42ff7adb..6bcd8187 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -1,14 +1,3 @@ -"""Tests for Agent.fake() / Agent.record() and the assert_prompted/reset helpers. - -``Agent.fake()`` registers a fixed, ordered list of replies as a deterministic -stand-in chat model (see ``Ai.fake()``) for the duration of a ``with`` block — -each call to ``prompt()``/``stream()`` replays the next reply, going through -the real message-building / pipeline / tool-execution path, only the model at -the bottom is swapped. ``Agent.record()`` binds a record-and-replay stand-in: -on a cassette miss it calls the real agent once and caches the response to -disk; on a hit it replays from the cassette without calling the agent again. -""" - import json import os import tempfile @@ -82,102 +71,88 @@ async def run(): self.assertEqual(result.content, "decorated reply") async def test_assert_prompted_passes_after_one_call(self): - agent = SimpleAgent() - with SimpleAgent.fake(["ok"]): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("first") agent.assert_prompted() async def test_assert_prompted_times_2_passes_after_exactly_2_calls(self): - agent = SimpleAgent() - with SimpleAgent.fake(["ok", "ok"]): + with SimpleAgent.fake(["ok", "ok"]) as agent: await agent.prompt("first") await agent.prompt("second") agent.assert_prompted(times=2) async def test_assert_prompted_times_fails_when_count_mismatch(self): - agent = SimpleAgent() - with SimpleAgent.fake(["ok"]): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("only once") with self.assertRaises(AssertionError): agent.assert_prompted(times=2) def test_assert_prompted_fails_when_never_called(self): - agent = SimpleAgent() - - with self.assertRaises(AssertionError): - agent.assert_prompted() + with SimpleAgent.fake([]) as agent: + with self.assertRaises(AssertionError): + agent.assert_prompted() def test_assert_prompted_times_zero_passes_when_never_called(self): - agent = SimpleAgent() - agent.assert_prompted(times=0) + with SimpleAgent.fake([]) as agent: + agent.assert_prompted(times=0) def test_assert_not_prompted_passes_when_no_calls_made(self): - agent = SimpleAgent() - agent.assert_not_prompted() + with SimpleAgent.fake([]) as agent: + agent.assert_not_prompted() async def test_assert_not_prompted_fails_after_one_call(self): - agent = SimpleAgent() - with SimpleAgent.fake(["ok"]): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("a prompt") with self.assertRaises(AssertionError): agent.assert_not_prompted() async def test_reset_clears_call_log(self): - agent = SimpleAgent() - with SimpleAgent.fake(["ok"]): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("first") - self.assertEqual(len(agent._call_log), 1) + self.assertEqual(len(agent._prompts), 1) - agent.reset() - self.assertEqual(agent._call_log, []) + agent.reset() + self.assertEqual(agent._prompts, []) - def test_reset_returns_agent_for_chaining(self): - agent = SimpleAgent() - result = agent.reset() - self.assertIs(result, agent) + def test_reset_returns_the_handle_for_chaining(self): + with SimpleAgent.fake([]) as agent: + self.assertIs(agent.reset(), agent) async def test_assert_not_prompted_passes_after_reset(self): - agent = SimpleAgent() - with SimpleAgent.fake(["ok"]): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("call before reset") - - agent.reset() - agent.assert_not_prompted() + agent.reset() + agent.assert_not_prompted() async def test_fake_rebinding_overrides_previous(self): - agent = SimpleAgent() - - with SimpleAgent.fake(["first fake"]): + with SimpleAgent.fake(["first fake"]) as agent: self.assertEqual((await agent.prompt("call")).content, "first fake") - with SimpleAgent.fake(["second fake"]): + with SimpleAgent.fake(["second fake"]) as agent: self.assertEqual((await agent.prompt("call again")).content, "second fake") async def test_stream_returns_fake_response(self): - agent = SimpleAgent() - with SimpleAgent.fake(["Faked stream!"]): + with SimpleAgent.fake(["Faked stream!"]) as agent: chunks = [chunk async for chunk in agent.stream("hello world")] - self.assertEqual("".join(chunks), "Faked stream!") - self.assertGreater(len(chunks), 1) - agent.assert_prompted(times=1) + self.assertEqual("".join(chunks), "Faked stream!") + self.assertGreater(len(chunks), 1) + agent.assert_prompted(times=1) async def test_stream_replays_the_registered_text_exactly(self): - agent = SimpleAgent() - with SimpleAgent.fake(["Hello there, friend"]): + with SimpleAgent.fake(["Hello there, friend"]) as agent: chunks = [chunk async for chunk in agent.stream("hi")] - self.assertEqual("".join(chunks), "Hello there, friend") + self.assertEqual("".join(chunks), "Hello there, friend") async def test_stream_records_one_call_not_two(self): - agent = SimpleAgent() - with SimpleAgent.fake(["x"]): + with SimpleAgent.fake(["x"]) as agent: [chunk async for chunk in agent.stream("once")] - # Streaming must log exactly one prompt — not one for stream + one for prompt. - agent.assert_prompted(times=1) + # Streaming must record exactly one call — not one for stream + one for prompt. + agent.assert_prompted(times=1) class TestAgentRecord(unittest.IsolatedAsyncioTestCase): @@ -188,7 +163,7 @@ async def fake_run(agent_self, message, **kwargs): calls.append(message) return AgentResponse(content=content) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) return calls @@ -197,8 +172,8 @@ async def test_first_run_records_response_to_cassette(self): calls = self.setup_agent("recorded reply") with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with SimpleAgent.record(cassette): - result = await SimpleAgent().prompt("hello") + with SimpleAgent.record(cassette) as agent: + result = await agent.prompt("hello") self.assertEqual(result.content, "recorded reply") self.assertEqual(calls, ["hello"]) @@ -211,10 +186,10 @@ async def test_second_run_replays_without_calling_run(self): calls = self.setup_agent("recorded reply") with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with SimpleAgent.record(cassette): - await SimpleAgent().prompt("hello") - with SimpleAgent.record(cassette): - replayed = await SimpleAgent().prompt("hello") + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + with SimpleAgent.record(cassette) as agent: + replayed = await agent.prompt("hello") self.assertEqual(replayed.content, "recorded reply") self.assertEqual(calls, ["hello"]) @@ -228,12 +203,12 @@ async def changed_run(s, m, **k): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with mock.patch.object(SimpleAgent, "_run", first_run): - with SimpleAgent.record(cassette): - await SimpleAgent().prompt("hello") - with mock.patch.object(SimpleAgent, "_run", changed_run): - with SimpleAgent.record(cassette): - result = await SimpleAgent().prompt("hello") + with mock.patch.object(SimpleAgent, "prompt", first_run): + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + with mock.patch.object(SimpleAgent, "prompt", changed_run): + with SimpleAgent.record(cassette) as agent: + result = await agent.prompt("hello") self.assertEqual(result.content, "from first record") @@ -241,9 +216,9 @@ async def test_distinct_messages_are_recorded_separately(self): calls = self.setup_agent("reply") with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with SimpleAgent.record(cassette): - await SimpleAgent().prompt("hello") - await SimpleAgent().prompt("goodbye") + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + await agent.prompt("goodbye") self.assertEqual(calls, ["hello", "goodbye"]) with open(cassette) as f: @@ -257,7 +232,7 @@ async def fake_stream(agent_self, message, **kwargs): for chunk in chunks: yield chunk - patcher = mock.patch.object(SimpleAgent, "_stream", fake_stream) + patcher = mock.patch.object(SimpleAgent, "stream", fake_stream) patcher.start() self.addCleanup(patcher.stop) return calls @@ -266,8 +241,8 @@ async def test_stream_first_run_records_chunk_list_to_cassette(self): calls = self.setup_stream(["Hel", "lo!"]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "s.json") - with SimpleAgent.record(cassette): - chunks = [c async for c in SimpleAgent().stream("hi")] + with SimpleAgent.record(cassette) as agent: + chunks = [c async for c in agent.stream("hi")] self.assertEqual(chunks, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) @@ -278,10 +253,10 @@ async def test_stream_second_run_replays_chunks_without_calling_stream(self): calls = self.setup_stream(["Hel", "lo!"]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "s.json") - with SimpleAgent.record(cassette): - [c async for c in SimpleAgent().stream("hi")] - with SimpleAgent.record(cassette): - replayed = [c async for c in SimpleAgent().stream("hi")] + with SimpleAgent.record(cassette) as agent: + [c async for c in agent.stream("hi")] + with SimpleAgent.record(cassette) as agent: + replayed = [c async for c in agent.stream("hi")] self.assertEqual(replayed, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) # real stream invoked only on the first run @@ -290,9 +265,9 @@ async def test_prompt_reads_a_stream_recorded_cassette_as_joined_content(self): self.setup_stream(["Hel", "lo!"]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "s.json") - with SimpleAgent.record(cassette): - [c async for c in SimpleAgent().stream("hi")] - with SimpleAgent.record(cassette): - response = await SimpleAgent().prompt("hi") + with SimpleAgent.record(cassette) as agent: + [c async for c in agent.stream("hi")] + with SimpleAgent.record(cassette) as agent: + response = await agent.prompt("hi") self.assertEqual(response.content, "Hello!") diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py index 45121c91..1c0d08fa 100644 --- a/fastapi_startkit/tests/ai/test_agent_record_fluent.py +++ b/fastapi_startkit/tests/ai/test_agent_record_fluent.py @@ -23,6 +23,7 @@ from langchain_core.messages import AIMessage, HumanMessage from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.runner import Runner from fastapi_startkit.ai.response import AgentResponse from fastapi_startkit.ai.testing import AgentRecordFake @@ -44,7 +45,7 @@ async def fake_run(agent_self, message, **kwargs): content, tool_calls = queue.pop(0) return AgentResponse(content=content, tool_calls=tool_calls or []) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) @@ -85,7 +86,7 @@ def setup_agent(self, content, tool_calls=None): async def fake_run(agent_self, message, **kwargs): return AgentResponse(content=content, tool_calls=tool_calls or []) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) @@ -116,7 +117,7 @@ def setup_agent(self, tool_calls): async def fake_run(agent_self, message, **kwargs): return AgentResponse(content="", tool_calls=tool_calls) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) @@ -156,7 +157,7 @@ def setup_agent(self, tool_calls): async def fake_run(agent_self, message, **kwargs): return AgentResponse(content="Hello!", tool_calls=tool_calls) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) @@ -181,7 +182,7 @@ def setup_agent(self): async def fake_run(agent_self, message, **kwargs): return AgentResponse(content="Hello!") - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) @@ -212,7 +213,7 @@ async def test_seed_messages_are_included_when_building_the_real_agents_messages seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] with tempfile.TemporaryDirectory() as tmp: with SimpleAgent.record(os.path.join(tmp, "c.json"), messages=seed) as agent: - built = agent._real._build_messages("suggest python developer jobs") + built = Runner(agent._real)._build_messages("suggest python developer jobs") self.assertEqual(built[0], seed[0]) self.assertEqual(built[1], seed[1]) @@ -225,7 +226,7 @@ async def test_same_followup_text_with_different_seed_history_does_not_collide(s async def run_a(agent_self, message, **kwargs): return AgentResponse(content="job list A") - with mock.patch.object(SimpleAgent, "_run", run_a): + with mock.patch.object(SimpleAgent, "prompt", run_a): with SimpleAgent.record(cassette) as agent: response_a = await agent.prompt("suggest python developer jobs") @@ -233,7 +234,7 @@ async def run_b(agent_self, message, **kwargs): return AgentResponse(content="job list B") seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] - with mock.patch.object(SimpleAgent, "_run", run_b): + with mock.patch.object(SimpleAgent, "prompt", run_b): with SimpleAgent.record(cassette, messages=seed) as agent: response_b = await agent.prompt("suggest python developer jobs") @@ -246,7 +247,7 @@ def setup_agent(self, content): async def fake_run(agent_self, message, **kwargs): return AgentResponse(content=content) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) @@ -339,8 +340,8 @@ async def fake_run(agent_self, message, **kwargs): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with mock.patch.object(SimpleAgent, "_run", fake_run): - with SimpleAgent.record(cassette): - result = await SimpleAgent().prompt("hello") + with mock.patch.object(SimpleAgent, "prompt", fake_run): + with SimpleAgent.record(cassette) as agent: + result = await agent.prompt("hello") self.assertEqual(result.content, "recorded reply") diff --git a/fastapi_startkit/tests/ai/test_agent_schema.py b/fastapi_startkit/tests/ai/test_agent_schema.py index 40306ea3..347abd7a 100644 --- a/fastapi_startkit/tests/ai/test_agent_schema.py +++ b/fastapi_startkit/tests/ai/test_agent_schema.py @@ -53,11 +53,11 @@ async def fake_run(self, message, **kwargs): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "user.json") - with mock.patch.object(UserAgent, "_run", fake_run): - with UserAgent.record(cassette): - recorded = await UserAgent().prompt("get the user") - with UserAgent.record(cassette): - replayed = await UserAgent().prompt("get the user") + with mock.patch.object(UserAgent, "prompt", fake_run): + with UserAgent.record(cassette) as agent: + recorded = await agent.prompt("get the user") + with UserAgent.record(cassette) as agent: + replayed = await agent.prompt("get the user") self.assertEqual(recorded.parsed, User(id="u-9", name="Sam")) self.assertEqual(replayed.parsed, User(id="u-9", name="Sam")) diff --git a/fastapi_startkit/tests/ai/test_graph_agent.py b/fastapi_startkit/tests/ai/test_graph_agent.py new file mode 100644 index 00000000..b5d0d858 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_graph_agent.py @@ -0,0 +1,77 @@ +"""Tests for GraphAgent — a LangGraph tool-calling loop layered on Agent. + +Uses the same ``Agent.fake()`` model-swap as the other agent tests: scripted +turns drive the real graph (message assembly -> llm node -> tool node -> loop), +only the model at the bottom is a fake. +""" + +import unittest + +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from langgraph.graph import END, START, StateGraph + +from fastapi_startkit.ai.ai import Ai +from fastapi_startkit.ai.graph import GraphAgent, GraphRunner, GraphState + + +@tool +def add(a: int, b: int) -> int: + """Add ``a`` and ``b``.""" + return a + b + + +class MathAgent(GraphAgent): + def instructions(self): + return "You are a helpful assistant tasked with performing arithmetic." + + def tools(self): + return [add] + + def graph(self, runner: GraphRunner) -> StateGraph: + graph = StateGraph(GraphState) + graph.add_node("llm", runner.llm) + graph.add_node("tools", runner.call_tools) + graph.add_edge(START, "llm") + graph.add_conditional_edges("llm", runner.route, ["tools", END]) + graph.add_edge("tools", "llm") + return graph + + +# A tool-calling turn (add 3 + 4), then the model's final natural-language answer. +SCRIPT = [ + AIMessage(content="", tool_calls=[{"name": "add", "args": {"a": 3, "b": 4}, "id": "call_0"}]), + "The answer is 7", +] + + +class TestGraphAgent(unittest.IsolatedAsyncioTestCase): + def tearDown(self): + Ai.reset_fakes() + + async def test_prompt_runs_the_tool_loop(self): + with MathAgent.fake(list(SCRIPT)): + response = await MathAgent().prompt("Add 3 and 4.") + + self.assertEqual(response.content, "The answer is 7") + self.assertTrue(any(call["name"] == "add" for call in response.tool_calls)) + + async def test_prompt_executes_the_tool_and_feeds_it_back(self): + with MathAgent.fake(list(SCRIPT)): + response = await MathAgent().prompt("Add 3 and 4.") + + # The add tool ran (3 + 4) and its result was fed back for the final turn. + add_call = next(call for call in response.tool_calls if call["name"] == "add") + self.assertEqual(add_call["args"], {"a": 3, "b": 4}) + + async def test_stream_yields_the_final_answer_token_by_token(self): + with MathAgent.fake(["The answer is 7"]): + chunks = [chunk async for chunk in MathAgent().stream("What is 3 plus 4?")] + + self.assertEqual("".join(chunks), "The answer is 7") + self.assertGreater(len(chunks), 1) + + async def test_records_the_prompt_call(self): + with MathAgent.fake(["7"]) as agent: + await agent.prompt("Add 3 and 4.") + agent.assert_prompted(times=1) diff --git a/fastapi_startkit/tests/ai/test_judge_agent.py b/fastapi_startkit/tests/ai/test_judge_agent.py index e1d5a964..bcb71569 100644 --- a/fastapi_startkit/tests/ai/test_judge_agent.py +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -13,6 +13,7 @@ from langchain_core.messages import AIMessage from fastapi_startkit.ai.judge import JudgeAgent, Verdict +from fastapi_startkit.ai.runner import Runner from fastapi_startkit.ai.response import AgentResponse @@ -43,7 +44,7 @@ async def ainvoke(self, messages): seen["messages"] = messages return AIMessage(content='{"passed": true, "reasoning": "ok"}') - with mock.patch.object(JudgeAgent, "_build_model", lambda self, *a, **k: Capturing()): + with mock.patch.object(Runner, "_build_model", lambda self, *a, **k: Capturing()): await JudgeAgent().judge("The llm should respond with greetings", "Hello there!") blob = " ".join(str(getattr(m, "content", m)) for m in seen["messages"]) @@ -71,7 +72,7 @@ async def fake_run(agent_self, message, **kwargs): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "judge.json") - with mock.patch.object(JudgeAgent, "_run", fake_run): + with mock.patch.object(JudgeAgent, "prompt", fake_run): with JudgeAgent.record(cassette) as agent: response = await agent.prompt("grade this") diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py index bfa60a46..d0957c64 100644 --- a/fastapi_startkit/tests/ai/test_structured_output.py +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -116,15 +116,6 @@ def tools(self): self.assertIs(result, fake) self.assertEqual(fake.calls, [("bind_tools", [noop])]) - def test_streaming_skips_structured_output(self): - fake = _FakeModel() - self._patch(fake) - - result = Ai().build(ToolMovieAgent(), structured=False) - - self.assertIs(result, fake) - self.assertEqual(fake.calls, [("bind_tools", [noop])]) - def test_no_schema_no_tools_returns_the_plain_model(self): fake = _FakeModel() self._patch(fake) @@ -142,7 +133,7 @@ class Model: async def ainvoke(self, messages): return payload - result = await Runner(MovieAgent(), Model()).run(["hi"]) + result = await Runner(MovieAgent())._invoke(Model(), ["hi"]) self.assertEqual(result, payload) @@ -151,7 +142,7 @@ class Model: async def ainvoke(self, messages): return AIMessage(content="", tool_calls=[_real_tool_call(query="hello")]) - result = await Runner(ToolMovieAgent(), Model()).run(["hi"]) + result = await Runner(ToolMovieAgent())._invoke(Model(), ["hi"]) self.assertEqual(result.content, "hello") @@ -160,7 +151,7 @@ class TestResponseMapping(unittest.TestCase): def test_unwraps_include_raw_into_parsed_and_content(self): parsed = Movie(title="Inception", year=2010) - response = MovieAgent()._to_agent_response( + response = Runner(MovieAgent())._to_agent_response( {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} ) diff --git a/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py b/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py index 675cf79a..f86fc1ef 100644 --- a/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py +++ b/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py @@ -9,10 +9,10 @@ class TestRelationships(TestCase): async def test_can_get_polymorphic_relation(self): likes = await Like.get() for like in likes: - record = await like.record + record = await like.log assert isinstance(record, (Articles, Product)) async def test_can_get_eager_load_polymorphic_relation(self): likes = await Like.with_("record").get() for like in likes: - assert isinstance(like.record, (Articles, Product)) + assert isinstance(like.log, (Articles, Product)) diff --git a/fastapi_startkit/uv.lock b/fastapi_startkit/uv.lock index c9ea4cba..032c1a36 100644 --- a/fastapi_startkit/uv.lock +++ b/fastapi_startkit/uv.lock @@ -543,6 +543,7 @@ dependencies = [ ai = [ { name = "langchain" }, { name = "langchain-core" }, + { name = "langgraph" }, ] database = [ { name = "faker" }, @@ -605,6 +606,7 @@ requires-dist = [ { name = "jinja2", marker = "extra == 'vite'", specifier = ">=3.1" }, { name = "langchain", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "langchain-core", marker = "extra == 'ai'", specifier = ">=1.0.0" }, + { name = "langgraph", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "markupsafe", marker = "extra == 'inertia'", specifier = ">=2.0" }, { name = "pendulum", specifier = ">=3.1.0,<4.0.0" }, { name = "pydantic", specifier = ">=2.12.5" },