From 556f61279cb6fb293f9933cb5b4f5ea73c33f875 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Thu, 16 Jul 2026 16:33:34 -0700 Subject: [PATCH 01/13] 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 02/13] 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 03/13] 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 04/13] 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 05/13] 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 06/13] 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 07/13] 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" }, From 58e2d1390b6d33c573f5ed70cdbed44b56a5432d Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 24 Jul 2026 13:48:39 -0700 Subject: [PATCH 08/13] feat: fix the ai chat --- .vscode/settings.json | 2 +- example/agents/.vscode/settings.json | 3 + example/agents/app/agents/chat.py | 24 +--- example/agents/app/agents/graph_agent.py | 33 +++++- .../app/providers/langchain_provider.py | 29 ++--- example/agents/app/requests/chat.py | 1 + .../agents/resources/js/components/Chat.tsx | 110 ++++++++++++++++++ .../agents/resources/js/pages/chat/Index.tsx | 99 +--------------- .../resources/js/pages/chat/sales/Index.tsx | 5 + example/agents/routes/api.py | 30 +++-- .../tests/features/test_chat_controller.py | 14 +-- .../tests/units/agents/basic_job_search.json | 15 +++ .../tests/units/agents/record_sales.json | 6 + .../tests/units/agents/record_stream.json | 10 -- .../tests/units/agents/stream_job_search.json | 21 ++++ .../units/agents/test_langchain_agent.py | 26 ----- .../tests/units/agents/test_router_agent.py | 6 +- .../tests/units/agents/test_sales_agent.py | 51 ++++++++ fastapi_startkit/.vscode/settings.json | 3 + .../src/fastapi_startkit/ai/__init__.py | 4 +- .../src/fastapi_startkit/ai/agent.py | 23 ++-- .../src/fastapi_startkit/ai/graph.py | 108 ++++++++++++++--- .../src/fastapi_startkit/ai/testing.py | 57 ++++++--- fastapi_startkit/tests/ai/test_agent_fake.py | 13 ++- fastapi_startkit/tests/ai/test_graph_agent.py | 37 +++++- 25 files changed, 479 insertions(+), 251 deletions(-) create mode 100644 example/agents/.vscode/settings.json create mode 100644 example/agents/resources/js/components/Chat.tsx create mode 100644 example/agents/resources/js/pages/chat/sales/Index.tsx create mode 100644 example/agents/tests/units/agents/basic_job_search.json create mode 100644 example/agents/tests/units/agents/record_sales.json delete mode 100644 example/agents/tests/units/agents/record_stream.json create mode 100644 example/agents/tests/units/agents/stream_job_search.json delete mode 100644 example/agents/tests/units/agents/test_langchain_agent.py create mode 100644 example/agents/tests/units/agents/test_sales_agent.py create mode 100644 fastapi_startkit/.vscode/settings.json diff --git a/.vscode/settings.json b/.vscode/settings.json index 83ddff40..ff5300ef 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,3 +1,3 @@ { - "python.languageServer": "Default" + "python.languageServer": "None" } \ No newline at end of file diff --git a/example/agents/.vscode/settings.json b/example/agents/.vscode/settings.json new file mode 100644 index 00000000..ff5300ef --- /dev/null +++ b/example/agents/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.languageServer": "None" +} \ No newline at end of file diff --git a/example/agents/app/agents/chat.py b/example/agents/app/agents/chat.py index 4e02a9e2..144d1101 100644 --- a/example/agents/app/agents/chat.py +++ b/example/agents/app/agents/chat.py @@ -1,31 +1,15 @@ -from typing import Callable - -from fastapi_startkit.ai import Agent, GraphAgent, Middleware +from fastapi_startkit.ai import Agent, Middleware +from langchain_core.tools import BaseTool from app.middleware.agent_logger import AgentLogger from app.tools.job_search_tool import job_search_tool -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() - +class ChatAgent(Agent): def middleware(self) -> list[Middleware]: return [AgentLogger()] - def tools(self) -> list[Callable]: + def tools(self) -> list[BaseTool]: return [job_search_tool] def instructions(self) -> str: diff --git a/example/agents/app/agents/graph_agent.py b/example/agents/app/agents/graph_agent.py index c5d4dd03..f149fa96 100644 --- a/example/agents/app/agents/graph_agent.py +++ b/example/agents/app/agents/graph_agent.py @@ -1,9 +1,30 @@ -from fastapi_startkit.ai import GraphAgent, GraphRunner -from langgraph.graph import StateGraph +from fastapi_startkit.ai import GraphAgent, GraphRunner, Middleware +from fastapi_startkit.ai.graph import AgentState +from langchain_core.tools import BaseTool +from langgraph.graph import END, START, StateGraph +from langgraph.types import Checkpointer + +from app.middleware.agent_logger import AgentLogger +from app.tools.job_search_tool import job_search_tool class SalesAgent(GraphAgent): - def checkpointer(self): - pass - def graph(self, runner: GraphRunner) -> StateGraph: - return StateGraph() + def middleware(self) -> list[Middleware]: + return [AgentLogger()] + + def tools(self) -> list[BaseTool]: + return [job_search_tool] + + async def graph(self, runner: GraphRunner) -> StateGraph: + from fastapi_startkit.application import app + + checkpointer: Checkpointer = await app().make("checkpointer") + + return ( + StateGraph(AgentState) + .add_node("llm", runner.llm) + .add_node("tools", runner.call_tools) + .add_edge(START, "llm") + .add_conditional_edges("llm", runner.route, ["tools", END]) + .add_edge("tools", "llm") + ).compile(checkpointer=checkpointer) diff --git a/example/agents/app/providers/langchain_provider.py b/example/agents/app/providers/langchain_provider.py index 67fcaeb3..ec7299c2 100644 --- a/example/agents/app/providers/langchain_provider.py +++ b/example/agents/app/providers/langchain_provider.py @@ -1,37 +1,26 @@ from fastapi_startkit.support import Provider from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver -from psycopg.rows import dict_row +from psycopg import AsyncConnection +from psycopg.rows import DictRow, 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( + self.pool = AsyncConnectionPool[AsyncConnection[DictRow]]( conninfo=uri, open=False, kwargs={"autocommit": True, "row_factory": dict_row}, ) - self._saver: AsyncPostgresSaver | None = None + self.saver: AsyncPostgresSaver | None = None async def resolve(self) -> AsyncPostgresSaver: - if self._saver is None: - await self._pool.open() - saver = AsyncPostgresSaver(self._pool) + if self.saver is None: + await self.pool.open() + saver = AsyncPostgresSaver(self.pool) await saver.setup() - self._saver = saver - return self._saver + self.saver = saver + return self.saver def __await__(self): return self.resolve().__await__() diff --git a/example/agents/app/requests/chat.py b/example/agents/app/requests/chat.py index 163422ac..706099db 100644 --- a/example/agents/app/requests/chat.py +++ b/example/agents/app/requests/chat.py @@ -3,3 +3,4 @@ class ChatRequest(BaseModel): message: str = Field(...) + thread_id: str = Field("default") diff --git a/example/agents/resources/js/components/Chat.tsx b/example/agents/resources/js/components/Chat.tsx new file mode 100644 index 00000000..58bd29d2 --- /dev/null +++ b/example/agents/resources/js/components/Chat.tsx @@ -0,0 +1,110 @@ +import { useState, useRef, useEffect } from "react" + +type Message = { + role: "user" | "assistant" + content: string +} + +type ChatProps = { + title?: string + endpoint?: string + placeholder?: string + emptyState?: string +} + +export default function Chat({ + title = "Chat", + endpoint = "/chat/stream", + placeholder = "Type a message...", + emptyState = "Send a message to start chatting.", +}: ChatProps) { + const [messages, setMessages] = useState([]) + const [input, setInput] = useState("") + const [loading, setLoading] = useState(false) + const bottomRef = useRef(null) + + useEffect(() => { + bottomRef.current?.scrollIntoView({ behavior: "smooth" }) + }, [messages]) + + const handleSubmit = async (e: { preventDefault(): void }) => { + e.preventDefault() + if (!input.trim() || loading) return + + const userMessage = input.trim() + setInput("") + setMessages(prev => [...prev, { role: "user", content: userMessage }]) + setLoading(true) + setMessages(prev => [...prev, { role: "assistant", content: "" }]) + + try { + const response = await fetch(endpoint, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ message: userMessage }), + }) + + const reader = response.body?.getReader() + const decoder = new TextDecoder() + if (!reader) return + + while (true) { + const { done, value } = await reader.read() + if (done) break + const chunk = decoder.decode(value, { stream: true }) + setMessages(prev => { + const updated = [...prev] + updated[updated.length - 1] = { + role: "assistant", + content: updated[updated.length - 1].content + chunk, + } + return updated + }) + } + } finally { + setLoading(false) + } + } + + return ( +
+

{title}

+ +
+ {messages.length === 0 && ( +

{emptyState}

+ )} + {messages.map((msg, i) => ( +
+
+ {msg.content || (loading && i === messages.length - 1 ? "▋" : "")} +
+
+ ))} +
+
+ +
+ setInput(e.target.value)} + placeholder={placeholder} + disabled={loading} + /> + +
+
+ ) +} diff --git a/example/agents/resources/js/pages/chat/Index.tsx b/example/agents/resources/js/pages/chat/Index.tsx index d3a112a4..79b2d82f 100644 --- a/example/agents/resources/js/pages/chat/Index.tsx +++ b/example/agents/resources/js/pages/chat/Index.tsx @@ -1,98 +1,5 @@ -import { useState, useRef, useEffect } from "react" +import Chat from "../../components/Chat" -type Message = { - role: "user" | "assistant" - content: string -} - -export default function Chat() { - const [messages, setMessages] = useState([]) - const [input, setInput] = useState("") - const [loading, setLoading] = useState(false) - const bottomRef = useRef(null) - - useEffect(() => { - bottomRef.current?.scrollIntoView({ behavior: "smooth" }) - }, [messages]) - - const handleSubmit = async (e: { preventDefault(): void }) => { - e.preventDefault() - if (!input.trim() || loading) return - - const userMessage = input.trim() - setInput("") - setMessages(prev => [...prev, { role: "user", content: userMessage }]) - setLoading(true) - setMessages(prev => [...prev, { role: "assistant", content: "" }]) - - try { - const response = await fetch("/chat/stream", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ message: userMessage }), - }) - - const reader = response.body?.getReader() - const decoder = new TextDecoder() - if (!reader) return - - while (true) { - const { done, value } = await reader.read() - if (done) break - const chunk = decoder.decode(value, { stream: true }) - setMessages(prev => { - const updated = [...prev] - updated[updated.length - 1] = { - role: "assistant", - content: updated[updated.length - 1].content + chunk, - } - return updated - }) - } - } finally { - setLoading(false) - } - } - - return ( -
-

Chat

- -
- {messages.length === 0 && ( -

Send a message to start chatting.

- )} - {messages.map((msg, i) => ( -
-
- {msg.content || (loading && i === messages.length - 1 ? "▋" : "")} -
-
- ))} -
-
- -
- setInput(e.target.value)} - placeholder="Type a message..." - disabled={loading} - /> - -
-
- ) +export default function Index() { + return } diff --git a/example/agents/resources/js/pages/chat/sales/Index.tsx b/example/agents/resources/js/pages/chat/sales/Index.tsx new file mode 100644 index 00000000..60853b8b --- /dev/null +++ b/example/agents/resources/js/pages/chat/sales/Index.tsx @@ -0,0 +1,5 @@ +import Chat from "../../../components/Chat" + +export default function Index() { + return +} diff --git a/example/agents/routes/api.py b/example/agents/routes/api.py index d78f5073..e3939fc5 100644 --- a/example/agents/routes/api.py +++ b/example/agents/routes/api.py @@ -1,9 +1,9 @@ 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.agents.chat import ChatAgent +from app.agents.graph_agent import SalesAgent from app.requests.chat import ChatRequest api = APIRouter() @@ -21,18 +21,32 @@ 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) + response = await ChatAgent().prompt(request.message) return {"content": response.content} @api.post("/chat/stream") async def chat_stream(request: ChatRequest): async def generate(): - async for chunk in RouterAgent().stream(request.message): + async for chunk in ChatAgent().stream(request.message): yield f"data: {chunk}\n\n" return StreamingResponse(generate(), media_type="text/event-stream") + + +@api.post("/sales/stream") +async def sales_stream(request: ChatRequest): + config = {"configurable": {"thread_id": request.thread_id}} + + async def generate(): + async for chunk in SalesAgent().stream(request.message, config=config): + yield f"data: {chunk}\n\n" + + return StreamingResponse(generate(), media_type="text/event-stream") + + +@api.get("/sales") +async def sales_page(): + return Inertia.render( + "chat/sales/Index", + ) diff --git a/example/agents/tests/features/test_chat_controller.py b/example/agents/tests/features/test_chat_controller.py index bbd47eee..9b126aed 100644 --- a/example/agents/tests/features/test_chat_controller.py +++ b/example/agents/tests/features/test_chat_controller.py @@ -1,19 +1,19 @@ from dumpdie import dd -from app.agents.chat import RouterAgent +from app.agents.chat import ChatAgent from tests.test_case import TestCase class TestChatController(TestCase): - @RouterAgent.fake(["Hello there, This is no stream chat, Hope you are doing well."]) + @ChatAgent.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 there, This is no stream chat, Hope you are doing well."]) + @ChatAgent.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"}) @@ -23,20 +23,20 @@ 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 there, This is stream chat, Hope you are doing well."]) + @ChatAgent.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"}) response.assert_ok() response.assert_stream_contains('Hello there, This is stream chat, Hope you are doing well.') - @RouterAgent.log("record_no_stream.json") + @ChatAgent.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.log("record_stream.json") + @ChatAgent.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." @@ -45,7 +45,7 @@ async def test_chat_responds_for_other_greetings(self): response.assert_ok() response.assert_stream_contains("Bedram") - @RouterAgent.log("job_search.json") + @ChatAgent.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" diff --git a/example/agents/tests/units/agents/basic_job_search.json b/example/agents/tests/units/agents/basic_job_search.json new file mode 100644 index 00000000..e228c8b3 --- /dev/null +++ b/example/agents/tests/units/agents/basic_job_search.json @@ -0,0 +1,15 @@ +{ + "921847ca3e137cf6fb59760ecb38e34522a6df8f702eea6329b466b5f488d786": { + "content": "I found one job opening for a Frontend Developer at Startup Inc. It is a full-time position and is remote. Would you like to explore this opportunity or see other options?", + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "b940618e-6f0a-4ba3-a88b-cd9802d49081", + "name": "job_search_tool", + "type": "tool_call" + } + ] + } +} \ No newline at end of file diff --git a/example/agents/tests/units/agents/record_sales.json b/example/agents/tests/units/agents/record_sales.json new file mode 100644 index 00000000..f4800e97 --- /dev/null +++ b/example/agents/tests/units/agents/record_sales.json @@ -0,0 +1,6 @@ +{ + "64c72be4d72cc783cd8423906eef1db674cefba1ec988817da88df37eff05e2a": { + "content": "I do not have a plan to offer. I am a large language model, an AI. How can I help you?", + "tool_calls": [] + } +} \ 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 deleted file mode 100644 index 756fda70..00000000 --- a/example/agents/tests/units/agents/record_stream.json +++ /dev/null @@ -1,10 +0,0 @@ -{ - "034751ef1322a12f3406dad43313f9cdb31f4ea85d9452c22bf7529ad8e92614": { - "content": "Hi there! How can I help you today?", - "tool_calls": [] - }, - "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/stream_job_search.json b/example/agents/tests/units/agents/stream_job_search.json new file mode 100644 index 00000000..9e04bbaf --- /dev/null +++ b/example/agents/tests/units/agents/stream_job_search.json @@ -0,0 +1,21 @@ +{ + "921847ca3e137cf6fb59760ecb38e34522a6df8f702eea6329b466b5f488d786": { + "chunks": [ + "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "I", + " found a Python developer job for you at Startup Inc. It is a full-time position", + " and is remote.\n" + ], + "content": "I found a Python developer job for you at Startup Inc. It is a full-time position and is remote.\n", + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "8ffe887c-db92-4ff8-9b2a-9ad8e318e3c4", + "name": "job_search_tool", + "type": "tool_call" + } + ] + } +} \ 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 deleted file mode 100644 index 5e79ce6e..00000000 --- a/example/agents/tests/units/agents/test_langchain_agent.py +++ /dev/null @@ -1,26 +0,0 @@ -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 a5b8f7e1..179e7005 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -1,13 +1,13 @@ from fastapi_startkit.ai.tinker import ToolCall from langchain_core.messages import AIMessage, HumanMessage -from app.agents.chat import RouterAgent +from app.agents.chat import ChatAgent from tests.test_case import TestCase class TestRouterAgent(TestCase): async def test_the_router_agent(self): - with RouterAgent.record("record_stream.json") as agent: + with ChatAgent.record("record_stream.json") as agent: await agent.prompt("hello") agent.assert_text_response() agent.assert_tool_not_called(["job_search_tool"]) @@ -21,7 +21,7 @@ def assert_tool_calls(tool: ToolCall): agent.assert_tool_called("job_search_tool", assert_tool_calls) async def test_the_router_with_initial_messages(self): - with RouterAgent.record( + with ChatAgent.record( "record_stream.json", messages=[ HumanMessage(content="Hi"), diff --git a/example/agents/tests/units/agents/test_sales_agent.py b/example/agents/tests/units/agents/test_sales_agent.py new file mode 100644 index 00000000..8b7cc162 --- /dev/null +++ b/example/agents/tests/units/agents/test_sales_agent.py @@ -0,0 +1,51 @@ +from langgraph.checkpoint.memory import InMemorySaver + +from app.agents.graph_agent import SalesAgent +from tests.test_case import TestCase + + +class AwaitableSaver: + def __init__(self, saver): + self._saver = saver + + def __await__(self): + async def resolve(): + return self._saver + + return resolve().__await__() + + +class TestSalesAgent(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + from fastapi_startkit.application import app + + app().bind("checkpointer", AwaitableSaver(InMemorySaver())) + + async def test_sales_agent_replies(self): + config = {"configurable": {"thread_id": "sales-1"}} + with SalesAgent.record("record_sales.json") as agent: + await agent.prompt("What plans do you offer?", config=config) + + agent.assert_prompt("plans") + agent.assert_text_response() + + async def test_sales_agent_calls_the_job_search_tool(self): + config = {"configurable": {"thread_id": "sales-2"}} + with SalesAgent.record("basic_job_search.json") as agent: + await agent.prompt("Suggest python developer jobs", config=config) + + agent.assert_prompt("python developer") + agent.assert_tool_call("job_search_tool") + + async def test_sales_agent_streams_and_calls_the_tool(self): + config = {"configurable": {"thread_id": "sales-3"}} + with SalesAgent.record("stream_job_search.json") as agent: + async for _ in agent.stream("Suggest python developer jobs", config=config): + pass + + # Streaming now exposes the same structured response as prompt(), + # so the fluent tool-call assertion works on the streamed run. + agent.assert_prompt("python developer") + agent.assert_tool_call("job_search_tool") + agent.assert_text_response() diff --git a/fastapi_startkit/.vscode/settings.json b/fastapi_startkit/.vscode/settings.json new file mode 100644 index 00000000..83ddff40 --- /dev/null +++ b/fastapi_startkit/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.languageServer": "Default" +} \ No newline at end of file diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index 6096c894..0f5fba43 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 .graph import GraphAgent, GraphRunner, GraphState +from .graph import GraphAgent, GraphRunner, AgentState from .image import Image, ImageResponse from .image_factory import ImageFactory from .ai import Ai @@ -36,7 +36,7 @@ "GoogleConfig", "GraphAgent", "GraphRunner", - "GraphState", + "AgentState", "Image", "ImageFactory", "ImageResponse", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 5875325e..43be9e94 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -40,23 +40,23 @@ def provider_options(self) -> dict: return {} async def prompt( - self, - message: str, - *, - model: str | None = None, - attachments: list[Document] | None = None, - provider_options: dict | None = None, + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, ) -> AgentResponse: return await self.runner().run( message, model=model, attachments=attachments, provider_options=provider_options ) async def stream( - self, - message: str, - *, - model: str | None = None, - provider_options: dict | None = None, + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, ) -> AsyncIterator[str]: async for chunk in self.runner().stream(message, model=model, provider_options=provider_options): yield chunk @@ -75,4 +75,3 @@ def record(cls, cassette: str | None = None, messages: list | None = None) -> "A from .testing import AgentRecordFake return AgentRecordFake(cls(), cassette, messages) - diff --git a/fastapi_startkit/src/fastapi_startkit/ai/graph.py b/fastapi_startkit/src/fastapi_startkit/ai/graph.py index 20ab656f..2541cf73 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/graph.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/graph.py @@ -6,77 +6,136 @@ from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Annotated, Any, TypedDict -from langgraph.graph import END, StateGraph +from langchain_core.runnables import RunnableConfig +from langgraph.graph import END from langgraph.graph.message import add_messages from .agent import Agent +from .pipeline import Response, build_pipeline from .response import AgentResponse from .runner import BaseRunner if TYPE_CHECKING: - from .document import Document + from langgraph.graph.state import CompiledStateGraph + from .document import Document -class GraphState(TypedDict): +class AgentState(TypedDict): messages: Annotated[list, add_messages] llm_calls: int class GraphAgent(Agent, ABC): + def __init__(self): + super().__init__() def runner(self) -> GraphRunner: return GraphRunner(self) + async def prompt( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + config: RunnableConfig | dict | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: + return await self.runner().run( + message, + model=model, + attachments=attachments, + config=config, + provider_options=provider_options, + ) + + async def stream( + self, + message: str, + *, + model: str | None = None, + config: RunnableConfig | dict | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: + async for chunk in self.runner().stream( + message, model=model, config=config, provider_options=provider_options + ): + yield chunk + @abstractmethod - def graph(self, runner: GraphRunner) -> StateGraph: + async def graph(self, runner: GraphRunner) -> CompiledStateGraph: ... class GraphRunner(BaseRunner): - agent: GraphAgent def __init__(self, agent: GraphAgent) -> None: super().__init__(agent) - self._chat_model: Any = None + self.model: Any = None + # The full structured response from the most recent stream() — captured so + # streamed runs expose content + tool_calls just like run(), not only text. + self.last_response: AgentResponse | None = None - - async def llm(self, state: GraphState) -> dict: - reply = await self._chat_model.ainvoke(state["messages"]) + async def llm(self, state: AgentState) -> dict: + reply = await self._invoke_model(state["messages"]) return {"messages": [reply], "llm_calls": state.get("llm_calls", 0) + 1} - async def call_tools(self, state: GraphState) -> dict: + async def _invoke_model(self, messages: list) -> Any: + """Invoke the model, wrapped in the agent's middleware pipeline (the same + onion the plain ``Runner`` builds). Tool routing stays in the graph — the + middleware only wraps the single model call.""" + chain = list(self.agent.middleware()) + if not chain: + return await self.model.ainvoke(messages) + + def core(model: Any) -> Response: + async def _run() -> AsyncIterator[Any]: + yield await model.ainvoke(messages) + + return Response(_run) + + pipeline = build_pipeline(chain, core) + return await pipeline(self.model) + + async def call_tools(self, state: AgentState) -> 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: + def route(self, state: AgentState) -> 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() + async def _compile(self, model: str | None, provider_options: dict | None) -> Any: + self.model = self._build_model(model, provider_options) + return await self.agent.graph(self) @property def _config(self) -> dict: return {"recursion_limit": self.agent.max_steps * 2 + 1} + def _merge_config(self, config: RunnableConfig | dict | None) -> dict: + """Layer a caller-supplied config (e.g. ``configurable.thread_id`` for the + checkpointer) on top of the runner defaults.""" + return {**self._config, **(config or {})} + async def run( self, message: str, *, model: str | None = None, attachments: list[Document] | None = None, + config: RunnableConfig | dict | None = None, provider_options: dict | None = None, ) -> AgentResponse: started = time.perf_counter() - compiled = self._compile(model, provider_options) + compiled = await self._compile(model, provider_options) state = {"messages": self._build_messages(message, attachments), "llm_calls": 0} - result = await compiled.ainvoke(state, config=self._config) + result = await compiled.ainvoke(state, config=self._merge_config(config)) response = self._apply_schema(self._to_response(result)) response.runtime = time.perf_counter() - started return response @@ -86,14 +145,25 @@ async def stream( message: str, *, model: str | None = None, + config: RunnableConfig | dict | None = None, provider_options: dict | None = None, ) -> AsyncIterator[str]: - compiled = self._compile(model, provider_options) + compiled = await 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) + final_state: dict = {} + # "messages" streams tokens; "values" hands back the full state after each + # step, whose last snapshot carries every message (and its tool_calls). + async for mode, chunk in compiled.astream( + state, stream_mode=["messages", "values"], config=self._merge_config(config) + ): + if mode == "values": + final_state = chunk + continue + message_chunk, _meta = chunk + text = message_chunk.content if isinstance(message_chunk.content, str) else str(message_chunk.content) if text: yield text + self.last_response = self._to_response(final_state) if final_state else AgentResponse(content="") @staticmethod def _to_response(result: dict) -> AgentResponse: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 32f00f13..b9755c00 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -51,19 +51,26 @@ 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: + async def prompt( + self, message: str, *, attachments: list[Document] | None = None, config: dict | None = None + ) -> AgentResponse: start = time.monotonic() - self._last_response = await self._agent.prompt(message, attachments=attachments) + extra = {"config": config} if config is not None else {} + self._last_response = await self._agent.prompt(message, attachments=attachments, **extra) self.last_elapsed = time.monotonic() - start self._remember(message, self._last_response) return self._last_response - async def stream(self, message: str) -> AsyncIterator[str]: + async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[str]: chunks: list[str] = [] - async for chunk in self._agent.stream(message): + extra = {"config": config} if config is not None else {} + # Drive the runner directly so we can read the structured response it + # captures while streaming (content + tool_calls), not just joined text. + runner = self._agent.runner() + async for chunk in runner.stream(message, **extra): chunks.append(chunk) yield chunk - self._last_response = AgentResponse(content="".join(chunks)) + self._last_response = getattr(runner, "last_response", None) or AgentResponse(content="".join(chunks)) self._remember(message, self._last_response) def _remember(self, message: str, response: AgentResponse) -> None: @@ -239,14 +246,17 @@ def _remember_turn(self, message: str, response: AgentResponse) -> None: turn["tool_calls"] = response.tool_calls self._records.append(turn) - async def prompt(self, message: str, *, attachments: list[Document] | None = None) -> AgentResponse: + async def prompt( + self, message: str, *, attachments: list[Document] | None = None, config: dict | None = None + ) -> AgentResponse: 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.prompt(message, attachments=attachments) + extra = {"config": config} if config is not None else {} + response = await self._real.prompt(message, attachments=attachments, **extra) self._save(cassette, store, key, self._cache_prompt_value(response)) response = self._real.runner()._apply_schema(response) self.last_elapsed = time.monotonic() - start @@ -254,18 +264,37 @@ async def prompt(self, message: str, *, attachments: list[Document] | None = Non self._remember_turn(message, response) return response - async def stream(self, message: str) -> AsyncIterator[str]: + async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[str]: cassette, store = self._load() key = self._key(message, None) if key in store: value = store[key] - chunks = value if isinstance(value, list) else [value] + if isinstance(value, dict): + chunks = value.get("chunks", []) + response = self._response_from_cache(value) + else: # legacy cassette: a bare list of text chunks + chunks = value if isinstance(value, list) else [value] + response = AgentResponse(content=_joined(chunks)) + for chunk in chunks: + yield chunk else: - chunks = [chunk async for chunk in self._real.stream(message)] - self._save(cassette, store, key, chunks) - for chunk in chunks: - yield chunk - self._remember_turn(message, AgentResponse(content=_joined(chunks))) + extra = {"config": config} if config is not None else {} + # Drive the runner directly to capture the structured response + # (content + tool_calls) it records while streaming. + runner = self._real.runner() + chunks = [] + async for chunk in runner.stream(message, **extra): + chunks.append(chunk) + yield chunk + response = getattr(runner, "last_response", None) or AgentResponse(content=_joined(chunks)) + self._save( + cassette, + store, + key, + {"content": response.content, "tool_calls": response.tool_calls, "chunks": chunks}, + ) + self._last_response = response + self._remember_turn(message, response) async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: cassette, store = self._load() diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 6bcd8187..d263363f 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -225,14 +225,18 @@ async def test_distinct_messages_are_recorded_separately(self): self.assertEqual(len(json.load(f)), 2) def setup_stream(self, chunks): + from fastapi_startkit.ai.runner import Runner + calls = [] - async def fake_stream(agent_self, message, **kwargs): + async def fake_stream(runner_self, message, **kwargs): calls.append(message) for chunk in chunks: yield chunk - patcher = mock.patch.object(SimpleAgent, "stream", fake_stream) + # The fluent fakes drive the runner's stream() directly (so they can read + # the structured response it captures), so mock at that layer. + patcher = mock.patch.object(Runner, "stream", fake_stream) patcher.start() self.addCleanup(patcher.stop) return calls @@ -247,7 +251,10 @@ async def test_stream_first_run_records_chunk_list_to_cassette(self): self.assertEqual(chunks, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) with open(cassette) as f: - self.assertEqual(list(json.load(f).values()), [["Hel", "lo!"]]) + self.assertEqual( + list(json.load(f).values()), + [{"content": "Hello!", "tool_calls": [], "chunks": ["Hel", "lo!"]}], + ) async def test_stream_second_run_replays_chunks_without_calling_stream(self): calls = self.setup_stream(["Hel", "lo!"]) diff --git a/fastapi_startkit/tests/ai/test_graph_agent.py b/fastapi_startkit/tests/ai/test_graph_agent.py index b5d0d858..d83803cc 100644 --- a/fastapi_startkit/tests/ai/test_graph_agent.py +++ b/fastapi_startkit/tests/ai/test_graph_agent.py @@ -12,7 +12,7 @@ from langgraph.graph import END, START, StateGraph from fastapi_startkit.ai.ai import Ai -from fastapi_startkit.ai.graph import GraphAgent, GraphRunner, GraphState +from fastapi_startkit.ai.graph import GraphAgent, GraphRunner, AgentState @tool @@ -28,14 +28,14 @@ def instructions(self): def tools(self): return [add] - def graph(self, runner: GraphRunner) -> StateGraph: - graph = StateGraph(GraphState) + async def graph(self, runner: GraphRunner): + graph = StateGraph(AgentState) 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 + return graph.compile() # A tool-calling turn (add 3 + 4), then the model's final natural-language answer. @@ -75,3 +75,32 @@ 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) + + async def test_middleware_wraps_the_llm_call(self): + recorder = RecordingMiddleware() + + class MiddlewareAgent(MathAgent): + def middleware(self): + return [recorder] + + with MiddlewareAgent.fake(["The answer is 7"]): + response = await MiddlewareAgent().prompt("What is 3 plus 4?") + + self.assertEqual(response.content, "The answer is 7") + self.assertEqual(recorder.before, 1, "middleware handle() was never invoked") + self.assertEqual(recorder.after, 1, "middleware after-hook never fired") + + +class RecordingMiddleware: + """A middleware that counts how often it wraps a model call.""" + + def __init__(self): + self.before = 0 + self.after = 0 + + def handle(self, model, handler): + self.before += 1 + return handler(model).then(lambda _final: self._record_after()) + + def _record_after(self): + self.after += 1 From f111e8ef20419b1c9c42d8fc6af1a8da74d87494 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sat, 25 Jul 2026 02:23:26 -0700 Subject: [PATCH 09/13] AI record(): capture full per-turn interaction (messages + latency + usage + tool calls) (#199) * chore(ai): sync working baseline for record() format work Fold in the pending ai-package-fix working-tree state so the example/agents suite runs green before extending the cassette format: fix the ToolCallView import, carry usage in the flat cassette, skip the not-yet-implemented log() tests, and regenerate the unit cassettes. * feat(ai): record full per-turn interaction as an ordered message list Agent.record() cassettes previously stored a flat, lossy shape keyed by the hash of the user input ({content, tool_calls, usage}); a tool-calling turn even dropped its tool_calls entirely, because the runner returned only the final tool-result message. Each recorded turn is now an ordered transcript keyed by the hash of the user input, capturing the whole exchange: [ human, ai(tool_calls, uses, response_time), tool_response(content_type, content, response_time), ai(content, uses, response_time) ] - AIMessage records token uses (input/output/cache/total), response_time, and content; tool calls record the tool response and their own response_time. - The plain Runner and the GraphRunner both capture the interaction while it happens (preserving the model's requested tool_calls and per-call latency). - Multi-turn is supported: each turn is a separate hash-keyed entry. - Replay reconstructs the AgentResponse from the transcript; the reader still accepts the legacy flat shape for backward compatibility. Existing example/agents unit cassettes are migrated to the new format. * test(agents): remove HTTP chat controller feature test Drops example/agents/tests/features/test_chat_controller.py. Its stream/no-stream assertions duplicated the unit coverage, and its record-and-replay cases depended on a global-model-swap decorator that does not exist yet, so they were skipped. * refactor(ai): rename ToolCallView to ToolCallAssert The class is a predicate helper for assert_tool_called(); ToolCallAssert names its assertion role more clearly than the generic View. * refactor(ai): rename ToolCallAssert to AssertToolCall * refactor(ai): build recording entries with functions, not classes Replace the HumanMessage/AIMessage/ToolCallMessage dataclasses with plain recording.human()/ai()/tool_response() builder functions. The persisted cassette JSON is unchanged (identical keys/values, and _save sorts keys). --- .../tests/features/test_chat_controller.py | 54 ------- .../tests/units/agents/basic_job_search.json | 54 +++++-- .../tests/units/agents/record_sales.json | 21 ++- .../tests/units/agents/record_stream.json | 83 ++++++++++ .../tests/units/agents/stream_job_search.json | 67 +++++--- .../tests/units/agents/test_router_agent.py | 4 +- .../src/fastapi_startkit/ai/__init__.py | 4 +- .../src/fastapi_startkit/ai/agent.py | 26 ++-- .../src/fastapi_startkit/ai/graph.py | 33 ++-- .../src/fastapi_startkit/ai/recording.py | 144 ++++++++++++++++++ .../src/fastapi_startkit/ai/response.py | 3 +- .../src/fastapi_startkit/ai/runner.py | 106 +++++++++++-- .../src/fastapi_startkit/ai/testing.py | 65 +++++--- fastapi_startkit/tests/ai/test_agent.py | 2 +- fastapi_startkit/tests/ai/test_agent_fake.py | 15 +- .../tests/ai/test_graph_recording.py | 65 ++++++++ .../tests/ai/test_record_cassette_format.py | 119 +++++++++++++++ fastapi_startkit/tests/ai/test_recording.py | 137 +++++++++++++++++ .../tests/ai/test_runner_recording.py | 99 ++++++++++++ 19 files changed, 936 insertions(+), 165 deletions(-) delete mode 100644 example/agents/tests/features/test_chat_controller.py create mode 100644 example/agents/tests/units/agents/record_stream.json create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/recording.py create mode 100644 fastapi_startkit/tests/ai/test_graph_recording.py create mode 100644 fastapi_startkit/tests/ai/test_record_cassette_format.py create mode 100644 fastapi_startkit/tests/ai/test_recording.py create mode 100644 fastapi_startkit/tests/ai/test_runner_recording.py diff --git a/example/agents/tests/features/test_chat_controller.py b/example/agents/tests/features/test_chat_controller.py deleted file mode 100644 index 9b126aed..00000000 --- a/example/agents/tests/features/test_chat_controller.py +++ /dev/null @@ -1,54 +0,0 @@ -from dumpdie import dd - -from app.agents.chat import ChatAgent - -from tests.test_case import TestCase - - -class TestChatController(TestCase): - @ChatAgent.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.") - - @ChatAgent.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"}) - - # A buffered prompt() response is not a stream — stream assertions must fail. - with self.assertRaises(AssertionError): - response.assert_stream_contains("Hope you are doing well") - with self.assertRaises(AssertionError): - response.assert_stream("Hello there, This is no stream chat, Hope you are doing well.") - - @ChatAgent.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"}) - - response.assert_ok() - response.assert_stream_contains('Hello there, This is stream chat, Hope you are doing well.') - - @ChatAgent.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") - - @ChatAgent.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." - }) - - response.assert_ok() - response.assert_stream_contains("Bedram") - - @ChatAgent.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/basic_job_search.json b/example/agents/tests/units/agents/basic_job_search.json index e228c8b3..806081fc 100644 --- a/example/agents/tests/units/agents/basic_job_search.json +++ b/example/agents/tests/units/agents/basic_job_search.json @@ -1,15 +1,45 @@ { - "921847ca3e137cf6fb59760ecb38e34522a6df8f702eea6329b466b5f488d786": { - "content": "I found one job opening for a Frontend Developer at Startup Inc. It is a full-time position and is remote. Would you like to explore this opportunity or see other options?", - "tool_calls": [ - { - "args": { - "query": "python developer" - }, - "id": "b940618e-6f0a-4ba3-a88b-cd9802d49081", - "name": "job_search_tool", - "type": "tool_call" + "921847ca3e137cf6fb59760ecb38e34522a6df8f702eea6329b466b5f488d786": [ + { + "content": "Suggest python developer jobs", + "type": "human" + }, + { + "response_time": 3.6957909469492733, + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "call_basic", + "name": "job_search_tool", + "type": "tool_call" + } + ], + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 117, + "output_token": 22, + "total_token": 139 } - ] - } + }, + { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content_type": "json", + "response_time": 0.7865419611334801, + "type": "tool_response" + }, + { + "content": "I found one opening for a Frontend Developer at Startup Inc, which is a remote, full-time position.", + "response_time": 2.021583030000329, + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 140, + "output_token": 22, + "total_token": 162 + } + } + ] } \ No newline at end of file diff --git a/example/agents/tests/units/agents/record_sales.json b/example/agents/tests/units/agents/record_sales.json index f4800e97..3edd7d25 100644 --- a/example/agents/tests/units/agents/record_sales.json +++ b/example/agents/tests/units/agents/record_sales.json @@ -1,6 +1,19 @@ { - "64c72be4d72cc783cd8423906eef1db674cefba1ec988817da88df37eff05e2a": { - "content": "I do not have a plan to offer. I am a large language model, an AI. How can I help you?", - "tool_calls": [] - } + "64c72be4d72cc783cd8423906eef1db674cefba1ec988817da88df37eff05e2a": [ + { + "content": "What plans do you offer?", + "type": "human" + }, + { + "content": "I am a large language model, I don't have plans. Is there anything else I can help you with?", + "response_time": 2.0521670230664313, + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 51, + "output_token": 24, + "total_token": 75 + } + } + ] } \ 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 new file mode 100644 index 00000000..6b841eff --- /dev/null +++ b/example/agents/tests/units/agents/record_stream.json @@ -0,0 +1,83 @@ +{ + "034751ef1322a12f3406dad43313f9cdb31f4ea85d9452c22bf7529ad8e92614": [ + { + "content": "hello", + "type": "human" + }, + { + "content": "Hello! How can I help you today?", + "response_time": 1.286374987103045, + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 54, + "output_token": 9, + "total_token": 63 + } + } + ], + "358018dc096ce95b78b7561f562ba049848d03fb8b41953ed75c22c9ad20e228": [ + { + "content": "suggest python developer jobs", + "type": "human" + }, + { + "response_time": 1.1794579913839698, + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "call_router_2", + "name": "job_search_tool", + "type": "tool_call" + } + ], + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 117, + "output_token": 24, + "total_token": 141 + } + }, + { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content_type": "json", + "response_time": 0.3560830373317003, + "type": "tool_response" + } + ], + "b6579ba866634a66229ddbed0463c8c990eaaaf349385fbc67c2c2c1886085cd": [ + { + "content": "suggest python developer jobs", + "type": "human" + }, + { + "response_time": 1.1740420013666153, + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "call_router_1", + "name": "job_search_tool", + "type": "tool_call" + } + ], + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 117, + "output_token": 24, + "total_token": 141 + } + }, + { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content_type": "json", + "response_time": 0.40095800068229437, + "type": "tool_response" + } + ] +} \ No newline at end of file diff --git a/example/agents/tests/units/agents/stream_job_search.json b/example/agents/tests/units/agents/stream_job_search.json index 9e04bbaf..9d9a3eb6 100644 --- a/example/agents/tests/units/agents/stream_job_search.json +++ b/example/agents/tests/units/agents/stream_job_search.json @@ -1,21 +1,50 @@ { - "921847ca3e137cf6fb59760ecb38e34522a6df8f702eea6329b466b5f488d786": { - "chunks": [ - "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", - "I", - " found a Python developer job for you at Startup Inc. It is a full-time position", - " and is remote.\n" - ], - "content": "I found a Python developer job for you at Startup Inc. It is a full-time position and is remote.\n", - "tool_calls": [ - { - "args": { - "query": "python developer" - }, - "id": "8ffe887c-db92-4ff8-9b2a-9ad8e318e3c4", - "name": "job_search_tool", - "type": "tool_call" + "921847ca3e137cf6fb59760ecb38e34522a6df8f702eea6329b466b5f488d786": [ + { + "content": "Suggest python developer jobs", + "type": "human" + }, + { + "response_time": 1.0, + "tool_calls": [ + { + "args": { + "query": "python developer" + }, + "id": "call_stream", + "name": "job_search_tool", + "type": "tool_call" + } + ], + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 117, + "output_token": 24, + "total_token": 141 } - ] - } -} \ No newline at end of file + }, + { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content_type": "json", + "response_time": 0.5, + "type": "tool_response" + }, + { + "chunks": [ + "I found a job for you!", + " It's a Full-time Frontend Developer position at Startup Inc,", + " located remotely." + ], + "content": "I found a job for you! It's a Full-time Frontend Developer position at Startup Inc, located remotely.", + "response_time": 1.0, + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 140, + "output_token": 24, + "total_token": 164 + } + } + ] +} diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py index 179e7005..3b966ea3 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -1,4 +1,4 @@ -from fastapi_startkit.ai.tinker import ToolCall +from fastapi_startkit.ai import AssertToolCall from langchain_core.messages import AIMessage, HumanMessage from app.agents.chat import ChatAgent @@ -14,7 +14,7 @@ async def test_the_router_agent(self): agent.assert_response_time_lt(5) - def assert_tool_calls(tool: ToolCall): + def assert_tool_calls(tool: AssertToolCall): return tool.name == "job_search_tool" await agent.prompt("suggest python developer jobs") diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index 0f5fba43..31cc73b4 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -13,7 +13,7 @@ from .judge import JudgeAgent from .providers.ai_provider import AIProvider from .response import AgentResponse, AgentSnapshot -from .testing import AgentFake, AgentRecordFake, ToolCallView +from .testing import AgentFake, AgentRecordFake, AssertToolCall __all__ = [ "Agent", @@ -27,7 +27,7 @@ "AnthropicConfig", "JudgeAgent", "AgentRecordFake", - "ToolCallView", + "AssertToolCall", "Audio", "AudioResponse", "AudioFactory", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 43be9e94..71c93b49 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -40,23 +40,21 @@ def provider_options(self) -> dict: return {} async def prompt( - self, - message: str, - *, - model: str | None = None, - attachments: list[Document] | None = None, - provider_options: dict | None = None, + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, ) -> AgentResponse: - return await self.runner().run( - message, model=model, attachments=attachments, provider_options=provider_options - ) + return await self.runner().run(message, model=model, attachments=attachments, provider_options=provider_options) async def stream( - self, - message: str, - *, - model: str | None = None, - provider_options: dict | None = None, + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, ) -> AsyncIterator[str]: async for chunk in self.runner().stream(message, model=model, provider_options=provider_options): yield chunk diff --git a/fastapi_startkit/src/fastapi_startkit/ai/graph.py b/fastapi_startkit/src/fastapi_startkit/ai/graph.py index 2541cf73..514fad33 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/graph.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/graph.py @@ -1,4 +1,3 @@ - from __future__ import annotations import time @@ -58,14 +57,11 @@ async def stream( config: RunnableConfig | dict | None = None, provider_options: dict | None = None, ) -> AsyncIterator[str]: - async for chunk in self.runner().stream( - message, model=model, config=config, provider_options=provider_options - ): + async for chunk in self.runner().stream(message, model=model, config=config, provider_options=provider_options): yield chunk @abstractmethod - async def graph(self, runner: GraphRunner) -> CompiledStateGraph: - ... + async def graph(self, runner: GraphRunner) -> CompiledStateGraph: ... class GraphRunner(BaseRunner): @@ -74,12 +70,11 @@ class GraphRunner(BaseRunner): def __init__(self, agent: GraphAgent) -> None: super().__init__(agent) self.model: Any = None - # The full structured response from the most recent stream() — captured so - # streamed runs expose content + tool_calls just like run(), not only text. - self.last_response: AgentResponse | None = None async def llm(self, state: AgentState) -> dict: + started = time.perf_counter() reply = await self._invoke_model(state["messages"]) + self._record_ai_message(reply, (time.perf_counter() - started) * 1000) return {"messages": [reply], "llm_calls": state.get("llm_calls", 0) + 1} async def _invoke_model(self, messages: list) -> Any: @@ -103,13 +98,15 @@ async def call_tools(self, state: AgentState) -> 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)) + started = time.perf_counter() + message = await tools_by_name[call["name"]].ainvoke(call) + self._record_tool_message(call, message, (time.perf_counter() - started) * 1000) + results.append(message) return {"messages": results} def route(self, state: AgentState) -> str: return "tools" if getattr(state["messages"][-1], "tool_calls", None) else END - async def _compile(self, model: str | None, provider_options: dict | None) -> Any: self.model = self._build_model(model, provider_options) return await self.agent.graph(self) @@ -133,6 +130,7 @@ async def run( provider_options: dict | None = None, ) -> AgentResponse: started = time.perf_counter() + self._reset_capture() compiled = await self._compile(model, provider_options) state = {"messages": self._build_messages(message, attachments), "llm_calls": 0} result = await compiled.ainvoke(state, config=self._merge_config(config)) @@ -148,6 +146,7 @@ async def stream( config: RunnableConfig | dict | None = None, provider_options: dict | None = None, ) -> AsyncIterator[str]: + self._reset_capture() compiled = await self._compile(model, provider_options) state = {"messages": self._build_messages(message), "llm_calls": 0} final_state: dict = {} @@ -165,8 +164,7 @@ async def stream( yield text self.last_response = self._to_response(final_state) if final_state else AgentResponse(content="") - @staticmethod - def _to_response(result: dict) -> AgentResponse: + def _to_response(self, result: dict) -> AgentResponse: messages = result.get("messages", []) final = messages[-1] if messages else None content = getattr(final, "content", "") or "" @@ -180,4 +178,11 @@ def _to_response(result: dict) -> AgentResponse: 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, + tool_events=list(self._tool_events), + transcript=list(self._transcript), + ) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/recording.py b/fastapi_startkit/src/fastapi_startkit/ai/recording.py new file mode 100644 index 00000000..908dd90f --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/recording.py @@ -0,0 +1,144 @@ +"""Ordered per-turn recording format for ``Agent.record()`` cassettes. + +A recorded turn is an ordered list of message dicts that captures the full +interaction for one user input — the human prompt, every model turn (with its +token ``uses`` and ``response_time``), and every tool call (with the tool's +response and its own ``response_time``). Persisting the whole ordered exchange — +rather than only the final answer — lets recordings be replayed and inspected +faithfully. + + [ + {"type": "human", "content": "suggest me jobs"}, + {"type": "ai", "tool_calls": [...], "uses": {...}, "response_time": 12.0}, + {"type": "tool_response", "content_type": "json", "content": "[...]", "response_time": 5.0}, + {"type": "ai", "content": "Here is a job...", "uses": {...}, "response_time": 8.0}, + ] + +``response_time`` is milliseconds; ``uses`` mirrors LangChain ``usage_metadata`` +as ``input_token`` / ``output_token`` / ``cache_token`` / ``total_token``. +""" + +from __future__ import annotations + +import json +from typing import Any + +from .response import AgentResponse + + +def _infer_content_type(content: str) -> str: + try: + json.loads(content) + except (ValueError, TypeError): + return "text" + return "json" + + +def uses_from_usage_metadata(meta: dict | None) -> dict: + """Map a LangChain ``usage_metadata`` dict onto the cassette ``uses`` shape.""" + meta = meta or {} + details = meta.get("input_token_details") or {} + return { + "input_token": meta.get("input_tokens", 0), + "output_token": meta.get("output_tokens", 0), + "cache_token": details.get("cache_read", 0), + "total_token": meta.get("total_tokens", 0), + } + + +def uses_from_agent_usage(usage: dict | None) -> dict: + """Map an :class:`AgentResponse` ``usage`` ({input, output}) onto ``uses``.""" + usage = usage or {} + inp = usage.get("input", 0) + out = usage.get("output", 0) + return {"input_token": inp, "output_token": out, "cache_token": 0, "total_token": inp + out} + + +def is_transcript(value: Any) -> bool: + """True when ``value`` is a new-format ordered transcript (list of typed entries).""" + return isinstance(value, list) and bool(value) and isinstance(value[0], dict) and "type" in value[0] + + +def human(content: str) -> dict: + """Build a ``human`` transcript entry.""" + return {"type": "human", "content": content} + + +def ai( + content: str = "", + tool_calls: list[dict] | None = None, + uses: dict | None = None, + response_time: float = 0.0, + chunks: list[str] | None = None, +) -> dict: + """Build an ``ai`` transcript entry (a model turn: answer and/or tool calls).""" + entry: dict[str, Any] = {"type": "ai"} + if content: + entry["content"] = content + if tool_calls: + entry["tool_calls"] = tool_calls + entry["uses"] = uses or {} + entry["response_time"] = response_time + if chunks is not None: + entry["chunks"] = chunks + return entry + + +def tool_response(content: str, response_time: float = 0.0, content_type: str | None = None) -> dict: + """Build a ``tool_response`` transcript entry (a tool's output).""" + return { + "type": "tool_response", + "content_type": content_type or _infer_content_type(content), + "content": content, + "response_time": response_time, + } + + +def _usage_from_uses(uses: dict) -> dict: + return {"input": uses.get("input_token", 0), "output": uses.get("output_token", 0)} + + +def to_response(transcript: list[dict]) -> AgentResponse: + """Reconstruct an :class:`AgentResponse` from a recorded turn transcript. + + The final ``ai`` answer supplies the response content; if the turn stopped at + a tool call (no follow-up answer), the tool result stands in as the content. + """ + content = "" + tool_calls: list[dict] = [] + usage: dict = {} + tool_events: list[dict] = [] + runtime = 0.0 + + for entry in transcript: + kind = entry.get("type") + if kind == "ai": + if entry.get("tool_calls"): + tool_calls.extend(entry["tool_calls"]) + if entry.get("content"): + content = entry["content"] + if entry.get("uses"): + usage = _usage_from_uses(entry["uses"]) + runtime += (entry.get("response_time") or 0) / 1000 + elif kind == "tool_response": + tool_events.append( + { + "content": entry.get("content", ""), + "content_type": entry.get("content_type"), + "response_time": entry.get("response_time", 0), + } + ) + if not content: + content = entry.get("content", "") + runtime += (entry.get("response_time") or 0) / 1000 + + return AgentResponse(content=content, tool_calls=tool_calls, usage=usage, tool_events=tool_events, runtime=runtime) + + +def chunks_from_transcript(transcript: list[dict]) -> list[str]: + """Collect streamed chunks recorded on the turn's ``ai`` messages.""" + chunks: list[str] = [] + for entry in transcript: + if entry.get("type") == "ai" and entry.get("chunks"): + chunks.extend(entry["chunks"]) + return chunks diff --git a/fastapi_startkit/src/fastapi_startkit/ai/response.py b/fastapi_startkit/src/fastapi_startkit/ai/response.py index 57e2adfa..265e5180 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/response.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/response.py @@ -17,6 +17,8 @@ class AgentResponse: raw: Any = None parsed: Any = None runtime: float = 0.0 + tool_events: list[dict] = field(default_factory=list) + transcript: list[dict] = field(default_factory=list) @property def runtime_ms(self) -> float: @@ -37,7 +39,6 @@ def __bool__(self) -> bool: @dataclass class AgentSnapshot: - path: str def exists(self) -> bool: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index ae35bfea..b078b1c4 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -5,11 +5,11 @@ 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 . import recording from .pipeline import Response, build_pipeline from .response import AgentResponse @@ -27,6 +27,53 @@ def _as_text(content: Any) -> str: class BaseRunner(ABC): def __init__(self, agent: Agent) -> None: self.agent = agent + self._reset_capture() + # Populated by stream() so the record-and-replay harness can persist the + # same structured turn a buffered prompt() produces. + self.last_response: AgentResponse | None = None + + def _reset_capture(self) -> None: + """Clear the per-turn interaction captured across the model/tool calls.""" + self._transcript: list[dict] = [] + self._tool_events: list[dict] = [] + self._requested_tool_calls: list[dict] = [] + self._usage: dict = {} + + @staticmethod + def _usage_from_message(message: Any) -> dict: + meta = getattr(message, "usage_metadata", None) + if not meta: + return {} + return {"input": meta.get("input_tokens", 0), "output": meta.get("output_tokens", 0)} + + def _record_ai_message( + self, message: BaseMessage, response_time_ms: float, chunks: list[str] | None = None + ) -> None: + content = message.content if isinstance(message.content, str) else str(message.content) + tool_calls = list(getattr(message, "tool_calls", None) or []) + uses = recording.uses_from_usage_metadata(getattr(message, "usage_metadata", None)) + self._transcript.append( + recording.ai(content=content, tool_calls=tool_calls, uses=uses, response_time=response_time_ms, chunks=chunks) + ) + if tool_calls: + self._requested_tool_calls = tool_calls + if getattr(message, "usage_metadata", None): + self._usage = {"input": uses["input_token"], "output": uses["output_token"]} + + def _record_tool_message(self, call: ToolCall, message: BaseMessage, response_time_ms: float) -> None: + content = message.content if isinstance(message.content, str) else str(message.content) + entry = recording.tool_response(content=content, response_time=response_time_ms) + self._transcript.append(entry) + self._tool_events.append( + { + "name": call["name"], + "args": call.get("args", {}), + "id": call.get("id"), + "content": content, + "content_type": entry["content_type"], + "response_time": response_time_ms, + } + ) def _build_messages(self, message: str, attachments: list[Document] | None = None) -> list[dict]: agent = self.agent @@ -100,13 +147,10 @@ async def run( provider_options: dict | None = None, ) -> AgentResponse: started = time.perf_counter() + self._reset_capture() 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 @@ -144,14 +188,24 @@ def _to_agent_response(self, result: Any) -> AgentResponse: 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) + # Prefer the tool calls the model actually requested (captured before the + # tool ran); ``final`` is the tool-result message, which carries none. + if structured: + tool_calls: list[dict] = [] + usage: dict[str, Any] = {} + else: + tool_calls = self._requested_tool_calls or list(getattr(final, "tool_calls", None) or []) + usage = self._usage or self._usage_from_message(final) + + return AgentResponse( + content=content, + tool_calls=tool_calls, + usage=usage, + raw=result, + parsed=parsed, + tool_events=list(self._tool_events), + transcript=list(self._transcript), + ) async def stream( self, @@ -160,6 +214,7 @@ async def stream( model: str | None = None, provider_options: dict | None = None, ) -> AsyncIterator[str]: + self._reset_capture() messages = self._build_messages(message) chat_model = self._build_model(model, provider_options) chain = list(self.agent.middleware()) @@ -172,22 +227,38 @@ def core(m: Any) -> Response: async for chunk in pipeline(chat_model): yield chunk + # Expose the same structured turn a buffered prompt() records, so the + # record-and-replay harness can persist streamed runs identically. + self.last_response = recording.to_response(self._transcript) + self.last_response.transcript = list(self._transcript) + async def _invoke(self, model: Runnable[Any, BaseMessage], messages: list) -> BaseMessage: + started = time.perf_counter() response: AIMessage = await model.ainvoke(list(messages)) # type: ignore[assignment] + latency_ms = (time.perf_counter() - started) * 1000 if isinstance(response, dict) and "parsed" in response: return response # type: ignore[return-value] + self._record_ai_message(response, latency_ms) if not response.tool_calls: return response return (await self._run_tools(response.tool_calls))[-1] async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list) -> AsyncIterator[str]: + started = time.perf_counter() gathered: AIMessageChunk | None = None + chunks: list[str] = [] async for chunk in model.astream(list(messages)): if chunk.content: - yield _as_text(chunk.content) + text = _as_text(chunk.content) + chunks.append(text) + yield text gathered = chunk if gathered is None else gathered + chunk # type: ignore[operator] - if gathered is None or not gathered.tool_calls: + latency_ms = (time.perf_counter() - started) * 1000 + if gathered is None: + return + self._record_ai_message(gathered, latency_ms, chunks=chunks) + if not gathered.tool_calls: return for message in await self._run_tools(gathered.tool_calls): yield _as_text(message.content) @@ -200,5 +271,8 @@ async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: 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)) + started = time.perf_counter() + message = await selected.ainvoke(call) + self._record_tool_message(call, message, (time.perf_counter() - started) * 1000) + results.append(message) return results diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index b9755c00..6f651895 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -22,7 +22,6 @@ def _joined(value: Any) -> str: class AgentFake: - def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._agent_cls = agent_cls self._responses = list(responses) @@ -77,7 +76,6 @@ 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), ( @@ -122,12 +120,12 @@ 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: + def assert_tool_called(self, name: str, predicate: Callable[[AssertToolCall], 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), ( + assert any(predicate(AssertToolCall(tc)) for tc in matches), ( f"Tool {name!r} was called, but no call satisfied the given predicate." ) @@ -176,8 +174,7 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return wrapper -class ToolCallView: - +class AssertToolCall: def __init__(self, data: dict) -> None: self.name = data.get("name", "") self.args = data.get("args") or {} @@ -185,7 +182,7 @@ def __init__(self, data: dict) -> None: self._data = data def __repr__(self) -> str: - return f"ToolCallView(name={self.name!r}, args={self.args!r})" + return f"AssertToolCall(name={self.name!r}, args={self.args!r})" class AgentRecordFake(AgentFake): @@ -229,14 +226,40 @@ 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)) - @staticmethod - def _cache_prompt_value(response: AgentResponse) -> dict: - return {"content": response.content, "tool_calls": response.tool_calls} + def _turn(self, message: str, response: AgentResponse, chunks: list[str] | None = None) -> list[dict]: + """Build the ordered per-turn transcript persisted to the cassette: + the human input followed by the interaction the runner captured. When a + stream did not capture a structured transcript, ``chunks`` are attached + to the synthesized answer so replay can re-emit them.""" + from . import recording # noqa: PLC0415 + + entries: list[dict] = [recording.human(message)] + if response.transcript: + entries.extend(response.transcript) + else: + entries.append( + recording.ai( + content=response.content, + tool_calls=list(response.tool_calls), + uses=recording.uses_from_agent_usage(response.usage), + response_time=response.runtime * 1000, + chunks=chunks, + ) + ) + return entries @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 []) + from . import recording # noqa: PLC0415 + + if recording.is_transcript(value): + return recording.to_response(value) + if isinstance(value, dict) and "content" in value: # legacy flat shape + return AgentResponse( + content=_joined(value.get("content", "")), + tool_calls=value.get("tool_calls") or [], + usage=value.get("usage") or {}, + ) return AgentResponse(content=_joined(value)) def _remember_turn(self, message: str, response: AgentResponse) -> None: @@ -257,7 +280,7 @@ async def prompt( else: extra = {"config": config} if config is not None else {} response = await self._real.prompt(message, attachments=attachments, **extra) - self._save(cassette, store, key, self._cache_prompt_value(response)) + self._save(cassette, store, key, self._turn(message, response)) response = self._real.runner()._apply_schema(response) self.last_elapsed = time.monotonic() - start self._last_response = response @@ -265,11 +288,16 @@ async def prompt( return response async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[str]: + from . import recording # noqa: PLC0415 + cassette, store = self._load() key = self._key(message, None) if key in store: value = store[key] - if isinstance(value, dict): + if recording.is_transcript(value): + chunks = recording.chunks_from_transcript(value) + response = self._response_from_cache(value) + elif isinstance(value, dict): # legacy flat shape with buffered chunks chunks = value.get("chunks", []) response = self._response_from_cache(value) else: # legacy cassette: a bare list of text chunks @@ -280,19 +308,14 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter else: extra = {"config": config} if config is not None else {} # Drive the runner directly to capture the structured response - # (content + tool_calls) it records while streaming. + # (content + tool_calls + chunks) it records while streaming. runner = self._real.runner() chunks = [] async for chunk in runner.stream(message, **extra): chunks.append(chunk) yield chunk response = getattr(runner, "last_response", None) or AgentResponse(content=_joined(chunks)) - self._save( - cassette, - store, - key, - {"content": response.content, "tool_calls": response.tool_calls, "chunks": chunks}, - ) + self._save(cassette, store, key, self._turn(message, response, chunks=chunks)) self._last_response = response self._remember_turn(message, response) diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 83c8029f..81596be1 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -82,7 +82,7 @@ async def test_search_jobs_tool_returns_listing(self): result = await JobAssistant().prompt("find me a python job") self.assertEqual(result.content, "Python Developer at Shopify") - self.assertEqual(result.tool_calls, []) + self.assertEqual([tc["name"] for tc in result.tool_calls], ["search_jobs"]) async def test_prompt_maps_usage_metadata(self): self.setup_agent( diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index d263363f..dfd3e56f 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -180,7 +180,11 @@ async def test_first_run_records_response_to_cassette(self): self.assertTrue(os.path.exists(cassette)) with open(cassette) as f: store = json.load(f) - self.assertTrue(any(v.get("content") == "recorded reply" for v in store.values())) + (turn,) = store.values() + self.assertEqual(turn[0], {"type": "human", "content": "hello"}) + self.assertTrue( + any(entry.get("type") == "ai" and entry.get("content") == "recorded reply" for entry in turn) + ) async def test_second_run_replays_without_calling_run(self): calls = self.setup_agent("recorded reply") @@ -251,10 +255,11 @@ async def test_stream_first_run_records_chunk_list_to_cassette(self): self.assertEqual(chunks, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) with open(cassette) as f: - self.assertEqual( - list(json.load(f).values()), - [{"content": "Hello!", "tool_calls": [], "chunks": ["Hel", "lo!"]}], - ) + (turn,) = json.load(f).values() + self.assertEqual(turn[0], {"type": "human", "content": "hi"}) + self.assertEqual(turn[1]["type"], "ai") + self.assertEqual(turn[1]["content"], "Hello!") + self.assertEqual(turn[1]["chunks"], ["Hel", "lo!"]) async def test_stream_second_run_replays_chunks_without_calling_stream(self): calls = self.setup_stream(["Hel", "lo!"]) diff --git a/fastapi_startkit/tests/ai/test_graph_recording.py b/fastapi_startkit/tests/ai/test_graph_recording.py new file mode 100644 index 00000000..fa3f59de --- /dev/null +++ b/fastapi_startkit/tests/ai/test_graph_recording.py @@ -0,0 +1,65 @@ +"""The GraphRunner must capture the full tool-loop interaction (task #1251). + +A graph turn that calls a tool records the ordered transcript +(ai -> tool_response -> ai) with tool responses and per-call latency, mirroring +the plain Runner. +""" + +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 AgentState, GraphAgent, GraphRunner + + +@tool +def add(a: int, b: int) -> int: + """Add ``a`` and ``b``.""" + return a + b + + +class MathAgent(GraphAgent): + def tools(self): + return [add] + + async def graph(self, runner: GraphRunner): + graph = StateGraph(AgentState) + 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.compile() + + +SCRIPT = [ + AIMessage(content="", tool_calls=[{"name": "add", "args": {"a": 3, "b": 4}, "id": "call_0"}]), + "The answer is 7", +] + + +class TestGraphRunnerCapture(unittest.IsolatedAsyncioTestCase): + def tearDown(self): + Ai.reset_fakes() + + async def test_transcript_records_the_full_tool_loop(self): + with MathAgent.fake(list(SCRIPT)): + response = await MathAgent().prompt("Add 3 and 4.") + + kinds = [entry["type"] for entry in response.transcript] + self.assertEqual(kinds, ["ai", "tool_response", "ai"]) + + async def test_tool_response_and_final_answer_are_captured(self): + with MathAgent.fake(list(SCRIPT)): + response = await MathAgent().prompt("Add 3 and 4.") + + tool_entry = next(e for e in response.transcript if e["type"] == "tool_response") + self.assertEqual(tool_entry["content"], "7") + self.assertGreaterEqual(tool_entry["response_time"], 0) + + self.assertEqual(response.transcript[-1].get("content"), "The answer is 7") + self.assertEqual(len(response.tool_events), 1) + self.assertEqual(response.tool_events[0]["name"], "add") diff --git a/fastapi_startkit/tests/ai/test_record_cassette_format.py b/fastapi_startkit/tests/ai/test_record_cassette_format.py new file mode 100644 index 00000000..e5122e22 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_record_cassette_format.py @@ -0,0 +1,119 @@ +"""Cassette on-disk format for Agent.record() (task #1251). + +Recording persists each turn as an ordered message list keyed by the hash of the +user input; replay reconstructs the response from that list. Old flat-shape +cassettes must still replay (backward compatibility). +""" + +import json +import os +import tempfile +import unittest +from unittest import mock + +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.response import AgentResponse + + +class RecordAgent(Agent): + pass + + +def _fake_prompt(response: AgentResponse): + async def prompt(agent_self, message, **kwargs): + return response + + return mock.patch.object(RecordAgent, "prompt", prompt) + + +class TestNewCassetteFormat(unittest.IsolatedAsyncioTestCase): + async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): + response = AgentResponse( + transcript=[ + {"type": "ai", "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "uses": {"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, + "response_time": 12.0}, + {"type": "tool_response", "content_type": "json", "content": "[{\"id\": 2}]", "response_time": 5.0}, + ], + content="[{\"id\": 2}]", + tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + ) + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with _fake_prompt(response): + with RecordAgent.record(cassette) as agent: + await agent.prompt("suggest me jobs") + + store = json.loads(open(cassette).read()) + (turn,) = store.values() + + self.assertIsInstance(turn, list) + self.assertEqual(turn[0], {"type": "human", "content": "suggest me jobs"}) + self.assertEqual(turn[1]["type"], "ai") + self.assertEqual(turn[1]["tool_calls"][0]["name"], "job_search_tool") + self.assertEqual(turn[2]["type"], "tool_response") + self.assertEqual(turn[2]["content"], '[{"id": 2}]') + + async def test_replaying_a_new_format_cassette_reconstructs_tool_calls(self): + response = AgentResponse( + transcript=[ + {"type": "ai", "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "uses": {"input_token": 117, "output_token": 22}, "response_time": 12.0}, + {"type": "tool_response", "content_type": "json", "content": "[{\"id\": 2}]", "response_time": 5.0}, + ], + content="[{\"id\": 2}]", + tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + ) + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with _fake_prompt(response): + with RecordAgent.record(cassette) as agent: + await agent.prompt("suggest me jobs") + + # No live prompt available on replay — must read the cassette. + with RecordAgent.record(cassette) as agent: + replayed = await agent.prompt("suggest me jobs") + agent.assert_tool_called("job_search_tool") + + self.assertEqual([tc["name"] for tc in replayed.tool_calls], ["job_search_tool"]) + self.assertEqual(replayed.content, '[{"id": 2}]') + + async def test_synthesizes_a_transcript_for_a_plain_response(self): + response = AgentResponse(content="Hi there!", usage={"input": 3, "output": 2}) + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with _fake_prompt(response): + with RecordAgent.record(cassette) as agent: + await agent.prompt("hello") + + store = json.loads(open(cassette).read()) + (turn,) = store.values() + + self.assertEqual(turn[0], {"type": "human", "content": "hello"}) + self.assertEqual(turn[1]["type"], "ai") + self.assertEqual(turn[1]["content"], "Hi there!") + self.assertEqual(turn[1]["uses"], {"input_token": 3, "output_token": 2, "cache_token": 0, "total_token": 5}) + + +class TestBackwardCompatibleReplay(unittest.IsolatedAsyncioTestCase): + async def test_old_flat_cassette_still_replays(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "legacy.json") + with RecordAgent.record(cassette) as agent: + key = agent._key("suggest me jobs", None) + legacy = { + key: { + "content": "I found a job.", + "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "usage": {"input": 117, "output": 22}, + } + } + with open(cassette, "w") as f: + json.dump(legacy, f) + + with RecordAgent.record(cassette) as agent: + replayed = await agent.prompt("suggest me jobs") + agent.assert_tool_called("job_search_tool") + + self.assertEqual(replayed.content, "I found a job.") + self.assertEqual([tc["name"] for tc in replayed.tool_calls], ["job_search_tool"]) diff --git a/fastapi_startkit/tests/ai/test_recording.py b/fastapi_startkit/tests/ai/test_recording.py new file mode 100644 index 00000000..3c8cec80 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_recording.py @@ -0,0 +1,137 @@ +"""Tests for the ordered per-turn recording format (task #1251). + +A recorded turn is an ORDERED list of message dicts capturing the full +interaction, keyed by the hash of the user input: + + [ + {"type": "human", "content": "suggest me jobs"}, + {"type": "ai", "tool_calls": [...], "uses": {...}, "response_time": 12.0}, + {"type": "tool_response", "content_type": "json", "content": "[...]", "response_time": 5.0}, + {"type": "ai", "content": "Here is a job...", "uses": {...}, "response_time": 8.0}, + ] +""" + +from fastapi_startkit.ai import recording + + +class TestEntryBuilders: + def test_human_entry(self): + assert recording.human("suggest me jobs") == {"type": "human", "content": "suggest me jobs"} + + def test_ai_entry_with_tool_calls_omits_empty_content(self): + entry = recording.ai( + tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1", "type": "tool_call"}], + uses={"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, + response_time=12.0, + ) + + assert entry["type"] == "ai" + assert "content" not in entry + assert entry["tool_calls"][0]["name"] == "job_search_tool" + assert entry["uses"]["input_token"] == 117 + assert entry["response_time"] == 12.0 + + def test_ai_entry_with_a_final_answer_carries_content(self): + entry = recording.ai(content="Here is a job.", uses={"output_token": 9}, response_time=8.0) + + assert entry["type"] == "ai" + assert entry["content"] == "Here is a job." + assert "tool_calls" not in entry + + def test_tool_response_infers_json_content_type(self): + entry = recording.tool_response(content='[{"id": 2}]', response_time=5.0) + + assert entry == { + "type": "tool_response", + "content_type": "json", + "content": '[{"id": 2}]', + "response_time": 5.0, + } + + def test_tool_response_infers_text_content_type(self): + entry = recording.tool_response(content="no jobs found", response_time=5.0) + + assert entry["content_type"] == "text" + + +class TestUsesFromUsageMetadata: + def test_maps_langchain_usage_metadata_to_the_uses_shape(self): + uses = recording.uses_from_usage_metadata( + { + "input_tokens": 117, + "output_tokens": 22, + "total_tokens": 139, + "input_token_details": {"cache_read": 40}, + } + ) + assert uses == {"input_token": 117, "output_token": 22, "cache_token": 40, "total_token": 139} + + def test_handles_missing_metadata(self): + assert recording.uses_from_usage_metadata(None) == { + "input_token": 0, + "output_token": 0, + "cache_token": 0, + "total_token": 0, + } + + +class TestReconstructResponse: + def test_final_ai_answer_wins_over_the_tool_result_for_content(self): + transcript = [ + recording.human("suggest me jobs"), + recording.ai( + tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + uses={"input_token": 117, "output_token": 22}, + response_time=12.0, + ), + recording.tool_response(content='[{"id": 2}]', response_time=5.0), + recording.ai(content="Here is a job.", uses={"input_token": 5, "output_token": 9}, response_time=8.0), + ] + + response = recording.to_response(transcript) + + assert response.content == "Here is a job." + assert [tc["name"] for tc in response.tool_calls] == ["job_search_tool"] + assert response.usage == {"input": 5, "output": 9} + assert len(response.tool_events) == 1 + assert response.tool_events[0]["content"] == '[{"id": 2}]' + + def test_falls_back_to_tool_result_when_there_is_no_final_ai_answer(self): + transcript = [ + recording.human("suggest me jobs"), + recording.ai( + tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + uses={"input_token": 117, "output_token": 22}, + response_time=12.0, + ), + recording.tool_response(content='[{"id": 2}]', response_time=5.0), + ] + + response = recording.to_response(transcript) + + assert response.content == '[{"id": 2}]' + assert [tc["name"] for tc in response.tool_calls] == ["job_search_tool"] + + def test_plain_text_turn_reconstructs_content_and_no_tool_calls(self): + transcript = [ + recording.human("hello"), + recording.ai(content="Hi there!", uses={"input_token": 3, "output_token": 2}, response_time=4.0), + ] + + response = recording.to_response(transcript) + + assert response.content == "Hi there!" + assert response.tool_calls == [] + + +class TestChunksRoundTrip: + def test_chunks_are_carried_on_the_ai_entry(self): + entry = recording.ai(content="Hi there!", chunks=["Hi", " there!"], response_time=4.0) + assert entry["chunks"] == ["Hi", " there!"] + + def test_chunks_from_transcript_collects_streamed_chunks(self): + transcript = [ + recording.human("hello"), + recording.ai(content="Hi there!", chunks=["Hi", " there!"], response_time=4.0), + ] + assert recording.chunks_from_transcript(transcript) == ["Hi", " there!"] diff --git a/fastapi_startkit/tests/ai/test_runner_recording.py b/fastapi_startkit/tests/ai/test_runner_recording.py new file mode 100644 index 00000000..78ca132e --- /dev/null +++ b/fastapi_startkit/tests/ai/test_runner_recording.py @@ -0,0 +1,99 @@ +"""The Runner must capture the full per-turn interaction (task #1251). + +A turn that calls a tool has to preserve the model's requested tool calls, the +tool's response and per-call latency, and the model usage — not collapse to the +last tool-result message (which drops ``tool_calls`` entirely). +""" + +import unittest + +from langchain_core.messages import AIMessage as LCAIMessage +from langchain_core.tools import tool + +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.runner import Runner + + +@tool +def echo_tool(x: str) -> str: + """Echo the input back.""" + return f"echoed:{x}" + + +class EchoAgent(Agent): + def tools(self): + return [echo_tool] + + +class FakeModel: + def __init__(self, response): + self._response = response + + async def ainvoke(self, messages): + return self._response + + +def _runner_with_model(response): + agent = EchoAgent() + runner = Runner(agent) + runner._build_model = lambda *a, **k: FakeModel(response) # type: ignore[method-assign] + return runner + + +class TestRunnerCapturesToolTurn(unittest.IsolatedAsyncioTestCase): + def _tool_ai_message(self): + return LCAIMessage( + content="", + tool_calls=[{"name": "echo_tool", "args": {"x": "hi"}, "id": "c1", "type": "tool_call"}], + usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + ) + + async def test_tool_calls_are_preserved_not_dropped(self): + runner = _runner_with_model(self._tool_ai_message()) + + response = await runner.run("hi") + + self.assertEqual([tc["name"] for tc in response.tool_calls], ["echo_tool"]) + + async def test_tool_response_and_latency_are_captured(self): + runner = _runner_with_model(self._tool_ai_message()) + + response = await runner.run("hi") + + self.assertEqual(len(response.tool_events), 1) + event = response.tool_events[0] + self.assertEqual(event["name"], "echo_tool") + self.assertEqual(event["content"], "echoed:hi") + self.assertGreaterEqual(event["response_time"], 0) + + async def test_usage_comes_from_the_requesting_ai_message(self): + runner = _runner_with_model(self._tool_ai_message()) + + response = await runner.run("hi") + + self.assertEqual(response.usage, {"input": 10, "output": 5}) + + async def test_transcript_is_an_ordered_ai_then_tool_response(self): + runner = _runner_with_model(self._tool_ai_message()) + + response = await runner.run("hi") + + kinds = [entry["type"] for entry in response.transcript] + self.assertEqual(kinds, ["ai", "tool_response"]) + self.assertEqual(response.transcript[0]["tool_calls"][0]["name"], "echo_tool") + self.assertEqual(response.transcript[1]["content"], "echoed:hi") + + +class TestRunnerPlainTurn(unittest.IsolatedAsyncioTestCase): + async def test_plain_answer_records_a_single_ai_message(self): + message = LCAIMessage( + content="Hi there!", usage_metadata={"input_tokens": 3, "output_tokens": 2, "total_tokens": 5} + ) + runner = _runner_with_model(message) + + response = await runner.run("hello") + + self.assertEqual(response.content, "Hi there!") + self.assertEqual(response.tool_calls, []) + self.assertEqual([e["type"] for e in response.transcript], ["ai"]) + self.assertEqual(response.transcript[0]["content"], "Hi there!") From cb9dd29c8ce93fc03a072c332676b79f222955a7 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sat, 25 Jul 2026 12:50:50 -0700 Subject: [PATCH 10/13] AI: add agent.assert_tokens() for accumulated token limits (#200) * feat(ai): add agent.assert_tokens() for accumulated token limits agent.assert_tokens(lambda x: x.where('input', '<=', 5000).where('output', '<=', 5000)) Tokens accumulate across every prompt()/stream() in a session (summing each ai message's token uses from the transcript, with a usage fallback), and the fluent TokenQuery evaluates where-clauses over input/output/cache/total. Works on both live-record and replay; to_response() now carries the transcript so replayed turns expose their per-message uses. * style(ai): ruff-format assert_tokens docstring * style(ai): ruff-format runner.py and test_record_cassette_format.py CI runs 'ruff format --check .' repo-wide; these two files (carried in from #199) were unformatted and failed the check. * test(orm): fix polymorphic relation tests to use the record relationship The Like model's MorphTo relationship is named 'record'; the tests referenced a nonexistent 'like.log' attribute, which returned None and failed to await / assert. Aligns with the eager-load case that already uses with_('record'). --- .../tests/units/agents/test_router_agent.py | 3 + .../src/fastapi_startkit/ai/recording.py | 22 ++++- .../src/fastapi_startkit/ai/runner.py | 4 +- .../src/fastapi_startkit/ai/testing.py | 69 ++++++++++++++ .../tests/ai/test_assert_tokens.py | 93 +++++++++++++++++++ .../tests/ai/test_record_cassette_format.py | 25 +++-- .../relationships/test_sqlite_polymorphic.py | 4 +- 7 files changed, 207 insertions(+), 13 deletions(-) create mode 100644 fastapi_startkit/tests/ai/test_assert_tokens.py diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py index 3b966ea3..0925ca12 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -20,6 +20,9 @@ def assert_tool_calls(tool: AssertToolCall): await agent.prompt("suggest python developer jobs") agent.assert_tool_called("job_search_tool", assert_tool_calls) + # Tokens accumulate across both turns of the session. + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + async def test_the_router_with_initial_messages(self): with ChatAgent.record( "record_stream.json", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/recording.py b/fastapi_startkit/src/fastapi_startkit/ai/recording.py index 908dd90f..39299a24 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/recording.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/recording.py @@ -132,7 +132,27 @@ def to_response(transcript: list[dict]) -> AgentResponse: content = entry.get("content", "") runtime += (entry.get("response_time") or 0) / 1000 - return AgentResponse(content=content, tool_calls=tool_calls, usage=usage, tool_events=tool_events, runtime=runtime) + return AgentResponse( + content=content, + tool_calls=tool_calls, + usage=usage, + tool_events=tool_events, + runtime=runtime, + transcript=list(transcript), + ) + + +def accumulate_uses(transcript: list[dict], totals: dict) -> None: + """Add every ``ai`` entry's token ``uses`` in ``transcript`` into ``totals`` + (keys: input, output, cache, total).""" + for entry in transcript or []: + if entry.get("type") != "ai": + continue + uses = entry.get("uses") or {} + totals["input"] += uses.get("input_token", 0) + totals["output"] += uses.get("output_token", 0) + totals["cache"] += uses.get("cache_token", 0) + totals["total"] += uses.get("total_token", 0) def chunks_from_transcript(transcript: list[dict]) -> list[str]: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index b078b1c4..2d846f19 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -53,7 +53,9 @@ def _record_ai_message( tool_calls = list(getattr(message, "tool_calls", None) or []) uses = recording.uses_from_usage_metadata(getattr(message, "usage_metadata", None)) self._transcript.append( - recording.ai(content=content, tool_calls=tool_calls, uses=uses, response_time=response_time_ms, chunks=chunks) + recording.ai( + content=content, tool_calls=tool_calls, uses=uses, response_time=response_time_ms, chunks=chunks + ) ) if tool_calls: self._requested_tool_calls = tool_calls diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 6f651895..46918704 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -4,6 +4,7 @@ import hashlib import inspect import json +import operator import sys import time from collections.abc import AsyncIterator @@ -21,6 +22,48 @@ def _joined(value: Any) -> str: return "".join(value) if isinstance(value, list) else value +def _new_token_totals() -> dict: + return {"input": 0, "output": 0, "cache": 0, "total": 0} + + +class TokenQuery: + """Fluent predicate over accumulated token totals, e.g.:: + + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + + Fields: ``input`` / ``output`` / ``cache`` / ``total``. + """ + + _OPS = { + "<=": operator.le, + ">=": operator.ge, + "<": operator.lt, + ">": operator.gt, + "==": operator.eq, + "=": operator.eq, + "!=": operator.ne, + } + + def __init__(self, totals: dict) -> None: + self._totals = totals + self._checks: list[tuple[str, str, float]] = [] + + def where(self, field: str, op: str, value: float) -> "TokenQuery": + self._checks.append((field, op, value)) + return self + + def failures(self) -> list[str]: + problems: list[str] = [] + for field, op, value in self._checks: + compare = self._OPS.get(op) + if compare is None: + raise ValueError(f"Unsupported token operator {op!r}") + actual = self._totals.get(field, 0) + if not compare(actual, value): + problems.append(f"{field}={actual} is not {op} {value}") + return problems + + class AgentFake: def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._agent_cls = agent_cls @@ -30,6 +73,7 @@ def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._records: list[dict] = [] self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None + self._tokens: dict = _new_token_totals() def _history(self) -> list: return self._records @@ -73,9 +117,31 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter self._remember(message, self._last_response) def _remember(self, message: str, response: AgentResponse) -> None: + self._accumulate_tokens(response) self._records.append({"role": "user", "content": message}) self._records.append({"role": "assistant", "content": response.content}) + def _accumulate_tokens(self, response: AgentResponse) -> None: + """Add a turn's token usage to the running totals. Prefer the per-message + ``uses`` on the transcript; fall back to the response's summary usage.""" + from . import recording # noqa: PLC0415 + + if response.transcript: + recording.accumulate_uses(response.transcript, self._tokens) + else: + self._tokens["input"] += response.usage.get("input", 0) + self._tokens["output"] += response.usage.get("output", 0) + + def assert_tokens(self, predicate: Callable[[TokenQuery], Any]) -> None: + """Assert on the tokens accumulated across every prompt()/stream() so far. + + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + """ + query = TokenQuery(dict(self._tokens)) + predicate(query) + failures = query.failures() + assert not failures, f"Token assertion failed: {'; '.join(failures)}. Accumulated tokens: {self._tokens}" + def assert_prompt(self, expected: str | Callable[[str], bool]) -> None: if callable(expected): assert any(expected(p) for p in self._prompts), ( @@ -107,6 +173,7 @@ def assert_not_prompted(self) -> None: def reset(self) -> "AgentFake": self._records.clear() self._last_response = None + self._tokens = _new_token_totals() return self def _require_response(self) -> AgentResponse: @@ -194,6 +261,7 @@ def __init__(self, real: Agent, cassette: str | None = None, messages: list | No self._real.messages = self._history # type: ignore[method-assign] self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None + self._tokens: dict = _new_token_totals() def _history(self) -> list: return self._seed_messages + self._records @@ -263,6 +331,7 @@ def _response_from_cache(value: Any) -> AgentResponse: return AgentResponse(content=_joined(value)) def _remember_turn(self, message: str, response: AgentResponse) -> None: + self._accumulate_tokens(response) self._records.append({"role": "user", "content": message}) turn: dict[str, Any] = {"role": "assistant", "content": response.content} if response.tool_calls: diff --git a/fastapi_startkit/tests/ai/test_assert_tokens.py b/fastapi_startkit/tests/ai/test_assert_tokens.py new file mode 100644 index 00000000..8cfe9447 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_assert_tokens.py @@ -0,0 +1,93 @@ +"""Tests for agent.assert_tokens() — assert on accumulated token usage (task #1251). + + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + +Tokens accumulate across every recorded turn (and every ai message within a +turn); the predicate builds where-clauses that must all hold. +""" + +import os +import tempfile +import unittest +from unittest import mock + +from fastapi_startkit.ai import recording +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.response import AgentResponse + + +class TokenAgent(Agent): + pass + + +def _ai(input_tokens: int, output_tokens: int) -> dict: + return recording.ai( + content="reply", + uses={ + "input_token": input_tokens, + "output_token": output_tokens, + "cache_token": 0, + "total_token": input_tokens + output_tokens, + }, + response_time=1.0, + ) + + +def _fake_prompt(responses: list): + queue = list(responses) + + async def prompt(agent_self, message, **kwargs): + return queue.pop(0) + + return mock.patch.object(TokenAgent, "prompt", prompt) + + +class TestAssertTokens(unittest.IsolatedAsyncioTestCase): + async def test_passes_when_accumulated_tokens_are_within_limits(self): + responses = [ + AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=[_ai(100, 20)]), + AgentResponse(content="b", usage={"input": 200, "output": 30}, transcript=[_ai(200, 30)]), + ] + with tempfile.TemporaryDirectory() as tmp: + with _fake_prompt(responses): + with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("a") + await agent.prompt("b") + # accumulated: input=300, output=50 + agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) + + async def test_fails_when_a_limit_is_exceeded(self): + responses = [AgentResponse(content="a", usage={"input": 400, "output": 20}, transcript=[_ai(400, 20)])] + with tempfile.TemporaryDirectory() as tmp: + with _fake_prompt(responses): + with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("a") + with self.assertRaises(AssertionError): + agent.assert_tokens(lambda x: x.where("input", "<=", 300)) + + async def test_supports_cache_and_total_fields(self): + transcript = [ + recording.ai( + content="a", + uses={"input_token": 100, "output_token": 20, "cache_token": 40, "total_token": 160}, + response_time=1.0, + ) + ] + with tempfile.TemporaryDirectory() as tmp: + with _fake_prompt([AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=transcript)]): + with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("a") + agent.assert_tokens(lambda x: x.where("cache", "<=", 40).where("total", ">=", 160)) + + async def test_accumulates_from_the_cassette_on_replay(self): + responses = [AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=[_ai(100, 20)])] + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with _fake_prompt(responses): + with TokenAgent.record(cassette) as agent: + await agent.prompt("a") + + # Replay: no live prompt available; tokens must come from the cassette. + with TokenAgent.record(cassette) as agent: + await agent.prompt("a") + agent.assert_tokens(lambda x: x.where("input", "==", 100).where("output", "==", 20)) diff --git a/fastapi_startkit/tests/ai/test_record_cassette_format.py b/fastapi_startkit/tests/ai/test_record_cassette_format.py index e5122e22..92e0b4b9 100644 --- a/fastapi_startkit/tests/ai/test_record_cassette_format.py +++ b/fastapi_startkit/tests/ai/test_record_cassette_format.py @@ -30,12 +30,15 @@ class TestNewCassetteFormat(unittest.IsolatedAsyncioTestCase): async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): response = AgentResponse( transcript=[ - {"type": "ai", "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - "uses": {"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, - "response_time": 12.0}, - {"type": "tool_response", "content_type": "json", "content": "[{\"id\": 2}]", "response_time": 5.0}, + { + "type": "ai", + "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "uses": {"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, + "response_time": 12.0, + }, + {"type": "tool_response", "content_type": "json", "content": '[{"id": 2}]', "response_time": 5.0}, ], - content="[{\"id\": 2}]", + content='[{"id": 2}]', tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], ) with tempfile.TemporaryDirectory() as tmp: @@ -57,11 +60,15 @@ async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): async def test_replaying_a_new_format_cassette_reconstructs_tool_calls(self): response = AgentResponse( transcript=[ - {"type": "ai", "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - "uses": {"input_token": 117, "output_token": 22}, "response_time": 12.0}, - {"type": "tool_response", "content_type": "json", "content": "[{\"id\": 2}]", "response_time": 5.0}, + { + "type": "ai", + "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], + "uses": {"input_token": 117, "output_token": 22}, + "response_time": 12.0, + }, + {"type": "tool_response", "content_type": "json", "content": '[{"id": 2}]', "response_time": 5.0}, ], - content="[{\"id\": 2}]", + content='[{"id": 2}]', tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], ) with tempfile.TemporaryDirectory() as tmp: 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 f86fc1ef..675cf79a 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.log + record = await like.record assert isinstance(record, (Articles, Product)) async def test_can_get_eager_load_polymorphic_relation(self): likes = await Like.with_("record").get() for like in likes: - assert isinstance(like.log, (Articles, Product)) + assert isinstance(like.record, (Articles, Product)) From 413194294213d723f0cf0468143d54bb3e257f8d Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Sun, 26 Jul 2026 00:01:40 -0700 Subject: [PATCH 11/13] test: relocate test_structures to tests/support and fix import Main added tests/utils/test_structures.py (task #1214) importing from fastapi_startkit.utils.structures, but this branch relocated the module to fastapi_startkit.support.structures. On the PR merge the old module was gone, breaking collection with ModuleNotFoundError. Move the test to tests/support/test_structures.py to mirror the source layout and point the import at fastapi_startkit.support.structures. Test content and coverage are preserved unchanged. --- fastapi_startkit/tests/{utils => support}/test_structures.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) rename fastapi_startkit/tests/{utils => support}/test_structures.py (94%) diff --git a/fastapi_startkit/tests/utils/test_structures.py b/fastapi_startkit/tests/support/test_structures.py similarity index 94% rename from fastapi_startkit/tests/utils/test_structures.py rename to fastapi_startkit/tests/support/test_structures.py index 5b0dab92..f512c33c 100644 --- a/fastapi_startkit/tests/utils/test_structures.py +++ b/fastapi_startkit/tests/support/test_structures.py @@ -1,10 +1,10 @@ -"""Unit tests for the data-structure helpers in utils/structures.py (task #1214).""" +"""Unit tests for the data-structure helpers in support/structures.py (task #1214).""" import pytest from dotty_dict import Dotty from fastapi_startkit.exceptions.exceptions import LoaderNotFound -from fastapi_startkit.utils.structures import data, data_get, data_set, load +from fastapi_startkit.support.structures import data, data_get, data_set, load class TestData: From e164604299cb53c091e90aa8bba43acbb74c94a7 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Mon, 27 Jul 2026 05:04:48 -0700 Subject: [PATCH 12/13] feat: add fixes --- example/agents/app/agents/agent.py | 69 ++++++++ example/agents/app/agents/job_search_graph.py | 163 ++++++++++++++++++ example/agents/app/agents/state.py | 43 +++++ example/agents/app/models/message.py | 20 +++ example/agents/app/models/thread.py | 13 ++ example/agents/app/tools/job_search_tool.py | 2 +- example/agents/bootstrap/application.py | 3 + example/agents/config/database.py | 27 +++ .../2026_07_27_000001_create_threads_table.py | 17 ++ ...2026_07_27_000002_create_messages_table.py | 23 +++ ...07_27_000003_add_type_to_messages_table.py | 21 +++ ...sage_meta_attachments_to_messages_table.py | 19 ++ example/agents/package-lock.json | 10 ++ example/agents/package.json | 1 + example/agents/pyproject.toml | 1 + .../agents/resources/js/components/Chat.tsx | 22 ++- .../resources/js/components/JobChat.tsx | 159 +++++++++++++++++ .../resources/js/pages/chat/jobs/Index.tsx | 5 + example/agents/routes/api.py | 140 ++++++++++++++- example/agents/tests/features/job_search.json | 6 - .../tests/features/record_no_stream.json | 6 - .../tests/units/agents/basic_job_search.json | 24 +-- .../tests/units/agents/record_sales.json | 8 +- .../tests/units/agents/record_stream.json | 36 ++-- .../tests/units/agents/stream_job_search.json | 31 ++-- .../tests/units/agents/test_router_agent.py | 10 ++ example/agents/tinker.py | 63 +------ example/agents/uv.lock | 42 +++++ .../src/fastapi_startkit/ai/agent.py | 7 +- .../src/fastapi_startkit/ai/ai.py | 6 +- .../src/fastapi_startkit/ai/graph.py | 17 +- .../src/fastapi_startkit/ai/pipeline.py | 13 +- .../src/fastapi_startkit/ai/runner.py | 26 +-- .../src/fastapi_startkit/ai/testing.py | 39 +++-- fastapi_startkit/tests/ai/test_agent.py | 34 ++-- fastapi_startkit/tests/ai/test_agent_fake.py | 10 +- .../tests/ai/test_agent_record_fluent.py | 2 +- .../tests/ai/test_assert_response_time.py | 106 ++++++++++++ fastapi_startkit/tests/ai/test_graph_agent.py | 2 +- 39 files changed, 1062 insertions(+), 184 deletions(-) create mode 100644 example/agents/app/agents/agent.py create mode 100644 example/agents/app/agents/job_search_graph.py create mode 100644 example/agents/app/agents/state.py create mode 100644 example/agents/app/models/message.py create mode 100644 example/agents/app/models/thread.py create mode 100644 example/agents/config/database.py create mode 100644 example/agents/databases/migrations/2026_07_27_000001_create_threads_table.py create mode 100644 example/agents/databases/migrations/2026_07_27_000002_create_messages_table.py create mode 100644 example/agents/databases/migrations/2026_07_27_000003_add_type_to_messages_table.py create mode 100644 example/agents/databases/migrations/2026_07_27_000004_add_usage_meta_attachments_to_messages_table.py create mode 100644 example/agents/resources/js/components/JobChat.tsx create mode 100644 example/agents/resources/js/pages/chat/jobs/Index.tsx delete mode 100644 example/agents/tests/features/job_search.json delete mode 100644 example/agents/tests/features/record_no_stream.json create mode 100644 fastapi_startkit/tests/ai/test_assert_response_time.py diff --git a/example/agents/app/agents/agent.py b/example/agents/app/agents/agent.py new file mode 100644 index 00000000..d56c43fe --- /dev/null +++ b/example/agents/app/agents/agent.py @@ -0,0 +1,69 @@ +import json + +from fastapi_startkit.ai import Agent, Middleware +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage +from langchain_core.runnables import RunnableConfig +from langchain_core.tools import BaseTool + +from app.agents.job_search_graph import USER_PROFILE +from app.agents.state import Context, OverallState +from app.middleware.agent_logger import AgentLogger +from app.models.message import Message +from app.tools.job_search_tool import job_search_tool + +JOB_SEARCH_PROMPT = """You find jobs for a user. ALWAYS call job_search_tool exactly once — +never reply with text only, and never ask clarifying questions. + +Derive role/skill/location keywords from the LATEST user message only (e.g. 'python +developer remote'); it replaces any earlier search. Use earlier messages solely for +preferences like location or remote. Ignore filler words like 'suggest', 'me', 'jobs'. +If the latest message has no usable keywords, take them from the user's profile.""" + + +class JobSearchAgent(Agent): + """Plain single-LLM agent (like ChatAgent) that owns its own context: the thread's + past user messages plus the user profile, injected only when the router asked for + it via `contexts`. messages() is async — the runner awaits it. + """ + + # The model must search, never chat: text-only replies here would be discarded + # by job_search_node (it only reads tool_events). + tool_choice = "any" + + def __init__(self, state: OverallState, config: RunnableConfig | None = None): + super().__init__() + self.state = state + self.thread_id = ((config or {}).get("configurable") or {}).get("thread_id", "") + self.contexts: list[Context] = state.get("contexts") or [] + + async def messages(self) -> list[BaseMessage]: + messages: list[BaseMessage] = [] + + if Context.INCLUDE_USER_PROFILE in self.contexts: + messages.append(SystemMessage(content=f"User profile: {json.dumps(USER_PROFILE)}")) + + rows = await Message.where("thread_id", self.thread_id).where("role", "user").order_by("id").get() + messages.extend(HumanMessage(content=row.data.get("text", "")) for row in rows) + + if Context.INCLUDE_LAST_JOB_SEARCH_RESPONSE in self.contexts and (last := await self._last_job_search()): + messages.append(AIMessage(content=f"Previous job search results: {json.dumps(last.data)}")) + + return messages + + async def _last_job_search(self) -> Message | None: + return ( + await Message.where("thread_id", self.thread_id) + .where("role", "ai") + .where("type", "tool_response") + .order_by("id", "desc") + .first() + ) + + def instructions(self) -> str: + return JOB_SEARCH_PROMPT + + def middleware(self) -> list[Middleware]: + return [AgentLogger()] + + def tools(self) -> list[BaseTool]: + return [job_search_tool] diff --git a/example/agents/app/agents/job_search_graph.py b/example/agents/app/agents/job_search_graph.py new file mode 100644 index 00000000..7c217538 --- /dev/null +++ b/example/agents/app/agents/job_search_graph.py @@ -0,0 +1,163 @@ +import json +from functools import cache +from typing import cast + +from langchain.chat_models import init_chat_model +from langchain_core.language_models.chat_models import BaseChatModel +from langchain_core.messages import HumanMessage, SystemMessage +from langchain_core.runnables import RunnableConfig +from langgraph.graph import END, START, StateGraph +from langgraph.graph.state import CompiledStateGraph + +from app.agents.state import InputState, OutputState, OverallState, RouterOutput + + +@cache +def llm() -> BaseChatModel: + """Built lazily so importing this module doesn't require the API key at import time.""" + return init_chat_model("google_genai:gemini-3.1-flash-lite") + + +# Stand-in for a real profile lookup; the router decides when it's needed. +USER_PROFILE = { + "name": "Alice", + "title": "Python Developer", + "location": "Remote", + "skills": ["Python", "FastAPI", "PostgreSQL"], +} + + +ROUTER_PROMPT = """You route a career assistant's queries. + +- job_search: the user wants job listings / openings / roles. +- company_research: the user asks about a specific company. +- chat: greetings (hi / hello), thanks, small talk, or anything conversational. + +Also decide which extra context the next node needs: +- include_user_profile: the query is vague about role, skills or location, so the user's own profile helps. +- include_last_job_search_response: the user is refining an earlier search. +""" + + +async def router_node(state: InputState) -> dict: + """Context: only the user input. No history, no profile, no tool output. + + Handles greetings itself: for the 'chat' intent it answers inline via + RouterOutput.reply, so no extra node/LLM call is needed. + """ + model = llm().with_structured_output(RouterOutput) + decision = cast(RouterOutput, await model.ainvoke([SystemMessage(ROUTER_PROMPT), HumanMessage(state["query"])])) + + out: dict = {"intent": decision.intent, "contexts": decision.contexts} + if decision.intent == "chat": + out |= {"type": "text", "data": decision.reply, "data_type": "string"} + return out + + +def route(state: OverallState) -> str: + # chat is already answered inline by the router; company_research isn't built yet. + return "job_search" if state["intent"] == "job_search" else END + + +# --- job search ------------------------------------------------------------ + + +async def job_search_node(state: OverallState, config: RunnableConfig) -> dict: + """Context: the thread's history + the user profile — both owned by JobSearchAgent. + + Delegates the tool-calling loop to JobSearchAgent.prompt(), then pulls the raw job + rows out of the response's tool_events (the summarizer node turns them into prose). + Also reports the calls it made — {"name", "args"} — so they can be shown/persisted. + """ + from app.agents.agent import JobSearchAgent + + response = await JobSearchAgent(state, config).prompt(state["query"]) + + jobs: list[dict] = [] + calls: list[dict] = [] + for event in response.tool_events: + if event["name"] == "job_search_tool": + jobs.extend(json.loads(event["content"])) + calls.append({"name": event["name"], "args": event["args"]}) + + return {"type": "tool_response", "data": jobs, "data_type": "json", "tool_calls": calls} + + +# --- ask user (human-in-the-loop) ------------------------------------------ + + +def ask_user_node(state: OverallState) -> dict: + """No jobs found: pause the graph and ask the frontend to refine the search. + + interrupt() freezes the run (state is persisted by the checkpointer) and + surfaces this payload to the caller. The graph resumes when the caller sends + Command(resume=), and interrupt() returns that reply here. + """ + from langgraph.types import interrupt + + reply = interrupt( + { + "type": "question", + "reason": "no_jobs", + "message": f"No jobs matched '{state['query']}'. Try different keywords, or send blank to stop.", + } + ) + + return {"query": str(reply).strip(), "data": None} + + +def after_ask(state: OverallState) -> str: + return "job_search" if state["query"] else "job_summarizer" + + +# --- job summarizer -------------------------------------------------------- + +SUMMARIZER_PROMPT = """Write a short, friendly summary of these job listings for the user. +Group by role type, mention company and location. If the list is empty, say so plainly. +Use only what is given - do not invent jobs.""" + + +async def job_summarizer_node(state: OverallState) -> dict: + """Context: only the tool response. Never sees the query or the profile.""" + if not state["data"] and state["query"]: + # ask_user still wants a refined query; don't spend a call summarizing nothing + return {} + + response = await llm().ainvoke([SystemMessage(SUMMARIZER_PROMPT), HumanMessage(json.dumps(state["data"] or []))]) + + return {"type": "text", "data": str(response.text), "data_type": "string"} + + +def ask_or_finish(state: OverallState) -> str: + return "ask_user" if not state["data"] and state["query"] else END + + +def build() -> StateGraph: + builder = StateGraph(OverallState, input_schema=InputState, output_schema=OutputState) + + builder.add_node("router", router_node) + builder.add_node("job_search", job_search_node) + builder.add_node("ask_user", ask_user_node) + builder.add_node("job_summarizer", job_summarizer_node) + + builder.add_edge(START, "router") + builder.add_conditional_edges("router", route, ["job_search", END]) + builder.add_edge("job_search", "job_summarizer") + builder.add_conditional_edges("job_summarizer", ask_or_finish, ["ask_user", END]) + builder.add_conditional_edges("ask_user", after_ask, ["job_search", "job_summarizer"]) + + return builder + + +_graph: CompiledStateGraph | None = None + + +async def job_graph() -> CompiledStateGraph: + """Compile once, backed by the container's checkpointer (needed for interrupts).""" + global _graph + if _graph is None: + from fastapi_startkit.application import app + + checkpointer = await app().make("checkpointer") + _graph = build().compile(checkpointer=checkpointer) + return _graph diff --git a/example/agents/app/agents/state.py b/example/agents/app/agents/state.py new file mode 100644 index 00000000..1e6b5c61 --- /dev/null +++ b/example/agents/app/agents/state.py @@ -0,0 +1,43 @@ +from enum import StrEnum +from typing import Literal, TypedDict + +from pydantic import BaseModel, Field + + +class Context(StrEnum): + INCLUDE_USER_PROFILE = "include_user_profile" + INCLUDE_LAST_JOB_SEARCH_RESPONSE = "include_last_job_search_response" + + +class InputState(TypedDict): + query: str + + +class OutputState(TypedDict): + type: Literal["text", "tool_response"] + data: str | list | dict + data_type: Literal["string", "json"] + + +class OverallState(TypedDict): + query: str + intent: str + contexts: list[Context] + type: Literal["text", "tool_response"] + data: str | list | dict + data_type: Literal["string", "json"] + tool_calls: list[dict] # the calls behind a tool_response: [{"name": ..., "args": ...}] + +class RouterOutput(BaseModel): + """How to handle the user's query.""" + + intent: Literal["job_search", "company_research", "chat"] + contexts: list[Context] = Field( + default_factory=list, + description="Extra context the downstream node needs to answer well.", + ) + reply: str = Field( + "", + description="A short, friendly reply to the user. Fill this ONLY when intent is 'chat' " + "(greetings, thanks, small talk); leave empty for job_search and company_research.", + ) diff --git a/example/agents/app/models/message.py b/example/agents/app/models/message.py new file mode 100644 index 00000000..5e1f6e67 --- /dev/null +++ b/example/agents/app/models/message.py @@ -0,0 +1,20 @@ +from fastapi_startkit.masoniteorm import BelongsTo +from fastapi_startkit.masoniteorm.models import Model + + +class Message(Model): + __table__ = "messages" + + thread_id: str + role: str # user | ai + type: str # text | tool_call | tool_response | web_search + data: dict + data_type: str # string | json + run_id: str | None + # Annotated as plain `dict` (not `dict | None`): only the exact annotation selects + # the ORM's JsonCast; unions fall through to the str cast. Nullable in the DB. + usage: dict # token counts: {"input": n, "output": n} + meta: dict # per-message extras: latency, model, node + attachments: dict # documents sent with the message + + thread = BelongsTo("Thread", local_key="thread_id", foreign_key="id") diff --git a/example/agents/app/models/thread.py b/example/agents/app/models/thread.py new file mode 100644 index 00000000..17ae8081 --- /dev/null +++ b/example/agents/app/models/thread.py @@ -0,0 +1,13 @@ +from fastapi_startkit.masoniteorm import HasMany +from fastapi_startkit.masoniteorm.models import Model + + +class Thread(Model): + __table__ = "threads" + __primary_key__ = "id" + __incrementing__ = False # id is the string conversation key, not auto-increment + + id: str + title: str | None + + messages = HasMany("Message", local_key="id", foreign_key="thread_id") diff --git a/example/agents/app/tools/job_search_tool.py b/example/agents/app/tools/job_search_tool.py index 63f96a53..556218cc 100644 --- a/example/agents/app/tools/job_search_tool.py +++ b/example/agents/app/tools/job_search_tool.py @@ -2,7 +2,7 @@ jobs = [ {"id": 1, "title": "Software Engineer", "location": "San Francisco", "company": "Acme Corp", "type": "Full-time"}, - {"id": 2, "title": "Frontend Developer", "location": "Remote", "company": "Startup Inc", "type": "Full-time"}, + {"id": 2, "title": "Python Developer", "location": "Remote", "company": "Startup Inc", "type": "Full-time"}, {"id": 3, "title": "Data Scientist", "location": "New York", "company": "DataCo", "type": "Full-time"}, {"id": 4, "title": "DevOps Engineer", "location": "Austin", "company": "CloudBase", "type": "Contract"}, {"id": 5, "title": "Product Manager", "location": "Remote", "company": "ProductHQ", "type": "Full-time"}, diff --git a/example/agents/bootstrap/application.py b/example/agents/bootstrap/application.py index 47011cfe..b8e0a6a7 100644 --- a/example/agents/bootstrap/application.py +++ b/example/agents/bootstrap/application.py @@ -3,10 +3,12 @@ from fastapi_startkit import Application from fastapi_startkit.inertia import InertiaProvider from fastapi_startkit.logging import LogProvider +from fastapi_startkit.masoniteorm import DatabaseProvider from fastapi_startkit.skills import AISkillProvider from fastapi_startkit.vite import ViteProvider from fastapi_startkit.ai import AIProvider +from config.database import DatabaseConfig from config.fastapi import FastAPIConfig from config.logging import LoggingConfig from config.vite import ViteConfig @@ -19,6 +21,7 @@ AISkillProvider, LangChainProvider, (LogProvider,LoggingConfig), + (DatabaseProvider, DatabaseConfig), (FastapiProvider, FastAPIConfig), AIProvider, (ViteProvider, ViteConfig), diff --git a/example/agents/config/database.py b/example/agents/config/database.py new file mode 100644 index 00000000..05a720d9 --- /dev/null +++ b/example/agents/config/database.py @@ -0,0 +1,27 @@ +from dataclasses import field +from typing import Any + +from fastapi_startkit.environment import env +from fastapi_startkit.masoniteorm import PostgresConfig +from pydantic.dataclasses import dataclass + + +@dataclass +class DatabaseConfig: + default: str = field(default_factory=lambda: env("DB_CONNECTION", "postgres")) + + connections: dict[str, dict[str, Any]] = field( + default_factory=lambda: { + # Same Postgres instance the LangGraph checkpointer uses, so a thread row + # here lines up with the checkpointer's thread_id. + "postgres": PostgresConfig( + driver="postgres", + host=env("DB_HOST", "localhost"), + port=env("DB_PORT", 5432), + database=env("DB_DATABASE", "agents"), + username=env("DB_USERNAME", "postgres"), + password=env("DB_PASSWORD", "postgres"), + sslmode=env("DB_SSLMODE", "disable"), + ), + } + ) diff --git a/example/agents/databases/migrations/2026_07_27_000001_create_threads_table.py b/example/agents/databases/migrations/2026_07_27_000001_create_threads_table.py new file mode 100644 index 00000000..47330243 --- /dev/null +++ b/example/agents/databases/migrations/2026_07_27_000001_create_threads_table.py @@ -0,0 +1,17 @@ +"""CreateThreadsTable Migration.""" + +from fastapi_startkit.masoniteorm import Migration + + +class CreateThreadsTable(Migration): + async def up(self): + """Run the migrations.""" + async with await self.schema.create("threads") as table: + # id is the conversation key (same value as the LangGraph checkpointer thread_id). + table.string("id").primary() + table.string("title").nullable() + table.timestamps() + + async def down(self): + """Revert the migrations.""" + await self.schema.drop("threads") diff --git a/example/agents/databases/migrations/2026_07_27_000002_create_messages_table.py b/example/agents/databases/migrations/2026_07_27_000002_create_messages_table.py new file mode 100644 index 00000000..ad05d970 --- /dev/null +++ b/example/agents/databases/migrations/2026_07_27_000002_create_messages_table.py @@ -0,0 +1,23 @@ +"""CreateMessagesTable Migration.""" + +from fastapi_startkit.masoniteorm import Migration + + +class CreateMessagesTable(Migration): + async def up(self): + """Run the migrations.""" + async with await self.schema.create("messages") as table: + table.increments("id") + table.string("thread_id") + table.string("role", length=10) # user | ai + table.json("data") # holds either a JSON payload (job cards) or a JSON string (prose) + table.string("data_type", length=10) # string | json + table.string("run_id").nullable() # LangGraph run id; null for user messages + table.timestamps() + + table.index("thread_id") + table.foreign("thread_id").references("id").on("threads").on_delete("cascade") + + async def down(self): + """Revert the migrations.""" + await self.schema.drop("messages") diff --git a/example/agents/databases/migrations/2026_07_27_000003_add_type_to_messages_table.py b/example/agents/databases/migrations/2026_07_27_000003_add_type_to_messages_table.py new file mode 100644 index 00000000..d60956e6 --- /dev/null +++ b/example/agents/databases/migrations/2026_07_27_000003_add_type_to_messages_table.py @@ -0,0 +1,21 @@ +"""AddTypeToMessagesTable Migration.""" + +from fastapi_startkit.masoniteorm import Migration + + +class AddTypeToMessagesTable(Migration): + async def up(self): + """Run the migrations.""" + async with await self.schema.table("messages") as table: + # What the row is: text | tool_call | tool_response | web_search. + # Existing rows default to text; tool responses are backfilled below. + table.string("type", length=20).default("text") + + from app.models.message import Message + + await Message.where("data_type", "json").update({"type": "tool_response"}) + + async def down(self): + """Revert the migrations.""" + async with await self.schema.table("messages") as table: + table.drop_column("type") diff --git a/example/agents/databases/migrations/2026_07_27_000004_add_usage_meta_attachments_to_messages_table.py b/example/agents/databases/migrations/2026_07_27_000004_add_usage_meta_attachments_to_messages_table.py new file mode 100644 index 00000000..40f9f5bb --- /dev/null +++ b/example/agents/databases/migrations/2026_07_27_000004_add_usage_meta_attachments_to_messages_table.py @@ -0,0 +1,19 @@ +"""AddUsageMetaAttachmentsToMessagesTable Migration.""" + +from fastapi_startkit.masoniteorm import Migration + + +class AddUsageMetaAttachmentsToMessagesTable(Migration): + async def up(self): + """Run the migrations.""" + async with await self.schema.table("messages") as table: + table.jsonb("usage").nullable() # token counts: {"input": n, "output": n} + table.jsonb("meta").nullable() # anything per-message: latency, model, node + table.jsonb("attachments").nullable() # documents sent with the message + + async def down(self): + """Revert the migrations.""" + async with await self.schema.table("messages") as table: + table.drop_column("usage") + table.drop_column("meta") + table.drop_column("attachments") diff --git a/example/agents/package-lock.json b/example/agents/package-lock.json index 64a422b0..e0880cad 100644 --- a/example/agents/package-lock.json +++ b/example/agents/package-lock.json @@ -7,6 +7,7 @@ "dependencies": { "@inertiajs/react": "^3.4.0", "@vitejs/plugin-react": "^6.0.2", + "eventsource-parser": "^3.1.0", "react": "^19.2.7", "react-dom": "^19.2.7" }, @@ -804,6 +805,15 @@ "benchmarks" ] }, + "node_modules/eventsource-parser": { + "version": "3.1.0", + "resolved": "https://registry.npmjs.org/eventsource-parser/-/eventsource-parser-3.1.0.tgz", + "integrity": "sha512-kJezFj9YFAMLeORyi7aCLxLbD5/qWMQnoMVlVPyHIll7lgRJCc3JVln9Vgl9nwQi0YkMnhdGTMNn7CkRRAptMg==", + "license": "MIT", + "engines": { + "node": ">=18.0.0" + } + }, "node_modules/fastapi-vite-plugin": { "version": "0.0.3", "resolved": "https://registry.npmjs.org/fastapi-vite-plugin/-/fastapi-vite-plugin-0.0.3.tgz", diff --git a/example/agents/package.json b/example/agents/package.json index 4296bade..96852080 100644 --- a/example/agents/package.json +++ b/example/agents/package.json @@ -18,6 +18,7 @@ "dependencies": { "@inertiajs/react": "^3.4.0", "@vitejs/plugin-react": "^6.0.2", + "eventsource-parser": "^3.1.0", "react": "^19.2.7", "react-dom": "^19.2.7" } diff --git a/example/agents/pyproject.toml b/example/agents/pyproject.toml index 4b1d78d7..eb5fb8b4 100644 --- a/example/agents/pyproject.toml +++ b/example/agents/pyproject.toml @@ -8,6 +8,7 @@ authors = [ readme = "README.md" requires-python = ">=3.12,<4.0" dependencies = [ + "asyncpg>=0.31.0", "cleo>=2.1.0", "fastapi-startkit[ai,database,fastapi,sqlite]==0.44.0", "langchain>=1.3.10", diff --git a/example/agents/resources/js/components/Chat.tsx b/example/agents/resources/js/components/Chat.tsx index 58bd29d2..f1ab2f27 100644 --- a/example/agents/resources/js/components/Chat.tsx +++ b/example/agents/resources/js/components/Chat.tsx @@ -1,4 +1,5 @@ import { useState, useRef, useEffect } from "react" +import { createParser, type EventSourceMessage } from "eventsource-parser" type Message = { role: "user" | "assistant" @@ -48,18 +49,29 @@ export default function Chat({ const decoder = new TextDecoder() if (!reader) return - while (true) { - const { done, value } = await reader.read() - if (done) break - const chunk = decoder.decode(value, { stream: true }) + const append = (text: string) => setMessages(prev => { const updated = [...prev] updated[updated.length - 1] = { role: "assistant", - content: updated[updated.length - 1].content + chunk, + content: updated[updated.length - 1].content + text, } return updated }) + + // SSE frames: deltas carry model tokens; tool_response frames carry tool output. + const parser = createParser({ + onEvent: (event: EventSourceMessage) => { + const frame = JSON.parse(event.data) + if (frame.type === "delta") append(frame.text) + else if (frame.type === "tool_response") append(frame.content) + }, + }) + + while (true) { + const { done, value } = await reader.read() + if (done) break + parser.feed(decoder.decode(value, { stream: true })) } } finally { setLoading(false) diff --git a/example/agents/resources/js/components/JobChat.tsx b/example/agents/resources/js/components/JobChat.tsx new file mode 100644 index 00000000..34e70a68 --- /dev/null +++ b/example/agents/resources/js/components/JobChat.tsx @@ -0,0 +1,159 @@ +import { useState, useRef, useEffect } from "react" +import { createParser, type EventSourceMessage } from "eventsource-parser" + +type Job = { + id: number + title: string + location: string + company: string + type: string +} + +type Frame = + | { kind: "delta"; node: string; text: string } + | { kind: "envelope"; node?: string; type: "tool_response"; data: Job[]; data_type: "json" } + | { kind: "envelope"; type: "text"; data: string; data_type: "string" } + | { kind: "interrupt"; type: string; reason: string; message: string } + +type Message = { + role: "user" | "assistant" + content: string + jobs?: Job[] + interrupt?: boolean +} + +type JobChatProps = { + title?: string + endpoint?: string + placeholder?: string + emptyState?: string +} + +export default function JobChat({ + title = "Jobs", + endpoint = "/jobs/stream", + placeholder = "Search for jobs...", + emptyState = "Ask me to find jobs.", +}: JobChatProps) { + const [messages, setMessages] = useState([]) + const [input, setInput] = useState("") + const [loading, setLoading] = useState(false) + const bottomRef = useRef(null) + + useEffect(() => { + bottomRef.current?.scrollIntoView({ behavior: "smooth" }) + }, [messages]) + + const patchLast = (patch: (last: Message) => Message) => + setMessages(prev => [...prev.slice(0, -1), patch(prev[prev.length - 1])]) + + const applyFrame = (frame: Frame) => { + if (frame.kind === "delta") { + // The backend streams every node; the router's deltas are raw routing + // JSON, so this UI chooses not to render them. + if (frame.node === "router") return + patchLast(last => ({ ...last, content: last.content + frame.text })) + } else if (frame.kind === "envelope" && frame.type === "tool_response") { + patchLast(last => ({ ...last, jobs: frame.data })) + } else if (frame.kind === "envelope" && frame.type === "text") { + patchLast(last => ({ ...last, content: frame.data })) + } else if (frame.kind === "interrupt") { + patchLast(last => ({ ...last, content: frame.message, interrupt: true })) + } + } + + const handleSubmit = async (e: { preventDefault(): void }) => { + e.preventDefault() + if (!input.trim() || loading) return + + const userMessage = input.trim() + setInput("") + setMessages(prev => [...prev, { role: "user", content: userMessage }]) + setLoading(true) + setMessages(prev => [...prev, { role: "assistant", content: "" }]) + + try { + const response = await fetch(endpoint, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ message: userMessage }), + }) + + const reader = response.body?.getReader() + const decoder = new TextDecoder() + if (!reader) return + + const parser = createParser({ + onEvent: (event: EventSourceMessage) => applyFrame(JSON.parse(event.data) as Frame), + }) + + while (true) { + const { done, value } = await reader.read() + if (done) break + parser.feed(decoder.decode(value, { stream: true })) + } + } finally { + setLoading(false) + } + } + + return ( +
+

{title}

+ +
+ {messages.length === 0 && ( +

{emptyState}

+ )} + {messages.map((msg, i) => ( +
+
+ {msg.jobs && msg.jobs.length > 0 && ( +
+ {msg.jobs.map(job => ( +
+

{job.title}

+

+ {job.company} · {job.location} · {job.type} +

+
+ ))} +
+ )} + {(msg.content || msg.role === "user" || (loading && i === messages.length - 1)) && ( +
+ {msg.content || (loading && i === messages.length - 1 ? "▋" : "")} +
+ )} +
+
+ ))} +
+
+ +
+ setInput(e.target.value)} + placeholder={placeholder} + disabled={loading} + /> + +
+
+ ) +} diff --git a/example/agents/resources/js/pages/chat/jobs/Index.tsx b/example/agents/resources/js/pages/chat/jobs/Index.tsx new file mode 100644 index 00000000..69f99357 --- /dev/null +++ b/example/agents/resources/js/pages/chat/jobs/Index.tsx @@ -0,0 +1,5 @@ +import JobChat from "../../../components/JobChat" + +export default function Index() { + return +} diff --git a/example/agents/routes/api.py b/example/agents/routes/api.py index e3939fc5..df429931 100644 --- a/example/agents/routes/api.py +++ b/example/agents/routes/api.py @@ -1,9 +1,15 @@ +import json + from fastapi import APIRouter, Request from fastapi.responses import StreamingResponse from fastapi_startkit.inertia import Inertia +from langgraph.types import Command from app.agents.chat import ChatAgent from app.agents.graph_agent import SalesAgent +from app.agents.job_search_graph import job_graph +from app.models.message import Message +from app.models.thread import Thread from app.requests.chat import ChatRequest api = APIRouter() @@ -28,8 +34,8 @@ async def chat(request: ChatRequest): @api.post("/chat/stream") async def chat_stream(request: ChatRequest): async def generate(): - async for chunk in ChatAgent().stream(request.message): - yield f"data: {chunk}\n\n" + async for frame in ChatAgent().stream(request.message): + yield f"data: {json.dumps(frame)}\n\n" return StreamingResponse(generate(), media_type="text/event-stream") @@ -39,8 +45,8 @@ async def sales_stream(request: ChatRequest): config = {"configurable": {"thread_id": request.thread_id}} async def generate(): - async for chunk in SalesAgent().stream(request.message, config=config): - yield f"data: {chunk}\n\n" + async for frame in SalesAgent().stream(request.message, config=config): + yield f"data: {json.dumps(frame)}\n\n" return StreamingResponse(generate(), media_type="text/event-stream") @@ -50,3 +56,129 @@ async def sales_page(): return Inertia.render( "chat/sales/Index", ) + + +@api.get("/jobs") +async def jobs_page(): + return Inertia.render("chat/jobs/Index") + + +@api.post("/jobs/stream") +async def jobs_stream(request: ChatRequest): + graph = await job_graph() + config = {"configurable": {"thread_id": request.thread_id}} + + def sse(frame: dict) -> str: + return f"data: {json.dumps(frame)}\n\n" + + async def generate(): + # A pending interrupt on this thread means the last turn asked the user to + # refine; this message is the reply, so resume. Otherwise it's a fresh query. + snapshot = await graph.aget_state(config) + payload = Command(resume=request.message) if snapshot.interrupts else {"query": request.message} + + run_id: str | None = None + tool_rows: list[tuple[str, dict]] = [] # (node, row) + streamed: list[str] = [] + reply_node = "ask_user" # overwritten by whichever node actually speaks + usage_by_node: dict[str, dict] = {} + model_by_node: dict[str, str] = {} + + async for event in graph.astream_events(payload, config): + run_id = run_id or event.get("run_id") + node = (event.get("metadata") or {}).get("langgraph_node", "") + + # Token usage per node, from every LLM call the run makes (including the + # buffered ones inside JobSearchAgent — their model events still surface). + if event["event"] == "on_chat_model_end" and node: + meta = getattr(event["data"].get("output"), "usage_metadata", None) or {} + if meta: + totals = usage_by_node.setdefault(node, {"input": 0, "output": 0}) + totals["input"] += meta.get("input_tokens", 0) + totals["output"] += meta.get("output_tokens", 0) + if model_name := (event.get("metadata") or {}).get("ls_model_name"): + model_by_node[node] = model_name + + # Capture every node that commits a tool_response envelope (job_search today, + # company_research tomorrow): downstream nodes overwrite state["data"], so the + # raw rows only exist here. Forward them so the UI can render before the prose. + if event["event"] == "on_chain_end" and node and event["name"] == node: + out = event["data"].get("output") or {} + if isinstance(out, dict) and out.get("type") == "tool_response": + rows = out.get("data") or [] + calls = out.get("tool_calls") or [] + for call in calls: + tool_rows.append((node, {"type": "tool_call", "data": call, "data_type": "json"})) + tool_rows.append((node, {"type": "tool_response", "data": rows, "data_type": "json"})) + yield sse( + { + "kind": "envelope", + "node": node, + "type": "tool_response", + "data": rows, + "data_type": "json", + "tool_calls": calls, + } + ) + + if event["event"] != "on_chat_model_stream": + continue + # Stream every node's tokens, tagged with the node; the frontend decides + # what to render (the router's are structured-output JSON, for instance). + if (chunk := event["data"].get("chunk")) and (text := str(chunk.text)): + if node != "router": # persisted reply = user-facing text only + streamed.append(text) + reply_node = node + yield sse({"kind": "delta", "node": node, "text": text}) + + snapshot = await graph.aget_state(config) + reply = "".join(streamed) + if snapshot.interrupts: + # No jobs -> the graph paused at ask_user; surface its question for the UI. + reply = snapshot.interrupts[0].value["message"] + reply_node = "ask_user" + yield sse({"kind": "interrupt", **snapshot.interrupts[0].value}) + elif not reply and snapshot.values.get("type") == "text": + # Greeting answered inline by the router (not token-streamed) -> emit it whole. + reply = str(snapshot.values.get("data") or "") + reply_node = "router" + yield sse({"kind": "envelope", "type": "text", "data": reply, "data_type": "string"}) + + def enrich(node: str) -> dict: + meta: dict = {"node": node} + if model_name := model_by_node.get(node): + meta["model"] = model_name + extra: dict = {"meta": meta} + if usage := usage_by_node.get(node): + extra["usage"] = usage + return extra + + # Persist the turn only now: inserting the user row before the run would make + # JobSearchAgent load the current query as history and duplicate it. + await Thread.first_or_create({"id": request.thread_id}, {"id": request.thread_id}) + await Message.create( + { + "thread_id": request.thread_id, + "role": "user", + "type": "text", + "data": {"text": request.message}, + "data_type": "string", + "run_id": None, + } + ) + for node, row in tool_rows: + await Message.create({"thread_id": request.thread_id, "role": "ai", "run_id": run_id, **enrich(node), **row}) + if reply: + await Message.create( + { + "thread_id": request.thread_id, + "role": "ai", + "type": "text", + "data": {"text": reply}, + "data_type": "string", + "run_id": run_id, + **enrich(reply_node), + } + ) + + return StreamingResponse(generate(), media_type="text/event-stream") diff --git a/example/agents/tests/features/job_search.json b/example/agents/tests/features/job_search.json deleted file mode 100644 index d102026f..00000000 --- a/example/agents/tests/features/job_search.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "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/record_no_stream.json b/example/agents/tests/features/record_no_stream.json deleted file mode 100644 index 369f84ea..00000000 --- a/example/agents/tests/features/record_no_stream.json +++ /dev/null @@ -1,6 +0,0 @@ -{ - "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/units/agents/basic_job_search.json b/example/agents/tests/units/agents/basic_job_search.json index 806081fc..3ae16aa6 100644 --- a/example/agents/tests/units/agents/basic_job_search.json +++ b/example/agents/tests/units/agents/basic_job_search.json @@ -5,13 +5,13 @@ "type": "human" }, { - "response_time": 3.6957909469492733, + "response_time": 496.14816601388156, "tool_calls": [ { "args": { "query": "python developer" }, - "id": "call_basic", + "id": "91311f6e-e343-4bbe-be96-45a0395a184a", "name": "job_search_tool", "type": "tool_call" } @@ -19,26 +19,26 @@ "type": "ai", "uses": { "cache_token": 0, - "input_token": 117, - "output_token": 22, - "total_token": 139 + "input_token": 49, + "output_token": 18, + "total_token": 67 } }, { - "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content": "[{\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", "content_type": "json", - "response_time": 0.7865419611334801, + "response_time": 1.1477499501779675, "type": "tool_response" }, { - "content": "I found one opening for a Frontend Developer at Startup Inc, which is a remote, full-time position.", - "response_time": 2.021583030000329, + "content": "I found a Python Developer position at Startup Inc. It is a full-time role and is remote.", + "response_time": 388.0674580577761, "type": "ai", "uses": { "cache_token": 0, - "input_token": 140, - "output_token": 22, - "total_token": 162 + "input_token": 117, + "output_token": 21, + "total_token": 138 } } ] diff --git a/example/agents/tests/units/agents/record_sales.json b/example/agents/tests/units/agents/record_sales.json index 3edd7d25..6623608b 100644 --- a/example/agents/tests/units/agents/record_sales.json +++ b/example/agents/tests/units/agents/record_sales.json @@ -5,14 +5,14 @@ "type": "human" }, { - "content": "I am a large language model, I don't have plans. Is there anything else I can help you with?", - "response_time": 2.0521670230664313, + "content": "I do not have a plan to offer. I am a large language model, an AI. How can I help you?", + "response_time": 547.4648340605199, "type": "ai", "uses": { "cache_token": 0, "input_token": 51, - "output_token": 24, - "total_token": 75 + "output_token": 25, + "total_token": 76 } } ] diff --git a/example/agents/tests/units/agents/record_stream.json b/example/agents/tests/units/agents/record_stream.json index 6b841eff..807872f1 100644 --- a/example/agents/tests/units/agents/record_stream.json +++ b/example/agents/tests/units/agents/record_stream.json @@ -6,7 +6,7 @@ }, { "content": "Hello! How can I help you today?", - "response_time": 1.286374987103045, + "response_time": 471.8274170299992, "type": "ai", "uses": { "cache_token": 0, @@ -22,13 +22,13 @@ "type": "human" }, { - "response_time": 1.1794579913839698, + "response_time": 581.3951250165701, "tool_calls": [ { "args": { "query": "python developer" }, - "id": "call_router_2", + "id": "eb87f523-9eb7-4b4c-bfe3-2729c5c0d40c", "name": "job_search_tool", "type": "tool_call" } @@ -36,15 +36,15 @@ "type": "ai", "uses": { "cache_token": 0, - "input_token": 117, - "output_token": 24, - "total_token": 141 + "input_token": 68, + "output_token": 18, + "total_token": 86 } }, { - "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content": "[{\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", "content_type": "json", - "response_time": 0.3560830373317003, + "response_time": 1.460916013456881, "type": "tool_response" } ], @@ -54,13 +54,13 @@ "type": "human" }, { - "response_time": 1.1740420013666153, + "response_time": 485.8719159383327, "tool_calls": [ { "args": { "query": "python developer" }, - "id": "call_router_1", + "id": "7c4726be-263f-47eb-ac87-596ec8d1ad21", "name": "job_search_tool", "type": "tool_call" } @@ -68,16 +68,20 @@ "type": "ai", "uses": { "cache_token": 0, - "input_token": 117, - "output_token": 24, - "total_token": 141 + "input_token": 69, + "output_token": 18, + "total_token": 87 } }, { - "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content": "[{\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", "content_type": "json", - "response_time": 0.40095800068229437, + "response_time": 1.8106659408658743, "type": "tool_response" } - ] + ], + "judge:b0e0599f655ed45c79573d23cdfd7c75d33508e07b6418b9377cf425458f9199": { + "passed": true, + "reasoning": "The agent provided a relevant Python Developer job recommendation as expected." + } } \ No newline at end of file diff --git a/example/agents/tests/units/agents/stream_job_search.json b/example/agents/tests/units/agents/stream_job_search.json index 9d9a3eb6..7fd85907 100644 --- a/example/agents/tests/units/agents/stream_job_search.json +++ b/example/agents/tests/units/agents/stream_job_search.json @@ -5,13 +5,13 @@ "type": "human" }, { - "response_time": 1.0, + "response_time": 435.28216693084687, "tool_calls": [ { "args": { "query": "python developer" }, - "id": "call_stream", + "id": "d6dfa6d2-ad11-48ad-a322-89d9dcf28fdd", "name": "job_search_tool", "type": "tool_call" } @@ -19,32 +19,27 @@ "type": "ai", "uses": { "cache_token": 0, - "input_token": 117, - "output_token": 24, - "total_token": 141 + "input_token": 49, + "output_token": 18, + "total_token": 67 } }, { - "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content": "[{\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", "content_type": "json", - "response_time": 0.5, + "response_time": 0.8195829577744007, "type": "tool_response" }, { - "chunks": [ - "I found a job for you!", - " It's a Full-time Frontend Developer position at Startup Inc,", - " located remotely." - ], - "content": "I found a job for you! It's a Full-time Frontend Developer position at Startup Inc, located remotely.", - "response_time": 1.0, + "content": "I found one Python developer job:\n\n* Company: Startup Inc\n* Title: Python Developer\n* Type: Full-time\n* Location: Remote", + "response_time": 476.57812503166497, "type": "ai", "uses": { "cache_token": 0, - "input_token": 140, - "output_token": 24, - "total_token": 164 + "input_token": 117, + "output_token": 31, + "total_token": 148 } } ] -} +} \ No newline at end of file diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py index 0925ca12..1dd3428f 100644 --- a/example/agents/tests/units/agents/test_router_agent.py +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -33,3 +33,13 @@ async def test_the_router_with_initial_messages(self): ) as agent: await agent.prompt("suggest python developer jobs") agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") + + agent.assert_response_time_lt(2) + agent.assert_tokens(lambda x: x.where("input", "<=", 1000).where("output", "<=", 5000)) + + await agent.assert_response_judged( + provider="google", + model="gemini-3.5-flash", + expectation="As use is asking for the python developer jobs, the agent " + "should suggest relevant job recommendations.", + ) diff --git a/example/agents/tinker.py b/example/agents/tinker.py index d2fa09f4..6d95e2b0 100644 --- a/example/agents/tinker.py +++ b/example/agents/tinker.py @@ -1,62 +1,5 @@ -import asyncio -import operator -from collections.abc import Callable -from typing import Annotated, TypedDict +from app.agents.job_search_graph import build -from dumpdie import dump -from langchain.agents import create_agent -from langchain_core.messages import AnyMessage -from langgraph.checkpoint.memory import InMemorySaver +graph = build().compile() -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"} -# ) +print(graph.get_graph().draw_mermaid()) diff --git a/example/agents/uv.lock b/example/agents/uv.lock index 302e78f4..23b8151e 100644 --- a/example/agents/uv.lock +++ b/example/agents/uv.lock @@ -47,6 +47,46 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/42/e921fccf5015463e32a3cf6ee7f980a6ed0f395ceeaa45060b61d86486c2/anyio-4.13.0-py3-none-any.whl", hash = "sha256:08b310f9e24a9594186fd75b4f73f4a4152069e3853f1ed8bfbf58369f4ad708", size = 114353, upload-time = "2026-03-24T12:59:08.246Z" }, ] +[[package]] +name = "asyncpg" +version = "0.31.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/fe/cc/d18065ce2380d80b1bcce927c24a2642efd38918e33fd724bc4bca904877/asyncpg-0.31.0.tar.gz", hash = "sha256:c989386c83940bfbd787180f2b1519415e2d3d6277a70d9d0f0145ac73500735", size = 993667, upload-time = "2025-11-24T23:27:00.812Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/a6/59d0a146e61d20e18db7396583242e32e0f120693b67a8de43f1557033e2/asyncpg-0.31.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b44c31e1efc1c15188ef183f287c728e2046abb1d26af4d20858215d50d91fad", size = 662042, upload-time = "2025-11-24T23:25:49.578Z" }, + { url = "https://files.pythonhosted.org/packages/36/01/ffaa189dcb63a2471720615e60185c3f6327716fdc0fc04334436fbb7c65/asyncpg-0.31.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0c89ccf741c067614c9b5fc7f1fc6f3b61ab05ae4aaa966e6fd6b93097c7d20d", size = 638504, upload-time = "2025-11-24T23:25:51.501Z" }, + { url = "https://files.pythonhosted.org/packages/9f/62/3f699ba45d8bd24c5d65392190d19656d74ff0185f42e19d0bbd973bb371/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:12b3b2e39dc5470abd5e98c8d3373e4b1d1234d9fbdedf538798b2c13c64460a", size = 3426241, upload-time = "2025-11-24T23:25:53.278Z" }, + { url = "https://files.pythonhosted.org/packages/8c/d1/a867c2150f9c6e7af6462637f613ba67f78a314b00db220cd26ff559d532/asyncpg-0.31.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:aad7a33913fb8bcb5454313377cc330fbb19a0cd5faa7272407d8a0c4257b671", size = 3520321, upload-time = "2025-11-24T23:25:54.982Z" }, + { url = "https://files.pythonhosted.org/packages/7a/1a/cce4c3f246805ecd285a3591222a2611141f1669d002163abef999b60f98/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:3df118d94f46d85b2e434fd62c84cb66d5834d5a890725fe625f498e72e4d5ec", size = 3316685, upload-time = "2025-11-24T23:25:57.43Z" }, + { url = "https://files.pythonhosted.org/packages/40/ae/0fc961179e78cc579e138fad6eb580448ecae64908f95b8cb8ee2f241f67/asyncpg-0.31.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:bd5b6efff3c17c3202d4b37189969acf8927438a238c6257f66be3c426beba20", size = 3471858, upload-time = "2025-11-24T23:25:59.636Z" }, + { url = "https://files.pythonhosted.org/packages/52/b2/b20e09670be031afa4cbfabd645caece7f85ec62d69c312239de568e058e/asyncpg-0.31.0-cp312-cp312-win32.whl", hash = "sha256:027eaa61361ec735926566f995d959ade4796f6a49d3bde17e5134b9964f9ba8", size = 527852, upload-time = "2025-11-24T23:26:01.084Z" }, + { url = "https://files.pythonhosted.org/packages/b5/f0/f2ed1de154e15b107dc692262395b3c17fc34eafe2a78fc2115931561730/asyncpg-0.31.0-cp312-cp312-win_amd64.whl", hash = "sha256:72d6bdcbc93d608a1158f17932de2321f68b1a967a13e014998db87a72ed3186", size = 597175, upload-time = "2025-11-24T23:26:02.564Z" }, + { url = "https://files.pythonhosted.org/packages/95/11/97b5c2af72a5d0b9bc3fa30cd4b9ce22284a9a943a150fdc768763caf035/asyncpg-0.31.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c204fab1b91e08b0f47e90a75d1b3c62174dab21f670ad6c5d0f243a228f015b", size = 661111, upload-time = "2025-11-24T23:26:04.467Z" }, + { url = "https://files.pythonhosted.org/packages/1b/71/157d611c791a5e2d0423f09f027bd499935f0906e0c2a416ce712ba51ef3/asyncpg-0.31.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:54a64f91839ba59008eccf7aad2e93d6e3de688d796f35803235ea1c4898ae1e", size = 636928, upload-time = "2025-11-24T23:26:05.944Z" }, + { url = "https://files.pythonhosted.org/packages/2e/fc/9e3486fb2bbe69d4a867c0b76d68542650a7ff1574ca40e84c3111bb0c6e/asyncpg-0.31.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0e0822b1038dc7253b337b0f3f676cadc4ac31b126c5d42691c39691962e403", size = 3424067, upload-time = "2025-11-24T23:26:07.957Z" }, + { url = "https://files.pythonhosted.org/packages/12/c6/8c9d076f73f07f995013c791e018a1cd5f31823c2a3187fc8581706aa00f/asyncpg-0.31.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bef056aa502ee34204c161c72ca1f3c274917596877f825968368b2c33f585f4", size = 3518156, upload-time = "2025-11-24T23:26:09.591Z" }, + { url = "https://files.pythonhosted.org/packages/ae/3b/60683a0baf50fbc546499cfb53132cb6835b92b529a05f6a81471ab60d0c/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:0bfbcc5b7ffcd9b75ab1558f00db2ae07db9c80637ad1b2469c43df79d7a5ae2", size = 3319636, upload-time = "2025-11-24T23:26:11.168Z" }, + { url = "https://files.pythonhosted.org/packages/50/dc/8487df0f69bd398a61e1792b3cba0e47477f214eff085ba0efa7eac9ce87/asyncpg-0.31.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:22bc525ebbdc24d1261ecbf6f504998244d4e3be1721784b5f64664d61fbe602", size = 3472079, upload-time = "2025-11-24T23:26:13.164Z" }, + { url = "https://files.pythonhosted.org/packages/13/a1/c5bbeeb8531c05c89135cb8b28575ac2fac618bcb60119ee9696c3faf71c/asyncpg-0.31.0-cp313-cp313-win32.whl", hash = "sha256:f890de5e1e4f7e14023619399a471ce4b71f5418cd67a51853b9910fdfa73696", size = 527606, upload-time = "2025-11-24T23:26:14.78Z" }, + { url = "https://files.pythonhosted.org/packages/91/66/b25ccb84a246b470eb943b0107c07edcae51804912b824054b3413995a10/asyncpg-0.31.0-cp313-cp313-win_amd64.whl", hash = "sha256:dc5f2fa9916f292e5c5c8b2ac2813763bcd7f58e130055b4ad8a0531314201ab", size = 596569, upload-time = "2025-11-24T23:26:16.189Z" }, + { url = "https://files.pythonhosted.org/packages/3c/36/e9450d62e84a13aea6580c83a47a437f26c7ca6fa0f0fd40b6670793ea30/asyncpg-0.31.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:f6b56b91bb0ffc328c4e3ed113136cddd9deefdf5f79ab448598b9772831df44", size = 660867, upload-time = "2025-11-24T23:26:17.631Z" }, + { url = "https://files.pythonhosted.org/packages/82/4b/1d0a2b33b3102d210439338e1beea616a6122267c0df459ff0265cd5807a/asyncpg-0.31.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:334dec28cf20d7f5bb9e45b39546ddf247f8042a690bff9b9573d00086e69cb5", size = 638349, upload-time = "2025-11-24T23:26:19.689Z" }, + { url = "https://files.pythonhosted.org/packages/41/aa/e7f7ac9a7974f08eff9183e392b2d62516f90412686532d27e196c0f0eeb/asyncpg-0.31.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:98cc158c53f46de7bb677fd20c417e264fc02b36d901cc2a43bd6cb0dc6dbfd2", size = 3410428, upload-time = "2025-11-24T23:26:21.275Z" }, + { url = "https://files.pythonhosted.org/packages/6f/de/bf1b60de3dede5c2731e6788617a512bc0ebd9693eac297ee74086f101d7/asyncpg-0.31.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9322b563e2661a52e3cdbc93eed3be7748b289f792e0011cb2720d278b366ce2", size = 3471678, upload-time = "2025-11-24T23:26:23.627Z" }, + { url = "https://files.pythonhosted.org/packages/46/78/fc3ade003e22d8bd53aaf8f75f4be48f0b460fa73738f0391b9c856a9147/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19857a358fc811d82227449b7ca40afb46e75b33eb8897240c3839dd8b744218", size = 3313505, upload-time = "2025-11-24T23:26:25.235Z" }, + { url = "https://files.pythonhosted.org/packages/bf/e9/73eb8a6789e927816f4705291be21f2225687bfa97321e40cd23055e903a/asyncpg-0.31.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:ba5f8886e850882ff2c2ace5732300e99193823e8107e2c53ef01c1ebfa1e85d", size = 3434744, upload-time = "2025-11-24T23:26:26.944Z" }, + { url = "https://files.pythonhosted.org/packages/08/4b/f10b880534413c65c5b5862f79b8e81553a8f364e5238832ad4c0af71b7f/asyncpg-0.31.0-cp314-cp314-win32.whl", hash = "sha256:cea3a0b2a14f95834cee29432e4ddc399b95700eb1d51bbc5bfee8f31fa07b2b", size = 532251, upload-time = "2025-11-24T23:26:28.404Z" }, + { url = "https://files.pythonhosted.org/packages/d3/2d/7aa40750b7a19efa5d66e67fc06008ca0f27ba1bd082e457ad82f59aba49/asyncpg-0.31.0-cp314-cp314-win_amd64.whl", hash = "sha256:04d19392716af6b029411a0264d92093b6e5e8285ae97a39957b9a9c14ea72be", size = 604901, upload-time = "2025-11-24T23:26:30.34Z" }, + { url = "https://files.pythonhosted.org/packages/ce/fe/b9dfe349b83b9dee28cc42360d2c86b2cdce4cb551a2c2d27e156bcac84d/asyncpg-0.31.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:bdb957706da132e982cc6856bb2f7b740603472b54c3ebc77fe60ea3e57e1bd2", size = 702280, upload-time = "2025-11-24T23:26:32Z" }, + { url = "https://files.pythonhosted.org/packages/6a/81/e6be6e37e560bd91e6c23ea8a6138a04fd057b08cf63d3c5055c98e81c1d/asyncpg-0.31.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:6d11b198111a72f47154fa03b85799f9be63701e068b43f84ac25da0bda9cb31", size = 682931, upload-time = "2025-11-24T23:26:33.572Z" }, + { url = "https://files.pythonhosted.org/packages/a6/45/6009040da85a1648dd5bc75b3b0a062081c483e75a1a29041ae63a0bf0dc/asyncpg-0.31.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:18c83b03bc0d1b23e6230f5bf8d4f217dc9bc08644ce0502a9d91dc9e634a9c7", size = 3581608, upload-time = "2025-11-24T23:26:35.638Z" }, + { url = "https://files.pythonhosted.org/packages/7e/06/2e3d4d7608b0b2b3adbee0d0bd6a2d29ca0fc4d8a78f8277df04e2d1fd7b/asyncpg-0.31.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:e009abc333464ff18b8f6fd146addffd9aaf63e79aa3bb40ab7a4c332d0c5e9e", size = 3498738, upload-time = "2025-11-24T23:26:37.275Z" }, + { url = "https://files.pythonhosted.org/packages/7d/aa/7d75ede780033141c51d83577ea23236ba7d3a23593929b32b49db8ed36e/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3b1fbcb0e396a5ca435a8826a87e5c2c2cc0c8c68eb6fadf82168056b0e53a8c", size = 3401026, upload-time = "2025-11-24T23:26:39.423Z" }, + { url = "https://files.pythonhosted.org/packages/ba/7a/15e37d45e7f7c94facc1e9148c0e455e8f33c08f0b8a0b1deb2c5171771b/asyncpg-0.31.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:8df714dba348efcc162d2adf02d213e5fab1bd9f557e1305633e851a61814a7a", size = 3429426, upload-time = "2025-11-24T23:26:41.032Z" }, + { url = "https://files.pythonhosted.org/packages/13/d5/71437c5f6ae5f307828710efbe62163974e71237d5d46ebd2869ea052d10/asyncpg-0.31.0-cp314-cp314t-win32.whl", hash = "sha256:1b41f1afb1033f2b44f3234993b15096ddc9cd71b21a42dbd87fc6a57b43d65d", size = 614495, upload-time = "2025-11-24T23:26:42.659Z" }, + { url = "https://files.pythonhosted.org/packages/3c/d7/8fb3044eaef08a310acfe23dae9a8e2e07d305edc29a53497e52bc76eca7/asyncpg-0.31.0-cp314-cp314t-win_amd64.whl", hash = "sha256:bd4107bb7cdd0e9e65fae66a62afd3a249663b844fa34d479f6d5b3bef9c04c3", size = 706062, upload-time = "2025-11-24T23:26:44.086Z" }, +] + [[package]] name = "certifi" version = "2026.4.22" @@ -497,6 +537,7 @@ name = "fastapi-startkit-application" version = "0.1.0" source = { virtual = "." } dependencies = [ + { name = "asyncpg" }, { name = "cleo" }, { name = "fastapi-startkit", extra = ["ai", "database", "fastapi", "sqlite"] }, { name = "langchain" }, @@ -514,6 +555,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "asyncpg", specifier = ">=0.31.0" }, { name = "cleo", specifier = ">=2.1.0" }, { name = "fastapi-startkit", extras = ["ai", "database", "fastapi", "sqlite"], editable = "../../fastapi_startkit" }, { name = "langchain", specifier = ">=1.3.10" }, diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 71c93b49..d46ba94a 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -20,6 +20,9 @@ class Agent: max_tokens: int = 4096 timeout: float = 30.0 top_p: float = 1.0 + # Forwarded to bind_tools() when set: "auto" (default), "any" (must call a tool), + # or a specific tool name. + tool_choice: str | None = None def messages(self) -> list[dict]: return [] @@ -55,7 +58,9 @@ async def stream( *, model: str | None = None, provider_options: dict | None = None, - ) -> AsyncIterator[str]: + ) -> AsyncIterator[dict]: + """Yields typed frames: {"type": "delta", "text": ...} for model tokens and + {"type": "tool_response", "name": ..., "content": ...} for tool results.""" async for chunk in self.runner().stream(message, model=model, provider_options=provider_options): yield chunk diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py index 28b76345..13d8e335 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/ai.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -82,7 +82,11 @@ def build( chat_model = init_chat_model(self._resolve_model(agent, model), **kwargs) tools = list(agent.tools()) - chat_model = chat_model.bind_tools(tools) if tools else chat_model + if tools: + if agent.tool_choice: + chat_model = chat_model.bind_tools(tools, tool_choice=agent.tool_choice) + else: + chat_model = chat_model.bind_tools(tools) schema = agent.schema() if schema is not None: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/graph.py b/fastapi_startkit/src/fastapi_startkit/ai/graph.py index 514fad33..0c6b1662 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/graph.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/graph.py @@ -5,6 +5,7 @@ from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Annotated, Any, TypedDict +from langchain_core.messages import ToolMessage from langchain_core.runnables import RunnableConfig from langgraph.graph import END from langgraph.graph.message import add_messages @@ -56,7 +57,7 @@ async def stream( model: str | None = None, config: RunnableConfig | dict | None = None, provider_options: dict | None = None, - ) -> AsyncIterator[str]: + ) -> AsyncIterator[dict]: async for chunk in self.runner().stream(message, model=model, config=config, provider_options=provider_options): yield chunk @@ -132,7 +133,7 @@ async def run( started = time.perf_counter() self._reset_capture() compiled = await self._compile(model, provider_options) - state = {"messages": self._build_messages(message, attachments), "llm_calls": 0} + state = {"messages": await self._build_messages(message, attachments), "llm_calls": 0} result = await compiled.ainvoke(state, config=self._merge_config(config)) response = self._apply_schema(self._to_response(result)) response.runtime = time.perf_counter() - started @@ -145,10 +146,10 @@ async def stream( model: str | None = None, config: RunnableConfig | dict | None = None, provider_options: dict | None = None, - ) -> AsyncIterator[str]: + ) -> AsyncIterator[dict]: self._reset_capture() compiled = await self._compile(model, provider_options) - state = {"messages": self._build_messages(message), "llm_calls": 0} + state = {"messages": await self._build_messages(message), "llm_calls": 0} final_state: dict = {} # "messages" streams tokens; "values" hands back the full state after each # step, whose last snapshot carries every message (and its tool_calls). @@ -160,8 +161,12 @@ async def stream( continue message_chunk, _meta = chunk text = message_chunk.content if isinstance(message_chunk.content, str) else str(message_chunk.content) - if text: - yield text + if not text: + continue + if isinstance(message_chunk, ToolMessage): + yield {"type": "tool_response", "name": message_chunk.name or "", "content": text} + else: + yield {"type": "delta", "text": text} self.last_response = self._to_response(final_state) if final_state else AgentResponse(content="") def _to_response(self, result: dict) -> AgentResponse: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py b/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py index 38fd8dc9..107e0369 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py @@ -40,13 +40,22 @@ async def _finish(self, final: Any) -> Any: await result return final + @staticmethod + def _merge(accumulated: Any, chunk: Any) -> Any: + # Stream frames are dicts ({"type": "delta"/"tool_response", ...}); the + # after-hooks want the joined text, not a dict sum. + if isinstance(chunk, dict): + text = chunk.get("text") or chunk.get("content") or "" + return text if accumulated is None else accumulated + text + return chunk if accumulated is None else accumulated + chunk + def __aiter__(self) -> AsyncIterator: async def _it() -> AsyncIterator: accumulated = None finished = False try: async for chunk in self._source(): - accumulated = chunk if accumulated is None else accumulated + chunk + accumulated = self._merge(accumulated, chunk) yield chunk finished = True await self._finish(accumulated) @@ -63,7 +72,7 @@ def __await__(self): async def _run() -> Any: accumulated = None async for chunk in self._source(): - accumulated = chunk if accumulated is None else accumulated + chunk + accumulated = self._merge(accumulated, chunk) return await self._finish(accumulated) return _run().__await__() diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index 2d846f19..a62604ca 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -1,5 +1,6 @@ from __future__ import annotations +import inspect import time from abc import ABC, abstractmethod from collections.abc import AsyncIterator @@ -77,7 +78,7 @@ def _record_tool_message(self, call: ToolCall, message: BaseMessage, response_ti } ) - def _build_messages(self, message: str, attachments: list[Document] | None = None) -> list[dict]: + async def _build_messages(self, message: str, attachments: list[Document] | None = None) -> list[dict]: agent = self.agent messages: list[dict] = [] @@ -85,7 +86,12 @@ def _build_messages(self, message: str, attachments: list[Document] | None = Non if instruction: messages.append({"role": "system", "content": instruction}) - messages.extend(agent.messages() or []) + # messages() may be sync (return a list) or async (agents that read + # their history from a database); support both. + history = agent.messages() + if inspect.isawaitable(history): + history = await history + messages.extend(history or []) if message: messages.append({"role": "user", "content": message}) @@ -136,7 +142,7 @@ def stream( *, model: str | None = None, provider_options: dict | None = None, - ) -> AsyncIterator[str]: ... + ) -> AsyncIterator[dict]: ... class Runner(BaseRunner): @@ -150,7 +156,7 @@ async def run( ) -> AgentResponse: started = time.perf_counter() self._reset_capture() - messages = self._build_messages(message, attachments) + messages = await self._build_messages(message, attachments) model = self._build_model(model, provider_options) response = await self._run_pipeline(model, messages) @@ -215,9 +221,9 @@ async def stream( *, model: str | None = None, provider_options: dict | None = None, - ) -> AsyncIterator[str]: + ) -> AsyncIterator[dict]: self._reset_capture() - messages = self._build_messages(message) + messages = await self._build_messages(message) chat_model = self._build_model(model, provider_options) chain = list(self.agent.middleware()) @@ -245,7 +251,7 @@ async def _invoke(self, model: Runnable[Any, BaseMessage], messages: list) -> Ba return response return (await self._run_tools(response.tool_calls))[-1] - async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list) -> AsyncIterator[str]: + async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list) -> AsyncIterator[dict]: started = time.perf_counter() gathered: AIMessageChunk | None = None chunks: list[str] = [] @@ -253,7 +259,7 @@ async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list if chunk.content: text = _as_text(chunk.content) chunks.append(text) - yield text + yield {"type": "delta", "text": text} gathered = chunk if gathered is None else gathered + chunk # type: ignore[operator] latency_ms = (time.perf_counter() - started) * 1000 @@ -262,8 +268,8 @@ async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list self._record_ai_message(gathered, latency_ms, chunks=chunks) if not gathered.tool_calls: return - for message in await self._run_tools(gathered.tool_calls): - yield _as_text(message.content) + for call, message in zip(gathered.tool_calls, await self._run_tools(gathered.tool_calls)): + yield {"type": "tool_response", "name": call["name"], "content": _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()} diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index 46918704..cbd4a634 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -22,6 +22,13 @@ def _joined(value: Any) -> str: return "".join(value) if isinstance(value, list) else value +def _frame_text(frame: Any) -> str: + """Text carried by a stream frame; cassettes keep storing plain strings.""" + if isinstance(frame, dict): + return frame.get("text") or frame.get("content") or "" + return str(frame) + + def _new_token_totals() -> dict: return {"input": 0, "output": 0, "cache": 0, "total": 0} @@ -74,6 +81,7 @@ def __init__(self, agent_cls: type[Agent], responses: list) -> None: self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None self._tokens: dict = _new_token_totals() + self._response_time: float = 0.0 def _history(self) -> list: return self._records @@ -104,20 +112,21 @@ async def prompt( self._remember(message, self._last_response) return self._last_response - async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[str]: + async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[dict]: chunks: list[str] = [] extra = {"config": config} if config is not None else {} # Drive the runner directly so we can read the structured response it # captures while streaming (content + tool_calls), not just joined text. runner = self._agent.runner() - async for chunk in runner.stream(message, **extra): - chunks.append(chunk) - yield chunk + async for frame in runner.stream(message, **extra): + chunks.append(_frame_text(frame)) + yield frame self._last_response = getattr(runner, "last_response", None) or AgentResponse(content="".join(chunks)) self._remember(message, self._last_response) def _remember(self, message: str, response: AgentResponse) -> None: self._accumulate_tokens(response) + self._response_time += response.runtime self._records.append({"role": "user", "content": message}) self._records.append({"role": "assistant", "content": response.content}) @@ -174,6 +183,7 @@ def reset(self) -> "AgentFake": self._records.clear() self._last_response = None self._tokens = _new_token_totals() + self._response_time = 0.0 return self def _require_response(self) -> AgentResponse: @@ -201,8 +211,15 @@ def assert_tool_not_called(self, names: list[str]) -> None: 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 on the response time accumulated across every prompt()/stream() so far. + + The time is read from the recorded transcript (each ai/tool step's + ``response_time``), not the wall-clock duration of the call — so it stays + correct on cassette replay, where the cache read is effectively instant. + """ 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" + total = self._response_time + assert total < seconds, f"Expected total response time < {seconds}s, took {total:.3f}s" async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None: response = self._require_response() @@ -262,6 +279,7 @@ def __init__(self, real: Agent, cassette: str | None = None, messages: list | No self._last_response: AgentResponse | None = None self.last_elapsed: float | None = None self._tokens: dict = _new_token_totals() + self._response_time: float = 0.0 def _history(self) -> list: return self._seed_messages + self._records @@ -332,6 +350,7 @@ def _response_from_cache(value: Any) -> AgentResponse: def _remember_turn(self, message: str, response: AgentResponse) -> None: self._accumulate_tokens(response) + self._response_time += response.runtime self._records.append({"role": "user", "content": message}) turn: dict[str, Any] = {"role": "assistant", "content": response.content} if response.tool_calls: @@ -356,7 +375,7 @@ async def prompt( self._remember_turn(message, response) return response - async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[str]: + async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[dict]: from . import recording # noqa: PLC0415 cassette, store = self._load() @@ -373,16 +392,16 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter chunks = value if isinstance(value, list) else [value] response = AgentResponse(content=_joined(chunks)) for chunk in chunks: - yield chunk + yield {"type": "delta", "text": chunk} else: extra = {"config": config} if config is not None else {} # Drive the runner directly to capture the structured response # (content + tool_calls + chunks) it records while streaming. runner = self._real.runner() chunks = [] - async for chunk in runner.stream(message, **extra): - chunks.append(chunk) - yield chunk + async for frame in runner.stream(message, **extra): + chunks.append(_frame_text(frame)) + yield frame response = getattr(runner, "last_response", None) or AgentResponse(content=_joined(chunks)) self._save(cassette, store, key, self._turn(message, response, chunks=chunks)) self._last_response = response diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 81596be1..9c77ea1a 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -116,9 +116,10 @@ class GoogleAgent(Agent): async def test_stream_yields_tokens_from_the_model(self): self.setup_agent([AIMessage(content="streamed reply")]) - chunks = [chunk async for chunk in Agent().stream("hello")] + frames = [frame async for frame in Agent().stream("hello")] - self.assertEqual("".join(chunks), "streamed reply") + self.assertTrue(all(frame["type"] == "delta" for frame in frames)) + self.assertEqual("".join(frame["text"] for frame in frames), "streamed reply") async def test_middleware_streams_token_by_token_and_runs_after_hook(self): events: list = [] @@ -134,7 +135,7 @@ def middleware(self): self.setup_agent([AIMessage(content="one two three")], LoggedAgent) - chunks = [chunk async for chunk in LoggedAgent().stream("hi")] + chunks = [frame["text"] async for frame in LoggedAgent().stream("hi")] # Middleware must not buffer: the model's tokens arrive as separate chunks... self.assertEqual("".join(chunks), "one two three") @@ -174,9 +175,12 @@ async def test_stream_yields_tool_result_without_calling_model_again(self): ) self.addCleanup(Ai.reset_fakes) - chunks = [chunk async for chunk in JobAssistant().stream("find me a python job")] + frames = [frame async for frame in JobAssistant().stream("find me a python job")] - self.assertEqual(chunks, ["Python Developer at Shopify"]) + self.assertEqual( + frames, + [{"type": "tool_response", "name": "search_jobs", "content": "Python Developer at Shopify"}], + ) def test_resolve_model_falls_back_to_lab_default(self): self.assertEqual(Ai()._resolve_model(Agent()), "gemini-2.5-flash-lite") @@ -189,40 +193,40 @@ class AnthropicAgent(Agent): 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 = Runner(JobAssistant())._build_messages("find me a job") + async def test_instructions_lead_the_message_list(self): + messages = await 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) - def test_instructions_can_be_a_method_override(self): + async def test_instructions_can_be_a_method_override(self): class DynamicAgent(Agent): def instructions(self) -> str: return "Computed identity." - messages = Runner(DynamicAgent())._build_messages("hi") + messages = await Runner(DynamicAgent())._build_messages("hi") self.assertEqual(messages[0], {"role": "system", "content": "Computed identity."}) - def test_no_instructions_prepends_no_system_message(self): - messages = Runner(Agent())._build_messages("hi") + async def test_no_instructions_prepends_no_system_message(self): + messages = await Runner(Agent())._build_messages("hi") self.assertTrue(all(m.get("role") != "system" for m in messages)) - def test_build_messages_inlines_text_attachment(self): + async def test_build_messages_inlines_text_attachment(self): doc = Document(content="Q3 revenue was $1.2M.", name="q3-report.txt") - messages = Runner(Agent())._build_messages("Summarise this report.", attachments=[doc]) + messages = await 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."}) self.assertEqual(user_content[1]["type"], "text") self.assertIn("q3-report.txt", user_content[1]["text"]) - def test_build_messages_encodes_binary_attachment_as_file_block(self): + async 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 = Runner(Agent())._build_messages("Summarise", attachments=[doc]) + messages = await 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 dfd3e56f..4ea8d6c5 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -135,7 +135,7 @@ async def test_fake_rebinding_overrides_previous(self): async def test_stream_returns_fake_response(self): with SimpleAgent.fake(["Faked stream!"]) as agent: - chunks = [chunk async for chunk in agent.stream("hello world")] + chunks = [frame["text"] async for frame in agent.stream("hello world")] self.assertEqual("".join(chunks), "Faked stream!") self.assertGreater(len(chunks), 1) @@ -143,7 +143,7 @@ async def test_stream_returns_fake_response(self): async def test_stream_replays_the_registered_text_exactly(self): with SimpleAgent.fake(["Hello there, friend"]) as agent: - chunks = [chunk async for chunk in agent.stream("hi")] + chunks = [frame["text"] async for frame in agent.stream("hi")] self.assertEqual("".join(chunks), "Hello there, friend") @@ -236,7 +236,7 @@ def setup_stream(self, chunks): async def fake_stream(runner_self, message, **kwargs): calls.append(message) for chunk in chunks: - yield chunk + yield {"type": "delta", "text": chunk} # The fluent fakes drive the runner's stream() directly (so they can read # the structured response it captures), so mock at that layer. @@ -250,7 +250,7 @@ async def test_stream_first_run_records_chunk_list_to_cassette(self): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "s.json") with SimpleAgent.record(cassette) as agent: - chunks = [c async for c in agent.stream("hi")] + chunks = [frame["text"] async for frame in agent.stream("hi")] self.assertEqual(chunks, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) @@ -268,7 +268,7 @@ async def test_stream_second_run_replays_chunks_without_calling_stream(self): 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")] + replayed = [frame["text"] async for frame in agent.stream("hi")] self.assertEqual(replayed, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) # real stream invoked only on the first run diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py index 1c0d08fa..513a6aa5 100644 --- a/fastapi_startkit/tests/ai/test_agent_record_fluent.py +++ b/fastapi_startkit/tests/ai/test_agent_record_fluent.py @@ -213,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 = Runner(agent._real)._build_messages("suggest python developer jobs") + built = await Runner(agent._real)._build_messages("suggest python developer jobs") self.assertEqual(built[0], seed[0]) self.assertEqual(built[1], seed[1]) diff --git a/fastapi_startkit/tests/ai/test_assert_response_time.py b/fastapi_startkit/tests/ai/test_assert_response_time.py new file mode 100644 index 00000000..5b988e95 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_assert_response_time.py @@ -0,0 +1,106 @@ +"""Tests for agent.assert_response_time_lt() — assert on recorded response time (task #1271). + + agent.assert_response_time_lt(0.1) + +The assertion must read the ``response_time`` recorded on the cassette transcript +(summed across every ai/tool step of every turn), NOT the wall-clock duration of +the replay. On cassette replay the cache read is near-instant, so a slow recorded +turn must still be seen as slow — otherwise a replayed call always looks fast and +the assertion silently passes regardless of what was recorded. +""" + +import os +import tempfile +import unittest +from unittest import mock + +from fastapi_startkit.ai import recording +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.response import AgentResponse + + +class TimingAgent(Agent): + pass + + +def _slow_turn_transcript() -> list[dict]: + """Mirror the reported repro: a slow ai turn (~660ms) plus a tool response + (~1.5ms). Total recorded time ~= 0.6616s.""" + return [ + recording.ai( + content="", + tool_calls=[ + {"args": {"query": "python developer"}, "id": "x", "name": "job_search_tool", "type": "tool_call"} + ], + uses={"input_token": 68, "output_token": 18, "cache_token": 0, "total_token": 86}, + response_time=660.0997089408338, + ), + recording.tool_response( + content='[{"id": 2, "title": "Frontend Developer"}]', response_time=1.4820829965174198 + ), + ] + + +def _fake_prompt(responses: list): + queue = list(responses) + + async def prompt(agent_self, message, **kwargs): + return queue.pop(0) + + return mock.patch.object(TimingAgent, "prompt", prompt) + + +def _record(cassette: str, responses: list, prompts: list[str]) -> None: + """Record a cassette by driving fake live prompts once.""" + + async def _drive(): + with _fake_prompt(responses): + with TimingAgent.record(cassette) as agent: + for message in prompts: + await agent.prompt(message) + + return _drive() + + +def _slow_response() -> AgentResponse: + return recording.to_response(_slow_turn_transcript()) + + +class TestAssertResponseTime(unittest.IsolatedAsyncioTestCase): + async def test_fails_when_recorded_time_exceeds_threshold_on_replay(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + await _record(cassette, [_slow_response()], ["suggest python developer jobs"]) + + # Replay: cache read is near-instant, but the recorded ~0.66s must win. + with TimingAgent.record(cassette) as agent: + await agent.prompt("suggest python developer jobs") + with self.assertRaises(AssertionError): + agent.assert_response_time_lt(0.1) + + async def test_passes_when_recorded_time_under_threshold_on_replay(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + await _record(cassette, [_slow_response()], ["suggest python developer jobs"]) + + with TimingAgent.record(cassette) as agent: + await agent.prompt("suggest python developer jobs") + agent.assert_response_time_lt(2.0) # 0.6616 < 2.0 + + async def test_accumulates_recorded_time_across_turns(self): + responses = [_slow_response(), _slow_response()] + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + await _record(cassette, responses, ["first", "second"]) + + with TimingAgent.record(cassette) as agent: + await agent.prompt("first") + await agent.prompt("second") + # Two turns ~0.6616s each -> ~1.32s total. + agent.assert_response_time_lt(2.0) + with self.assertRaises(AssertionError): + agent.assert_response_time_lt(1.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/fastapi_startkit/tests/ai/test_graph_agent.py b/fastapi_startkit/tests/ai/test_graph_agent.py index d83803cc..236fd1f1 100644 --- a/fastapi_startkit/tests/ai/test_graph_agent.py +++ b/fastapi_startkit/tests/ai/test_graph_agent.py @@ -66,7 +66,7 @@ async def test_prompt_executes_the_tool_and_feeds_it_back(self): 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?")] + chunks = [frame["text"] async for frame in MathAgent().stream("What is 3 plus 4?")] self.assertEqual("".join(chunks), "The answer is 7") self.assertGreater(len(chunks), 1) From d3f39ac8e4f5afb39bdf702063f5c44dabc4a4f5 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Wed, 29 Jul 2026 15:04:32 -0700 Subject: [PATCH 13/13] feat(ai)!: agents speak the LangChain create_agent contract MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Delete AgentResponse/AgentSnapshot and make the ai package's agents return and stream exactly what langchain's create_agent produces: - prompt() returns the state dict {"messages": [HumanMessage, AIMessage, ToolMessage, ...]}, plus "structured_response" for schema agents. AI/tool messages carry their response_time (ms) in additional_kwargs. - stream() yields astream_events-shaped StandardStreamEvent dicts (on_chat_model_start/stream/end, on_tool_start/end); GraphRunner.stream delegates to the compiled graph's astream_events verbatim. - New fastapi_startkit.ai.state helpers (text, tool_calls, tool_events, usage, runtime, structured) replace the response wrapper's accessors. - Testing harness (AgentFake/AgentRecordFake) operates on states; every assert_* keeps its API. Cassette on-disk format unchanged — old cassettes (including legacy flat shapes) still replay. - usage now sums every model call in a turn instead of only the last. - Example app (remember mixin, router/search/summarizer graph nodes, SSE routes) migrated to the state helpers; frontend SSE frames unchanged via a route-level event→frame mapping. Co-Authored-By: Claude Fable 5 --- example/agents/app/agents/agent.py | 74 ++++++- example/agents/app/agents/job_search_graph.py | 48 ++--- example/agents/app/agents/remember.py | 60 ++++++ example/agents/app/tools/job_search_tool.py | 12 +- example/agents/routes/api.py | 140 +++++-------- .../agents/tests/features/test_jobs_stream.py | 120 +++++++++++ .../integrations/agents/job_search_agent.json | 66 ++++++ .../agents/job_summarizer_agent.json | 40 ++++ .../agents/test_job_search_integration.py | 96 +++++++++ example/agents/tests/support/fakes.py | 34 ++++ example/agents/tests/test_case.py | 15 ++ .../units/agents/test_job_search_node.py | 159 +++++++++++++++ example/agents/tinker.py | 31 ++- .../src/fastapi_startkit/ai/__init__.py | 5 +- .../src/fastapi_startkit/ai/agent.py | 12 +- .../src/fastapi_startkit/ai/graph.py | 73 +++---- .../src/fastapi_startkit/ai/judge.py | 4 +- .../src/fastapi_startkit/ai/pipeline.py | 21 +- .../src/fastapi_startkit/ai/recording.py | 134 +++++++------ .../src/fastapi_startkit/ai/response.py | 76 ------- .../src/fastapi_startkit/ai/runner.py | 175 +++++++++------- .../src/fastapi_startkit/ai/state.py | 88 ++++++++ .../src/fastapi_startkit/ai/testing.py | 188 +++++++++--------- fastapi_startkit/tests/ai/test_agent.py | 45 +++-- fastapi_startkit/tests/ai/test_agent_fake.py | 38 ++-- .../tests/ai/test_agent_record_fluent.py | 36 ++-- .../tests/ai/test_agent_response.py | 104 ---------- .../tests/ai/test_agent_schema.py | 21 +- fastapi_startkit/tests/ai/test_agent_state.py | 83 ++++++++ fastapi_startkit/tests/ai/test_ai_fake.py | 5 +- .../tests/ai/test_assert_response_time.py | 5 +- .../tests/ai/test_assert_tokens.py | 38 ++-- fastapi_startkit/tests/ai/test_graph_agent.py | 15 +- .../tests/ai/test_graph_recording.py | 20 +- fastapi_startkit/tests/ai/test_judge_agent.py | 8 +- .../tests/ai/test_record_cassette_format.py | 69 ++++--- fastapi_startkit/tests/ai/test_recording.py | 28 +-- .../tests/ai/test_runner_recording.py | 29 +-- .../tests/ai/test_structured_output.py | 20 +- 39 files changed, 1466 insertions(+), 769 deletions(-) create mode 100644 example/agents/app/agents/remember.py create mode 100644 example/agents/tests/features/test_jobs_stream.py create mode 100644 example/agents/tests/integrations/agents/job_search_agent.json create mode 100644 example/agents/tests/integrations/agents/job_summarizer_agent.json create mode 100644 example/agents/tests/integrations/agents/test_job_search_integration.py create mode 100644 example/agents/tests/support/fakes.py create mode 100644 example/agents/tests/units/agents/test_job_search_node.py delete mode 100644 fastapi_startkit/src/fastapi_startkit/ai/response.py create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/state.py delete mode 100644 fastapi_startkit/tests/ai/test_agent_response.py create mode 100644 fastapi_startkit/tests/ai/test_agent_state.py diff --git a/example/agents/app/agents/agent.py b/example/agents/app/agents/agent.py index d56c43fe..cf8f1aee 100644 --- a/example/agents/app/agents/agent.py +++ b/example/agents/app/agents/agent.py @@ -1,12 +1,13 @@ import json -from fastapi_startkit.ai import Agent, Middleware +from fastapi_startkit.ai import Middleware +from fastapi_startkit.ai import state as ai_state from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage -from langchain_core.runnables import RunnableConfig from langchain_core.tools import BaseTool from app.agents.job_search_graph import USER_PROFILE -from app.agents.state import Context, OverallState +from app.agents.remember import RememberMixin, ThreadAgent +from app.agents.state import Context, RouterOutput from app.middleware.agent_logger import AgentLogger from app.models.message import Message from app.tools.job_search_tool import job_search_tool @@ -19,8 +20,59 @@ preferences like location or remote. Ignore filler words like 'suggest', 'me', 'jobs'. If the latest message has no usable keywords, take them from the user's profile.""" +SUMMARIZER_PROMPT = """Write a short, friendly summary of these job listings for the user. +Group by role type, mention company and location. If the list is empty, say so plainly. +Use only what is given - do not invent jobs.""" -class JobSearchAgent(Agent): +ROUTER_PROMPT = """You route a career assistant's queries. + +- job_search: the user wants job listings / openings / roles. +- company_research: the user asks about a specific company. +- chat: greetings (hi / hello), thanks, small talk, or anything conversational. + +Also decide which extra context the next node needs: +- include_user_profile: the query is vague about role, skills or location, so the user's own profile helps. +- include_last_job_search_response: the user is refining an earlier search. +""" + + +class JobSearchRouterAgent(RememberMixin, ThreadAgent): + """Classifies a query into a RouterOutput via structured output, answering + greetings inline through its `reply` field. Context: only the user input.""" + + model = "gemini-3.1-flash-lite" # match the graph's model, not the lab default + + def instructions(self) -> str: + return ROUTER_PROMPT + + def schema(self) -> type: + return RouterOutput + + def middleware(self) -> list[Middleware]: + return [AgentLogger()] + + async def remember(self, state: dict) -> None: + # The router's output is routing state, except when it answers a greeting + # inline — that reply is conversation, so it goes on the thread. + decision = ai_state.structured(state) + if decision and decision.intent == "chat" and decision.reply: + await self.remember_row("text", {"text": decision.reply}, state) + + +class JobSummarizerAgent(RememberMixin, ThreadAgent): + """Summarizes a job list into prose. Context: only the tool response it is + prompted with — no history, no profile, no tools.""" + + model = "gemini-3.1-flash-lite" # match the graph's model, not the lab default + + def instructions(self) -> str: + return SUMMARIZER_PROMPT + + def middleware(self) -> list[Middleware]: + return [AgentLogger()] + + +class JobSearchAgent(RememberMixin, ThreadAgent): """Plain single-LLM agent (like ChatAgent) that owns its own context: the thread's past user messages plus the user profile, injected only when the router asked for it via `contexts`. messages() is async — the runner awaits it. @@ -30,11 +82,9 @@ class JobSearchAgent(Agent): # by job_search_node (it only reads tool_events). tool_choice = "any" - def __init__(self, state: OverallState, config: RunnableConfig | None = None): - super().__init__() - self.state = state - self.thread_id = ((config or {}).get("configurable") or {}).get("thread_id", "") - self.contexts: list[Context] = state.get("contexts") or [] + @property + def contexts(self) -> list[Context]: + return self.state.get("contexts") or [] async def messages(self) -> list[BaseMessage]: messages: list[BaseMessage] = [] @@ -42,7 +92,11 @@ async def messages(self) -> list[BaseMessage]: if Context.INCLUDE_USER_PROFILE in self.contexts: messages.append(SystemMessage(content=f"User profile: {json.dumps(USER_PROFILE)}")) - rows = await Message.where("thread_id", self.thread_id).where("role", "user").order_by("id").get() + rows = list(await Message.where("thread_id", self.thread_id).where("role", "user").order_by("id").get()) + # The current turn's user row is persisted before the run; the runner appends + # the live query itself, so drop it from history to avoid sending it twice. + if rows and rows[-1].data.get("text") == self.state.get("query"): + rows = rows[:-1] messages.extend(HumanMessage(content=row.data.get("text", "")) for row in rows) if Context.INCLUDE_LAST_JOB_SEARCH_RESPONSE in self.contexts and (last := await self._last_job_search()): diff --git a/example/agents/app/agents/job_search_graph.py b/example/agents/app/agents/job_search_graph.py index 7c217538..5b661536 100644 --- a/example/agents/app/agents/job_search_graph.py +++ b/example/agents/app/agents/job_search_graph.py @@ -1,23 +1,13 @@ import json -from functools import cache from typing import cast -from langchain.chat_models import init_chat_model -from langchain_core.language_models.chat_models import BaseChatModel -from langchain_core.messages import HumanMessage, SystemMessage +from fastapi_startkit.ai import state as ai_state from langchain_core.runnables import RunnableConfig from langgraph.graph import END, START, StateGraph from langgraph.graph.state import CompiledStateGraph from app.agents.state import InputState, OutputState, OverallState, RouterOutput - -@cache -def llm() -> BaseChatModel: - """Built lazily so importing this module doesn't require the API key at import time.""" - return init_chat_model("google_genai:gemini-3.1-flash-lite") - - # Stand-in for a real profile lookup; the router decides when it's needed. USER_PROFILE = { "name": "Alice", @@ -27,26 +17,18 @@ def llm() -> BaseChatModel: } -ROUTER_PROMPT = """You route a career assistant's queries. - -- job_search: the user wants job listings / openings / roles. -- company_research: the user asks about a specific company. -- chat: greetings (hi / hello), thanks, small talk, or anything conversational. - -Also decide which extra context the next node needs: -- include_user_profile: the query is vague about role, skills or location, so the user's own profile helps. -- include_last_job_search_response: the user is refining an earlier search. -""" - - -async def router_node(state: InputState) -> dict: +async def router_node(state: InputState, config: RunnableConfig) -> dict: """Context: only the user input. No history, no profile, no tool output. Handles greetings itself: for the 'chat' intent it answers inline via RouterOutput.reply, so no extra node/LLM call is needed. """ - model = llm().with_structured_output(RouterOutput) - decision = cast(RouterOutput, await model.ainvoke([SystemMessage(ROUTER_PROMPT), HumanMessage(state["query"])])) + from app.agents.agent import JobSearchRouterAgent + + response = await JobSearchRouterAgent(state, config).prompt(state["query"]) + decision = cast( + RouterOutput, ai_state.structured(response) or RouterOutput.model_validate_json(ai_state.text(response)) + ) out: dict = {"intent": decision.intent, "contexts": decision.contexts} if decision.intent == "chat": @@ -75,7 +57,7 @@ async def job_search_node(state: OverallState, config: RunnableConfig) -> dict: jobs: list[dict] = [] calls: list[dict] = [] - for event in response.tool_events: + for event in ai_state.tool_events(response): if event["name"] == "job_search_tool": jobs.extend(json.loads(event["content"])) calls.append({"name": event["name"], "args": event["args"]}) @@ -112,20 +94,18 @@ def after_ask(state: OverallState) -> str: # --- job summarizer -------------------------------------------------------- -SUMMARIZER_PROMPT = """Write a short, friendly summary of these job listings for the user. -Group by role type, mention company and location. If the list is empty, say so plainly. -Use only what is given - do not invent jobs.""" - -async def job_summarizer_node(state: OverallState) -> dict: +async def job_summarizer_node(state: OverallState, config: RunnableConfig) -> dict: """Context: only the tool response. Never sees the query or the profile.""" if not state["data"] and state["query"]: # ask_user still wants a refined query; don't spend a call summarizing nothing return {} - response = await llm().ainvoke([SystemMessage(SUMMARIZER_PROMPT), HumanMessage(json.dumps(state["data"] or []))]) + from app.agents.agent import JobSummarizerAgent + + response = await JobSummarizerAgent(state, config).prompt(json.dumps(state["data"] or [])) - return {"type": "text", "data": str(response.text), "data_type": "string"} + return {"type": "text", "data": ai_state.text(response), "data_type": "string"} def ask_or_finish(state: OverallState) -> str: diff --git a/example/agents/app/agents/remember.py b/example/agents/app/agents/remember.py new file mode 100644 index 00000000..940cd01d --- /dev/null +++ b/example/agents/app/agents/remember.py @@ -0,0 +1,60 @@ +import json + +from fastapi_startkit.ai import Agent +from fastapi_startkit.ai import state as ai_state +from langchain_core.runnables import RunnableConfig + +from app.agents.state import OverallState +from app.models.message import Message +from app.models.thread import Thread + + +class ThreadAgent(Agent): + """An agent bound to a conversation thread via the graph's (state, config).""" + + def __init__(self, state: OverallState | None = None, config: RunnableConfig | None = None): + super().__init__() + self.state: dict = dict(state or {}) + configurable = (config or {}).get("configurable") or {} + self.thread_id: str = configurable.get("thread_id", "") + self.turn_id: str | None = configurable.get("turn_id") + + +class RememberMixin: + """Persists what the agent did — tool calls, tool responses and text replies — + onto its thread, right after each prompt(). No thread_id -> remembers nothing. + """ + + async def prompt(self, message: str, **kwargs) -> dict: + state = await super().prompt(message, **kwargs) + await self.remember(state) + return state + + async def remember(self, state: dict) -> None: + events = ai_state.tool_events(state) + for event in events: + await self.remember_row("tool_call", {"name": event["name"], "args": event["args"]}, state) + await self.remember_row("tool_response", json.loads(event["content"]), state) + + # A plain text reply; structured output is state, not conversation. + text = ai_state.text(state) + if text and not events and ai_state.structured(state) is None: + await self.remember_row("text", {"text": text}, state) + + async def remember_row(self, type_: str, data: dict | list, state: dict) -> None: + if not self.thread_id: # not bound to a conversation -> nothing to remember + return + await Thread.first_or_create({"id": self.thread_id}, {"id": self.thread_id}) + row: dict = { + "thread_id": self.thread_id, + "role": "ai", + "type": type_, + "data": data, + "data_type": "string" if type_ == "text" else "json", + "run_id": self.turn_id, + "meta": {"agent": type(self).__name__, **({"model": self.model} if self.model else {})}, + } + usage = ai_state.usage(state) + if usage["input"] or usage["output"]: + row["usage"] = usage + await Message.create(row) diff --git a/example/agents/app/tools/job_search_tool.py b/example/agents/app/tools/job_search_tool.py index 556218cc..82e5de92 100644 --- a/example/agents/app/tools/job_search_tool.py +++ b/example/agents/app/tools/job_search_tool.py @@ -9,12 +9,22 @@ ] +# Generic words that describe "a job" rather than which job — searching them +# would match nothing even though the user clearly wants to see openings. +NOISE_TERMS = {"job", "jobs", "role", "roles", "position", "positions", "opening", "openings", "work", "career"} + + @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 - patterns = [f"*{term}*" for term in query.lower().split()] + patterns = [f"*{term}*" for term in query.lower().split() if term not in NOISE_TERMS] + + if not patterns: + # A blank or all-noise query ("jobs") means "show me the board", not + # "match nothing" — the model sends these when there are no usable keywords. + return jobs return [ job diff --git a/example/agents/routes/api.py b/example/agents/routes/api.py index df429931..6f033b1e 100644 --- a/example/agents/routes/api.py +++ b/example/agents/routes/api.py @@ -1,7 +1,9 @@ import json +import uuid from fastapi import APIRouter, Request from fastapi.responses import StreamingResponse +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.inertia import Inertia from langgraph.types import Command @@ -28,14 +30,25 @@ async def index(request: Request): @api.post("/chat") async def chat(request: ChatRequest): response = await ChatAgent().prompt(request.message) - return {"content": response.content} + return {"content": ai_state.text(response)} + + +def _frames(event: dict) -> dict | None: + """Map a StandardStreamEvent onto the flat SSE frame the frontend renders.""" + if event["event"] == "on_chat_model_stream" and (text := str(event["data"]["chunk"].text)): + return {"type": "delta", "text": text} + if event["event"] == "on_tool_end": + output = event["data"].get("output") + return {"type": "tool_response", "name": event["name"], "content": str(getattr(output, "content", output))} + return None @api.post("/chat/stream") async def chat_stream(request: ChatRequest): async def generate(): - async for frame in ChatAgent().stream(request.message): - yield f"data: {json.dumps(frame)}\n\n" + async for event in ChatAgent().stream(request.message): + if frame := _frames(event): + yield f"data: {json.dumps(frame)}\n\n" return StreamingResponse(generate(), media_type="text/event-stream") @@ -45,8 +58,9 @@ async def sales_stream(request: ChatRequest): config = {"configurable": {"thread_id": request.thread_id}} async def generate(): - async for frame in SalesAgent().stream(request.message, config=config): - yield f"data: {json.dumps(frame)}\n\n" + async for event in SalesAgent().stream(request.message, config=config): + if frame := _frames(event): + yield f"data: {json.dumps(frame)}\n\n" return StreamingResponse(generate(), media_type="text/event-stream") @@ -66,7 +80,8 @@ async def jobs_page(): @api.post("/jobs/stream") async def jobs_stream(request: ChatRequest): graph = await job_graph() - config = {"configurable": {"thread_id": request.thread_id}} + # turn_id groups every row this turn produces; RememberMixin stamps it as run_id. + config = {"configurable": {"thread_id": request.thread_id, "turn_id": str(uuid.uuid4())}} def sse(frame: dict) -> str: return f"data: {json.dumps(frame)}\n\n" @@ -77,108 +92,57 @@ async def generate(): snapshot = await graph.aget_state(config) payload = Command(resume=request.message) if snapshot.interrupts else {"query": request.message} - run_id: str | None = None - tool_rows: list[tuple[str, dict]] = [] # (node, row) - streamed: list[str] = [] - reply_node = "ask_user" # overwritten by whichever node actually speaks - usage_by_node: dict[str, dict] = {} - model_by_node: dict[str, str] = {} + # Each agent persists its own side of the turn (RememberMixin); the route only + # stores what no agent owns: the user's message, before the run so agents can + # tell history from the current turn, and the interrupt question after it. + await Thread.first_or_create({"id": request.thread_id}, {"id": request.thread_id}) + await Message.create( + { + "thread_id": request.thread_id, + "role": "user", + "type": "text", + "data": {"text": request.message}, + "data_type": "string", + "run_id": None, + } + ) + streamed = False async for event in graph.astream_events(payload, config): - run_id = run_id or event.get("run_id") node = (event.get("metadata") or {}).get("langgraph_node", "") - # Token usage per node, from every LLM call the run makes (including the - # buffered ones inside JobSearchAgent — their model events still surface). - if event["event"] == "on_chat_model_end" and node: - meta = getattr(event["data"].get("output"), "usage_metadata", None) or {} - if meta: - totals = usage_by_node.setdefault(node, {"input": 0, "output": 0}) - totals["input"] += meta.get("input_tokens", 0) - totals["output"] += meta.get("output_tokens", 0) - if model_name := (event.get("metadata") or {}).get("ls_model_name"): - model_by_node[node] = model_name - - # Capture every node that commits a tool_response envelope (job_search today, - # company_research tomorrow): downstream nodes overwrite state["data"], so the - # raw rows only exist here. Forward them so the UI can render before the prose. + # Forward every tool_response envelope a node commits (job_search today, + # company_research tomorrow) so the UI can render before the prose lands. if event["event"] == "on_chain_end" and node and event["name"] == node: out = event["data"].get("output") or {} if isinstance(out, dict) and out.get("type") == "tool_response": - rows = out.get("data") or [] - calls = out.get("tool_calls") or [] - for call in calls: - tool_rows.append((node, {"type": "tool_call", "data": call, "data_type": "json"})) - tool_rows.append((node, {"type": "tool_response", "data": rows, "data_type": "json"})) - yield sse( - { - "kind": "envelope", - "node": node, - "type": "tool_response", - "data": rows, - "data_type": "json", - "tool_calls": calls, - } - ) - - if event["event"] != "on_chat_model_stream": - continue + yield sse({"kind": "envelope", "node": node, **out}) + # Stream every node's tokens, tagged with the node; the frontend decides # what to render (the router's are structured-output JSON, for instance). - if (chunk := event["data"].get("chunk")) and (text := str(chunk.text)): - if node != "router": # persisted reply = user-facing text only - streamed.append(text) - reply_node = node - yield sse({"kind": "delta", "node": node, "text": text}) + if event["event"] == "on_chat_model_stream": + if (chunk := event["data"].get("chunk")) and (text := str(chunk.text)): + streamed = streamed or node != "router" + yield sse({"kind": "delta", "node": node, "text": text}) snapshot = await graph.aget_state(config) - reply = "".join(streamed) if snapshot.interrupts: # No jobs -> the graph paused at ask_user; surface its question for the UI. - reply = snapshot.interrupts[0].value["message"] - reply_node = "ask_user" - yield sse({"kind": "interrupt", **snapshot.interrupts[0].value}) - elif not reply and snapshot.values.get("type") == "text": - # Greeting answered inline by the router (not token-streamed) -> emit it whole. - reply = str(snapshot.values.get("data") or "") - reply_node = "router" - yield sse({"kind": "envelope", "type": "text", "data": reply, "data_type": "string"}) - - def enrich(node: str) -> dict: - meta: dict = {"node": node} - if model_name := model_by_node.get(node): - meta["model"] = model_name - extra: dict = {"meta": meta} - if usage := usage_by_node.get(node): - extra["usage"] = usage - return extra - - # Persist the turn only now: inserting the user row before the run would make - # JobSearchAgent load the current query as history and duplicate it. - await Thread.first_or_create({"id": request.thread_id}, {"id": request.thread_id}) - await Message.create( - { - "thread_id": request.thread_id, - "role": "user", - "type": "text", - "data": {"text": request.message}, - "data_type": "string", - "run_id": None, - } - ) - for node, row in tool_rows: - await Message.create({"thread_id": request.thread_id, "role": "ai", "run_id": run_id, **enrich(node), **row}) - if reply: + question = snapshot.interrupts[0].value + yield sse({"kind": "interrupt", **question}) await Message.create( { "thread_id": request.thread_id, "role": "ai", "type": "text", - "data": {"text": reply}, + "data": {"text": question["message"]}, "data_type": "string", - "run_id": run_id, - **enrich(reply_node), + "run_id": config["configurable"]["turn_id"], + "meta": {"agent": "ask_user"}, } ) + elif not streamed and snapshot.values.get("type") == "text": + # Greeting answered inline by the router (not token-streamed) -> emit it whole. + yield sse({"kind": "envelope", "type": "text", "data": snapshot.values.get("data") or "", "data_type": "string"}) return StreamingResponse(generate(), media_type="text/event-stream") diff --git a/example/agents/tests/features/test_jobs_stream.py b/example/agents/tests/features/test_jobs_stream.py new file mode 100644 index 00000000..c6730f36 --- /dev/null +++ b/example/agents/tests/features/test_jobs_stream.py @@ -0,0 +1,120 @@ +import json +import uuid + +from fastapi_startkit.ai.ai import Ai +from langchain_core.messages import AIMessage +from langgraph.checkpoint.memory import InMemorySaver + +import app.agents.job_search_graph as graph_module +from app.agents.state import RouterOutput +from app.models.message import Message +from app.models.thread import Thread +from tests.support.fakes import StreamingToolFake, tool_call_message +from tests.test_case import TestCase + + +class AwaitableSaver: + def __init__(self, saver): + self._saver = saver + + def __await__(self): + async def resolve(): + return self._saver + + return resolve().__await__() + + +class TestJobsStream(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + self.thread_id = f"test-{uuid.uuid4()}" + + # Fresh in-memory checkpointer per test: reset the cached compiled graph so + # job_graph() recompiles against it instead of Postgres. + from fastapi_startkit.application import app + + app().bind("checkpointer", AwaitableSaver(InMemorySaver())) + graph_module._graph = None + self.addCleanup(setattr, graph_module, "_graph", None) + self.addCleanup(Ai.reset_fakes) + + async def asyncTearDown(self): + await Message.where("thread_id", self.thread_id).delete() + await Thread.where("id", self.thread_id).delete() + await super().asyncTearDown() + + def route_to(self, *decisions: RouterOutput): + # The router is an agent too now: fake its structured-output turns as JSON + # content the runner parses back through schema(). + Ai.fake("JobSearchRouterAgent", [AIMessage(content=d.model_dump_json()) for d in decisions]) + + async def frames(self, message: str) -> list[dict]: + response = await self.client.post("/jobs/stream", json={"message": message, "thread_id": self.thread_id}) + frames = [] + for block in response.text.split("\n\n"): + data = "".join(line[6:] for line in block.split("\n") if line.startswith("data: ")) + if data: + frames.append(json.loads(data)) + return frames + + async def rows(self) -> list[Message]: + return list(await Message.where("thread_id", self.thread_id).order_by("id").get()) + + def fake_search_calls(self, *messages: AIMessage) -> None: + Ai._fakes["JobSearchAgent"] = StreamingToolFake(messages=iter(messages)) + + async def test_search_streams_the_envelope_and_persists_the_turn(self): + self.route_to(RouterOutput(intent="job_search")) + self.fake_search_calls(tool_call_message("python developer")) + Ai.fake("JobSummarizerAgent", [AIMessage(content="Here is the Python Developer role.")]) + + frames = await self.frames("find python jobs") + + envelope = next(f for f in frames if f["kind"] == "envelope") + assert envelope["node"] == "job_search" + assert [job["title"] for job in envelope["data"]] == ["Python Developer"] + assert envelope["tool_calls"] == [{"name": "job_search_tool", "args": {"query": "python developer"}}] + deltas = "".join(f["text"] for f in frames if f["kind"] == "delta" and f["node"] == "job_summarizer") + assert deltas == "Here is the Python Developer role." + + rows = await self.rows() + assert [(m.role, m.type) for m in rows] == [ + ("user", "text"), + ("ai", "tool_call"), + ("ai", "tool_response"), + ("ai", "text"), + ] + assert rows[1].data == {"name": "job_search_tool", "args": {"query": "python developer"}} + assert rows[3].data == {"text": "Here is the Python Developer role."} + assert rows[3].meta["agent"] == "JobSummarizerAgent" + assert rows[1].run_id is not None and rows[0].run_id is None + + async def test_greeting_is_answered_inline_by_the_router(self): + self.route_to(RouterOutput(intent="chat", reply="Hello! How can I help?")) + + frames = await self.frames("hi") + + assert {"kind": "envelope", "type": "text", "data": "Hello! How can I help?", "data_type": "string"} in frames + + rows = await self.rows() + assert [(m.role, m.type) for m in rows] == [("user", "text"), ("ai", "text")] + assert rows[1].meta["agent"] == "JobSearchRouterAgent" + + async def test_no_jobs_interrupts_then_a_reply_resumes_the_search(self): + self.route_to(RouterOutput(intent="job_search")) + self.fake_search_calls(tool_call_message("cobol mainframe antarctica"), tool_call_message("python developer")) + Ai.fake("JobSummarizerAgent", [AIMessage(content="Found it after refining.")]) + + frames = await self.frames("any cobol roles?") + + interrupt = next(f for f in frames if f["kind"] == "interrupt") + assert interrupt["reason"] == "no_jobs" + assert (await self.rows())[-1].meta["agent"] == "ask_user" + + # The next message on this thread resumes the paused graph with a refined query. + frames = await self.frames("python developer") + + envelope = next(f for f in frames if f["kind"] == "envelope") + assert [job["title"] for job in envelope["data"]] == ["Python Developer"] + assert not [f for f in frames if f["kind"] == "interrupt"] + assert (await self.rows())[-1].data == {"text": "Found it after refining."} diff --git a/example/agents/tests/integrations/agents/job_search_agent.json b/example/agents/tests/integrations/agents/job_search_agent.json new file mode 100644 index 00000000..6cd88887 --- /dev/null +++ b/example/agents/tests/integrations/agents/job_search_agent.json @@ -0,0 +1,66 @@ +{ + "5ef8e6a4df0f363d14e6d4ccd9bbb79693bc9440fc00aa3946ccc32d895b3e1a": [ + { + "content": "suggest me jobs", + "type": "human" + }, + { + "response_time": 596.5392499929294, + "tool_calls": [ + { + "args": { + "query": " " + }, + "id": "4b9174a8-c418-49b4-b6e8-01709d36725b", + "name": "job_search_tool", + "type": "tool_call" + } + ], + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 155, + "output_token": 16, + "total_token": 171 + } + }, + { + "content": "[{\"id\": 1, \"title\": \"Software Engineer\", \"location\": \"San Francisco\", \"company\": \"Acme Corp\", \"type\": \"Full-time\"}, {\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}, {\"id\": 3, \"title\": \"Data Scientist\", \"location\": \"New York\", \"company\": \"DataCo\", \"type\": \"Full-time\"}, {\"id\": 4, \"title\": \"DevOps Engineer\", \"location\": \"Austin\", \"company\": \"CloudBase\", \"type\": \"Contract\"}, {\"id\": 5, \"title\": \"Product Manager\", \"location\": \"Remote\", \"company\": \"ProductHQ\", \"type\": \"Full-time\"}]", + "content_type": "json", + "response_time": 1.1912910267710686, + "type": "tool_response" + } + ], + "8f9e73965490770e00b6bde6b5ceed9ce44ca9a9e5ca6f45d1460b4b6278dc84": [ + { + "content": "find roles that fit me", + "type": "human" + }, + { + "response_time": 562.1582079911605, + "tool_calls": [ + { + "args": { + "query": "software developer" + }, + "id": "ad782a7d-bc5c-43f0-991d-14e4564cd58a", + "name": "job_search_tool", + "type": "tool_call" + } + ], + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 157, + "output_token": 18, + "total_token": 175 + } + }, + { + "content": "[{\"id\": 1, \"title\": \"Software Engineer\", \"location\": \"San Francisco\", \"company\": \"Acme Corp\", \"type\": \"Full-time\"}, {\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "content_type": "json", + "response_time": 1.2913750251755118, + "type": "tool_response" + } + ] +} \ No newline at end of file diff --git a/example/agents/tests/integrations/agents/job_summarizer_agent.json b/example/agents/tests/integrations/agents/job_summarizer_agent.json new file mode 100644 index 00000000..a44869d4 --- /dev/null +++ b/example/agents/tests/integrations/agents/job_summarizer_agent.json @@ -0,0 +1,40 @@ +{ + "60b56ef7d9e8a2e0bea169dd8df0fe6da53ecdc08e1ace9d01c030494acda01d": [ + { + "content": "[]", + "type": "human" + }, + { + "content": "[{'type': 'text', 'text': 'There are currently no job listings to summarize.', 'extras': {'signature': 'EjQKMgERTTIP53WUt3GgvqNEeWshZKHUf0EVT+L2i2D6BGcdUgYP73hZpIqpY+whbgiJMfro'}}]", + "response_time": 764.6852920297533, + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 49, + "output_token": 9, + "total_token": 58 + } + } + ], + "e201c72a97325bb8d3d885bfd735e5442e4fddbd48a5675b734860bf17ab362a": [ + { + "content": "[{\"id\": 2, \"title\": \"Python Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}, {\"id\": 5, \"title\": \"Product Manager\", \"location\": \"Remote\", \"company\": \"ProductHQ\", \"type\": \"Full-time\"}]", + "type": "human" + }, + { + "content": "[{'type': 'text', 'text': 'Here is a summary of the current job listings:\\n\\n**Development**\\n* **Python Developer** at Startup Inc (Remote)\\n\\n**Product Management**\\n* **Product Manager** at ProductHQ (Remote)', 'extras': {'signature': 'EjQKMgERTTIPhq3ltEUqORsD7yOhUgrr7N5bXv423tiAdkFrogUgx1imoHYIgAQsJG8tYY5y'}}]", + "response_time": 797.7974159875885, + "type": "ai", + "uses": { + "cache_token": 0, + "input_token": 118, + "output_token": 45, + "total_token": 163 + } + } + ], + "judge:bf29ceac94d0cae9617aafb3d708ada88d7e2f4f67bec3a379117e051cfe7917": { + "passed": true, + "reasoning": "The response accurately and friendly summarizes exactly the two specified roles with their respective companies and remote locations without inventing any other jobs." + } +} \ No newline at end of file diff --git a/example/agents/tests/integrations/agents/test_job_search_integration.py b/example/agents/tests/integrations/agents/test_job_search_integration.py new file mode 100644 index 00000000..ac7be35f --- /dev/null +++ b/example/agents/tests/integrations/agents/test_job_search_integration.py @@ -0,0 +1,96 @@ +"""Integration tests: real prompt assembly, real tool execution, and real model +behaviour captured with the record/replay harness. The first run hits the live +API and writes the cassettes next to this file; every run after replays them. +""" + +import json + +from fastapi_startkit.ai import AssertToolCall +from fastapi_startkit.ai import state as ai_state +from fastapi_startkit.ai.testing import AgentRecordFake + +from app.agents.agent import JobSearchAgent, JobSummarizerAgent +from app.agents.state import Context +from app.models.message import Message +from app.models.thread import Thread +from tests.test_case import TestCase + +THREAD_ID = "integration-job-search" + +JOBS = [ + {"id": 2, "title": "Python Developer", "location": "Remote", "company": "Startup Inc", "type": "Full-time"}, + {"id": 5, "title": "Product Manager", "location": "Remote", "company": "ProductHQ", "type": "Full-time"}, +] + + +class TestJobSearchAgentIntegration(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + # The same seed every run: on the recording run it shapes the real prompt; + # on replay runs the cassette answers and the rows are simply unused. + await Thread.first_or_create({"id": THREAD_ID}, {"id": THREAD_ID}) + await Message.where("thread_id", THREAD_ID).delete() + await Message.create( + { + "thread_id": THREAD_ID, + "role": "user", + "type": "text", + "data": {"text": "I prefer fully remote roles."}, + "data_type": "string", + "run_id": None, + } + ) + + async def asyncTearDown(self): + await Message.where("thread_id", THREAD_ID).delete() + await Thread.where("id", THREAD_ID).delete() + await super().asyncTearDown() + + def recorded_agent(self, contexts: list[Context]) -> AgentRecordFake: + real = JobSearchAgent({"contexts": contexts}, {"configurable": {"thread_id": THREAD_ID}}) + return AgentRecordFake(real, "job_search_agent.json") + + async def test_searches_and_finds_roles_for_a_natural_query(self): + with self.recorded_agent([Context.INCLUDE_USER_PROFILE]) as agent: + response = await agent.prompt("find roles that fit me") + + def called_with_a_real_query(tool: AssertToolCall) -> bool: + return tool.args.get("query", "").strip() != "" + + agent.assert_tool_called("job_search_tool", called_with_a_real_query) + assert json.loads(ai_state.text(response)) != [] + agent.assert_tokens(lambda x: x.where("input", "<=", 20_000).where("output", "<=", 2_000)) + agent.assert_response_time_lt(30) + + async def test_always_searches_even_for_a_vague_query(self): + with self.recorded_agent([Context.INCLUDE_USER_PROFILE]) as agent: + response = await agent.prompt("suggest me jobs") + + # tool_choice="any": a text-only reply here would be a regression. The + # recording run showed the model may send a blank query for vague asks, + # which the tool now treats as "show the whole board". + agent.assert_tool_called("job_search_tool") + assert json.loads(ai_state.text(response)) != [] + + +class TestJobSummarizerAgentIntegration(TestCase): + async def test_summarizes_only_what_it_is_given(self): + with JobSummarizerAgent.record("job_summarizer_agent.json") as agent: + response = await agent.prompt(json.dumps(JOBS)) + + agent.assert_text_response() + assert "Python Developer" in ai_state.text(response) + assert "Product Manager" in ai_state.text(response) + await agent.assert_response_judged( + provider="google", + model="gemini-3.5-flash", + expectation="A friendly summary of exactly two roles - Python Developer at Startup Inc and " + "Product Manager at ProductHQ, both remote - inventing no other jobs.", + ) + + async def test_says_so_plainly_when_there_are_no_jobs(self): + with JobSummarizerAgent.record("job_summarizer_agent.json") as agent: + response = await agent.prompt(json.dumps([])) + + agent.assert_text_response() + assert "Python Developer" not in ai_state.text(response) diff --git a/example/agents/tests/support/fakes.py b/example/agents/tests/support/fakes.py new file mode 100644 index 00000000..2f3be8a9 --- /dev/null +++ b/example/agents/tests/support/fakes.py @@ -0,0 +1,34 @@ +import json +import re +from typing import cast + +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, AIMessageChunk +from langchain_core.outputs import ChatGenerationChunk + + +def tool_call_message(query: str) -> AIMessage: + """A model turn that calls job_search_tool with the given query.""" + return AIMessage( + content="", + tool_calls=[{"name": "job_search_tool", "args": {"query": query}, "id": "c1", "type": "tool_call"}], + ) + + +class StreamingToolFake(GenericFakeChatModel): + """GenericFakeChatModel can't stream a content-less tool-call message (it yields + zero chunks -> 'No generations found'). Under astream_events models auto-stream, + so emit proper tool-call chunks, mirroring the framework's own test fake.""" + + 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)) diff --git a/example/agents/tests/test_case.py b/example/agents/tests/test_case.py index b1a97db6..182ec428 100644 --- a/example/agents/tests/test_case.py +++ b/example/agents/tests/test_case.py @@ -6,3 +6,18 @@ def get_application(self): from bootstrap.application import app return app + + async def asyncSetUp(self): + # Each test runs in a fresh event loop, but the ORM caches engines whose + # asyncpg connections are bound to the loop that created them. Drop the + # cache so this test's loop builds its own engine. + from fastapi_startkit.masoniteorm.models import Model + + for connection in list(Model.db_manager.connections.values()): + try: + await connection.engine.dispose() + except Exception: # noqa: BLE001 - engine belonged to a dead loop + pass + Model.db_manager.connections.clear() + + await super().asyncSetUp() diff --git a/example/agents/tests/units/agents/test_job_search_node.py b/example/agents/tests/units/agents/test_job_search_node.py new file mode 100644 index 00000000..ac5bd313 --- /dev/null +++ b/example/agents/tests/units/agents/test_job_search_node.py @@ -0,0 +1,159 @@ +import uuid + +from fastapi_startkit.ai import Ai +from langchain_core.messages import AIMessage, SystemMessage +from langgraph.graph import END + +from app.agents.agent import JobSearchAgent +from app.agents.job_search_graph import after_ask, ask_or_finish, build, job_search_node, route +from app.agents.state import Context, RouterOutput +from app.models.message import Message +from app.models.thread import Thread +from tests.support.fakes import tool_call_message +from tests.test_case import TestCase + + +def router_decision(**kwargs) -> AIMessage: + """A faked structured-output turn: the runner parses the JSON content back + into RouterOutput via the agent's schema().""" + return AIMessage(content=RouterOutput(**kwargs).model_dump_json()) + + +class TestJobSearchGraph(TestCase): + """The whole graph, every LLM faked at the agent seam, real tool + real edges.""" + + async def asyncSetUp(self): + await super().asyncSetUp() + self.addCleanup(Ai.reset_fakes) + + async def test_it_responds_to_greetings(self): + Ai.fake("JobSearchRouterAgent", [router_decision(intent="chat", reply="hi, how are you doing")]) + + graph = build().compile() + out = await graph.ainvoke({"query": "hello"}) + + assert out == {"type": "text", "data": "hi, how are you doing", "data_type": "string"} + + async def test_it_searches_and_summarizes(self): + # A fake holds one scripted turn per graph run, consumed in order. + Ai.fake( + "JobSearchRouterAgent", + [ + router_decision(intent="job_search"), + router_decision(intent="chat", reply="Hello again! Ready for more roles?"), + ], + ) + Ai.fake("JobSearchAgent", [tool_call_message("python developer")]) + Ai.fake("JobSummarizerAgent", [AIMessage(content="One Python role, fully remote.")]) + + graph = build().compile() + + response = await graph.ainvoke({"query": "find python jobs"}) + assert response == {"type": "text", "data": "One Python role, fully remote.", "data_type": "string"} + + response = await graph.ainvoke({"query": "Hello"}) + assert response == {"type": "text", "data": "Hello again! Ready for more roles?", "data_type": "string"} + + + + +class TestJobSearchNode(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + self.thread_id = f"test-{uuid.uuid4()}" + self.config = {"configurable": {"thread_id": self.thread_id}} + self.addCleanup(Ai.reset_fakes) + + async def asyncTearDown(self): + await Message.where("thread_id", self.thread_id).delete() + await Thread.where("id", self.thread_id).delete() + await super().asyncTearDown() + + async def test_returns_matching_jobs_and_the_calls_it_made(self): + Ai.fake("JobSearchAgent", [tool_call_message("python developer")]) + + out = await job_search_node({"query": "find python jobs", "contexts": []}, self.config) + + assert out["type"] == "tool_response" + assert out["data_type"] == "json" + assert [job["title"] for job in out["data"]] == ["Python Developer"] + assert out["tool_calls"] == [{"name": "job_search_tool", "args": {"query": "python developer"}}] + + async def test_returns_empty_data_when_nothing_matches(self): + Ai.fake("JobSearchAgent", [tool_call_message("cobol mainframe antarctica")]) + + out = await job_search_node({"query": "cobol roles?", "contexts": []}, self.config) + + assert out["data"] == [] + assert out["tool_calls"] == [{"name": "job_search_tool", "args": {"query": "cobol mainframe antarctica"}}] + + +class TestGraphRouting(TestCase): + def test_route_sends_job_search_to_the_search_branch(self): + assert route({"intent": "job_search"}) == "job_search" + + def test_route_ends_for_chat_and_unbuilt_branches(self): + assert route({"intent": "chat"}) == END + assert route({"intent": "company_research"}) == END + + def test_ask_or_finish_asks_only_when_empty_with_a_query(self): + assert ask_or_finish({"data": [], "query": "python"}) == "ask_user" + assert ask_or_finish({"data": [{"id": 1}], "query": "python"}) == END + assert ask_or_finish({"data": [], "query": ""}) == END + + def test_after_ask_retries_with_a_reply_and_gives_up_on_blank(self): + assert after_ask({"query": "python developer"}) == "job_search" + assert after_ask({"query": ""}) == "job_summarizer" + + +class TestJobSearchAgentContext(TestCase): + async def asyncSetUp(self): + await super().asyncSetUp() + self.thread_id = f"test-{uuid.uuid4()}" + await Thread.first_or_create({"id": self.thread_id}, {"id": self.thread_id}) + + async def asyncTearDown(self): + await Message.where("thread_id", self.thread_id).delete() + await Thread.where("id", self.thread_id).delete() + await super().asyncTearDown() + + def agent(self, contexts: list[Context]) -> JobSearchAgent: + return JobSearchAgent({"contexts": contexts}, {"configurable": {"thread_id": self.thread_id}}) + + async def seed(self, role: str, type_: str, data: dict | list, data_type: str) -> None: + await Message.create( + { + "thread_id": self.thread_id, + "role": role, + "type": type_, + "data": data, + "data_type": data_type, + "run_id": None, + } + ) + + async def test_loads_only_user_messages_as_history(self): + await self.seed("user", "text", {"text": "remote only please"}, "string") + await self.seed("ai", "text", {"text": "Got it."}, "string") + + messages = await self.agent([]).messages() + + assert [m.content for m in messages] == ["remote only please"] + + async def test_injects_the_profile_only_when_asked(self): + assert await self.agent([]).messages() == [] + + messages = await self.agent([Context.INCLUDE_USER_PROFILE]).messages() + + assert isinstance(messages[0], SystemMessage) + assert "Alice" in messages[0].content + + async def test_injects_the_last_tool_response_only_when_asked(self): + await self.seed("ai", "tool_response", [{"id": 2, "title": "Python Developer"}], "json") + + assert await self.agent([]).messages() == [] + + messages = await self.agent([Context.INCLUDE_LAST_JOB_SEARCH_RESPONSE]).messages() + + assert len(messages) == 1 + assert "Python Developer" in messages[0].content diff --git a/example/agents/tinker.py b/example/agents/tinker.py index 6d95e2b0..fa0bb052 100644 --- a/example/agents/tinker.py +++ b/example/agents/tinker.py @@ -1,5 +1,30 @@ -from app.agents.job_search_graph import build +import asyncio -graph = build().compile() +from dumpdie import dd +from langchain.agents import create_agent +from langchain_core.tools import tool -print(graph.get_graph().draw_mermaid()) +from app.agents.agent import JobSearchAgent +from bootstrap.application import app # NOQA + + +@tool(description="Use this tool to search for jobs", return_direct=True) +def job_search_tool(query: str): + return [] + + +async def main(): + responses = await JobSearchAgent().prompt("suggest me python developer jobs") + dd(responses) + + +asyncio.run(main()) +# +# agent = create_agent(model="google_genai:gemini-3.1-flash-lite", tools=[job_search_tool]) +# +# # agent.invoke({"messages": [{"role": "user", "content": "hi"}]}) +# +# responses = agent.invoke( +# {"messages": [{"role": "user", "content": "suggest me python developer jobs"}]} +# ) +# dd(responses) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index 31cc73b4..9ca39a00 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -12,7 +12,7 @@ from .ai import Ai from .judge import JudgeAgent from .providers.ai_provider import AIProvider -from .response import AgentResponse, AgentSnapshot +from . import state from .testing import AgentFake, AgentRecordFake, AssertToolCall __all__ = [ @@ -20,8 +20,7 @@ "Ai", "Middleware", "AgentFake", - "AgentResponse", - "AgentSnapshot", + "state", "AIConfig", "AIProvider", "AnthropicConfig", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index d46ba94a..a6fb6837 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -3,7 +3,6 @@ from typing import TYPE_CHECKING, AsyncIterator, Optional, Type from .document import Document -from .response import AgentResponse from .runner import BaseRunner, Runner if TYPE_CHECKING: @@ -49,7 +48,10 @@ async def prompt( model: str | None = None, attachments: list[Document] | None = None, provider_options: dict | None = None, - ) -> AgentResponse: + ) -> dict: + """Returns the turn as ``{"messages": [HumanMessage, AIMessage, ToolMessage, + ...]}`` — the same state langchain's ``create_agent().invoke()`` yields — + plus ``"structured_response"`` when the agent defines a schema().""" return await self.runner().run(message, model=model, attachments=attachments, provider_options=provider_options) async def stream( @@ -59,8 +61,10 @@ async def stream( model: str | None = None, provider_options: dict | None = None, ) -> AsyncIterator[dict]: - """Yields typed frames: {"type": "delta", "text": ...} for model tokens and - {"type": "tool_response", "name": ..., "content": ...} for tool results.""" + """Yields StandardStreamEvent dicts — the same shape LangChain's + astream_events produces: on_chat_model_(start|stream|end) for the model + turn (data.chunk carries each AIMessageChunk) and on_tool_(start|end) + around each tool execution.""" async for chunk in self.runner().stream(message, model=model, provider_options=provider_options): yield chunk diff --git a/fastapi_startkit/src/fastapi_startkit/ai/graph.py b/fastapi_startkit/src/fastapi_startkit/ai/graph.py index 0c6b1662..51f76483 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/graph.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/graph.py @@ -5,14 +5,12 @@ from collections.abc import AsyncIterator from typing import TYPE_CHECKING, Annotated, Any, TypedDict -from langchain_core.messages import ToolMessage from langchain_core.runnables import RunnableConfig from langgraph.graph import END from langgraph.graph.message import add_messages from .agent import Agent from .pipeline import Response, build_pipeline -from .response import AgentResponse from .runner import BaseRunner if TYPE_CHECKING: @@ -41,7 +39,7 @@ async def prompt( attachments: list[Document] | None = None, config: RunnableConfig | dict | None = None, provider_options: dict | None = None, - ) -> AgentResponse: + ) -> dict: return await self.runner().run( message, model=model, @@ -129,15 +127,13 @@ async def run( attachments: list[Document] | None = None, config: RunnableConfig | dict | None = None, provider_options: dict | None = None, - ) -> AgentResponse: - started = time.perf_counter() + ) -> dict: self._reset_capture() compiled = await self._compile(model, provider_options) state = {"messages": await self._build_messages(message, attachments), "llm_calls": 0} result = await compiled.ainvoke(state, config=self._merge_config(config)) - response = self._apply_schema(self._to_response(result)) - response.runtime = time.perf_counter() - started - return response + result["messages"] = self._with_response_times(result.get("messages", [])) + return self._apply_schema(result) async def stream( self, @@ -147,47 +143,28 @@ async def stream( config: RunnableConfig | dict | None = None, provider_options: dict | None = None, ) -> AsyncIterator[dict]: + """Yields the graph's astream_events verbatim — StandardStreamEvent dicts + (on_chain_*, on_chat_model_*, on_tool_*). The root chain's end event + carries the final state, exposed as ``last_response``.""" self._reset_capture() compiled = await self._compile(model, provider_options) state = {"messages": await self._build_messages(message), "llm_calls": 0} final_state: dict = {} - # "messages" streams tokens; "values" hands back the full state after each - # step, whose last snapshot carries every message (and its tool_calls). - async for mode, chunk in compiled.astream( - state, stream_mode=["messages", "values"], config=self._merge_config(config) - ): - if mode == "values": - final_state = chunk - continue - message_chunk, _meta = chunk - text = message_chunk.content if isinstance(message_chunk.content, str) else str(message_chunk.content) - if not text: - continue - if isinstance(message_chunk, ToolMessage): - yield {"type": "tool_response", "name": message_chunk.name or "", "content": text} - else: - yield {"type": "delta", "text": text} - self.last_response = self._to_response(final_state) if final_state else AgentResponse(content="") - - def _to_response(self, 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, - tool_events=list(self._tool_events), - transcript=list(self._transcript), - ) + async for event in compiled.astream_events(state, config=self._merge_config(config)): + if event["event"] == "on_chain_end" and not event.get("parent_ids"): + output = event["data"].get("output") + final_state = output if isinstance(output, dict) else {} + yield event + if final_state: + final_state["messages"] = self._with_response_times(final_state.get("messages", [])) + self.last_response = final_state or {"messages": []} + + def _with_response_times(self, messages: list) -> list: + """Stamp each recorded model/tool call's ``response_time`` (ms) onto the + turn's trailing AI/tool messages — the graph state also holds history, so + the recorded entries align with the last N ai/tool messages.""" + entries = [e for e in self._transcript if e["type"] in ("ai", "tool_response")] + targets = [m for m in messages if getattr(m, "type", "") in ("ai", "tool")] + for entry, message in zip(entries, targets[max(0, len(targets) - len(entries)) :]): + message.additional_kwargs["response_time"] = entry.get("response_time", 0.0) + return list(messages) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/judge.py b/fastapi_startkit/src/fastapi_startkit/ai/judge.py index 6923bd3d..fb4d75dc 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/judge.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/judge.py @@ -30,5 +30,5 @@ 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() + state = await self.prompt(f"Expectation: {expectation}\n\nResponse to grade:\n{content}") + return state["structured_response"].model_dump() diff --git a/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py b/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py index 107e0369..bae45993 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/pipeline.py @@ -42,11 +42,22 @@ async def _finish(self, final: Any) -> Any: @staticmethod def _merge(accumulated: Any, chunk: Any) -> Any: - # Stream frames are dicts ({"type": "delta"/"tool_response", ...}); the - # after-hooks want the joined text, not a dict sum. - if isinstance(chunk, dict): - text = chunk.get("text") or chunk.get("content") or "" - return text if accumulated is None else accumulated + text + # Stream events (astream_events shape) merge into joined text for the + # after-hooks: model chunks and tool outputs contribute their text, the + # start/end envelopes add nothing. Any other dict (e.g. a + # structured-output result carrying "raw"/"parsed") passes through whole. + if isinstance(chunk, dict) and "event" in chunk: + data = chunk.get("data") or {} + if chunk["event"] == "on_chat_model_stream": + text = getattr(data.get("chunk"), "content", "") or "" + elif chunk["event"] == "on_tool_end": + text = getattr(data.get("output"), "content", "") or "" + else: + text = "" + text = text if isinstance(text, str) else str(text) + if accumulated is None: + return text or None + return accumulated + text return chunk if accumulated is None else accumulated + chunk def __aiter__(self) -> AsyncIterator: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/recording.py b/fastapi_startkit/src/fastapi_startkit/ai/recording.py index 39299a24..0135c79a 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/recording.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/recording.py @@ -23,7 +23,7 @@ import json from typing import Any -from .response import AgentResponse +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage def _infer_content_type(content: str) -> str: @@ -46,14 +46,6 @@ def uses_from_usage_metadata(meta: dict | None) -> dict: } -def uses_from_agent_usage(usage: dict | None) -> dict: - """Map an :class:`AgentResponse` ``usage`` ({input, output}) onto ``uses``.""" - usage = usage or {} - inp = usage.get("input", 0) - out = usage.get("output", 0) - return {"input_token": inp, "output_token": out, "cache_token": 0, "total_token": inp + out} - - def is_transcript(value: Any) -> bool: """True when ``value`` is a new-format ordered transcript (list of typed entries).""" return isinstance(value, list) and bool(value) and isinstance(value[0], dict) and "type" in value[0] @@ -94,65 +86,91 @@ def tool_response(content: str, response_time: float = 0.0, content_type: str | } -def _usage_from_uses(uses: dict) -> dict: - return {"input": uses.get("input_token", 0), "output": uses.get("output_token", 0)} +def usage_metadata_from_uses(uses: dict) -> dict: + """Map a cassette ``uses`` dict back onto LangChain ``usage_metadata``.""" + return { + "input_tokens": uses.get("input_token", 0), + "output_tokens": uses.get("output_token", 0), + "total_tokens": uses.get("total_token", 0), + "input_token_details": {"cache_read": uses.get("cache_token", 0)}, + } + + +def _as_text(content: Any) -> str: + return content if isinstance(content, str) else str(content) -def to_response(transcript: list[dict]) -> AgentResponse: - """Reconstruct an :class:`AgentResponse` from a recorded turn transcript. +def entries_from_messages(messages: list[BaseMessage]) -> list[dict]: + """Convert a turn's AI/tool LangChain messages into cassette entries. - The final ``ai`` answer supplies the response content; if the turn stopped at - a tool call (no follow-up answer), the tool result stands in as the content. - """ - content = "" - tool_calls: list[dict] = [] - usage: dict = {} - tool_events: list[dict] = [] - runtime = 0.0 + Human (and system) messages are skipped — the cassette turn opens with its + own ``human`` entry. ``response_time`` is read from ``additional_kwargs``, + where the runner stamps it on every AI and tool message.""" + entries: list[dict] = [] + for message in messages or []: + kind = getattr(message, "type", "") + response_time = (getattr(message, "additional_kwargs", None) or {}).get("response_time", 0.0) + if kind == "ai": + entries.append( + ai( + content=_as_text(message.content), + tool_calls=list(getattr(message, "tool_calls", None) or []), + uses=uses_from_usage_metadata(getattr(message, "usage_metadata", None)), + response_time=response_time, + ) + ) + elif kind == "tool": + entries.append(tool_response(content=_as_text(message.content), response_time=response_time)) + return entries + +def messages_from_transcript(transcript: list[dict]) -> list[BaseMessage]: + """Rebuild the turn's LangChain messages from cassette entries, so a replayed + response carries the same ``messages`` a live run produces. Cassette entries + don't store the tool's name/call id, so those stay empty on replay.""" + messages: list[BaseMessage] = [] for entry in transcript: kind = entry.get("type") - if kind == "ai": - if entry.get("tool_calls"): - tool_calls.extend(entry["tool_calls"]) - if entry.get("content"): - content = entry["content"] - if entry.get("uses"): - usage = _usage_from_uses(entry["uses"]) - runtime += (entry.get("response_time") or 0) / 1000 + if kind == "human": + messages.append(HumanMessage(content=entry.get("content", ""))) + elif kind == "ai": + uses = entry.get("uses") or {} + messages.append( + AIMessage( + content=entry.get("content", ""), + tool_calls=entry.get("tool_calls") or [], + additional_kwargs={"response_time": entry.get("response_time", 0.0)}, + usage_metadata=usage_metadata_from_uses(uses) if uses else None, + ) + ) elif kind == "tool_response": - tool_events.append( - { - "content": entry.get("content", ""), - "content_type": entry.get("content_type"), - "response_time": entry.get("response_time", 0), - } + messages.append( + ToolMessage( + content=entry.get("content", ""), + tool_call_id="", + additional_kwargs={"response_time": entry.get("response_time", 0.0)}, + ) ) - if not content: - content = entry.get("content", "") - runtime += (entry.get("response_time") or 0) / 1000 - - return AgentResponse( - content=content, - tool_calls=tool_calls, - usage=usage, - tool_events=tool_events, - runtime=runtime, - transcript=list(transcript), - ) - - -def accumulate_uses(transcript: list[dict], totals: dict) -> None: - """Add every ``ai`` entry's token ``uses`` in ``transcript`` into ``totals`` + return messages + + +def to_state(transcript: list[dict]) -> dict: + """Reconstruct a turn's ``{"messages": [...]}`` state from cassette entries, + so a replayed turn has the same shape a live run returns.""" + return {"messages": messages_from_transcript(transcript)} + + +def accumulate_uses(messages: list, totals: dict) -> None: + """Add every AI message's token usage in a turn's ``messages`` into ``totals`` (keys: input, output, cache, total).""" - for entry in transcript or []: - if entry.get("type") != "ai": + for message in messages or []: + if getattr(message, "type", "") != "ai": continue - uses = entry.get("uses") or {} - totals["input"] += uses.get("input_token", 0) - totals["output"] += uses.get("output_token", 0) - totals["cache"] += uses.get("cache_token", 0) - totals["total"] += uses.get("total_token", 0) + uses = uses_from_usage_metadata(getattr(message, "usage_metadata", None)) + totals["input"] += uses["input_token"] + totals["output"] += uses["output_token"] + totals["cache"] += uses["cache_token"] + totals["total"] += uses["total_token"] def chunks_from_transcript(transcript: list[dict]) -> list[str]: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/response.py b/fastapi_startkit/src/fastapi_startkit/ai/response.py deleted file mode 100644 index 265e5180..00000000 --- a/fastapi_startkit/src/fastapi_startkit/ai/response.py +++ /dev/null @@ -1,76 +0,0 @@ -from __future__ import annotations - -import json -import os -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, Any - -if TYPE_CHECKING: - from .agent import Agent - - -@dataclass -class AgentResponse: - 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 - tool_events: list[dict] = field(default_factory=list) - transcript: list[dict] = field(default_factory=list) - - @property - def runtime_ms(self) -> float: - return self.runtime * 1000 - - def text(self) -> str: - return self.content - - def json(self) -> Any: - return json.loads(self.content) - - def __str__(self) -> str: - return self.content - - def __bool__(self) -> bool: - return bool(self.content) - - -@dataclass -class AgentSnapshot: - path: str - - def exists(self) -> bool: - return os.path.exists(self.path) - - def load(self) -> AgentResponse: - 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: - os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True) - with open(self.path, "w") as f: - json.dump( - { - "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: - if self.exists(): - return self.load() - 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 a62604ca..e00273ef 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -12,7 +12,6 @@ from . import recording from .pipeline import Response, build_pipeline -from .response import AgentResponse if TYPE_CHECKING: from .agent import Agent @@ -25,13 +24,18 @@ def _as_text(content: Any) -> str: return content if isinstance(content, str) else str(content) +def _stream_event(event: str, name: str, run_id: str, data: dict) -> dict: + """A StandardStreamEvent dict, the shape LangChain's astream_events yields.""" + return {"event": event, "name": name, "run_id": run_id, "tags": [], "metadata": {}, "parent_ids": [], "data": data} + + class BaseRunner(ABC): def __init__(self, agent: Agent) -> None: self.agent = agent self._reset_capture() # Populated by stream() so the record-and-replay harness can persist the # same structured turn a buffered prompt() produces. - self.last_response: AgentResponse | None = None + self.last_response: dict | None = None def _reset_capture(self) -> None: """Clear the per-turn interaction captured across the model/tool calls.""" @@ -78,6 +82,44 @@ def _record_tool_message(self, call: ToolCall, message: BaseMessage, response_ti } ) + def _turn_messages(self, message: str) -> list[BaseMessage]: + """Rebuild the turn as LangChain messages — [HumanMessage, AIMessage, + ToolMessage, ...] — mirroring create_agent().invoke() output. Each AI and + tool message carries its ``response_time`` (ms) in ``additional_kwargs``.""" + from langchain_core.messages import HumanMessage, ToolMessage # noqa: PLC0415 + + messages: list[BaseMessage] = [HumanMessage(content=message)] if message else [] + tool_events = iter(self._tool_events) + for entry in self._transcript: + if entry["type"] == "ai": + uses = entry.get("uses") or {} + messages.append( + AIMessage( + content=entry.get("content", ""), + tool_calls=entry.get("tool_calls") or [], + additional_kwargs={"response_time": entry.get("response_time", 0.0)}, + usage_metadata={ + "input_tokens": uses.get("input_token", 0), + "output_tokens": uses.get("output_token", 0), + "total_tokens": uses.get("total_token", 0), + "input_token_details": {"cache_read": uses.get("cache_token", 0)}, + } + if uses + else None, + ) + ) + elif entry["type"] == "tool_response": + event = next(tool_events, {}) + messages.append( + ToolMessage( + content=entry.get("content", ""), + name=event.get("name"), + tool_call_id=event.get("id") or "", + additional_kwargs={"response_time": entry.get("response_time", 0.0)}, + ) + ) + return messages + async def _build_messages(self, message: str, attachments: list[Document] | None = None) -> list[dict]: agent = self.agent messages: list[dict] = [] @@ -109,11 +151,13 @@ def _build_model(self, model: str | None = None, provider_options: dict | None = return Ai().get_model_for(self.agent, model, provider_options) - def _apply_schema(self, response: AgentResponse) -> AgentResponse: + def _apply_schema(self, state: dict) -> dict: + from . import state as agent_state # noqa: PLC0415 + 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 + if schema is not None and state.get("structured_response") is None and agent_state.text(state): + state["structured_response"] = self._build_schema(schema, agent_state.text(state)) + return state @staticmethod def _build_schema(schema: Any, content: str) -> Any: @@ -133,7 +177,7 @@ async def run( model: str | None = None, attachments: list[Document] | None = None, provider_options: dict | None = None, - ) -> AgentResponse: ... + ) -> dict: ... @abstractmethod def stream( @@ -153,22 +197,24 @@ async def run( model: str | None = None, attachments: list[Document] | None = None, provider_options: dict | None = None, - ) -> AgentResponse: - started = time.perf_counter() + ) -> dict: self._reset_capture() messages = await self._build_messages(message, attachments) model = self._build_model(model, provider_options) - response = await self._run_pipeline(model, messages) - response = self._apply_schema(response) - response.runtime = time.perf_counter() - started - return response + parsed = await self._run_pipeline(model, messages) + state: dict = {"messages": self._turn_messages(message)} + if parsed is not None: + state["structured_response"] = parsed + return self._apply_schema(state) - async def _run_pipeline(self, chat_model: Any, messages: list) -> AgentResponse: + async def _run_pipeline(self, chat_model: Any, messages: list) -> Any: + """Run the turn through the middleware pipeline; the interaction lands on + the capture buffers. Returns the parsed structured output, if any.""" chain = list(self.agent.middleware()) if not chain: raw = await self._invoke(chat_model, messages) - return self._to_agent_response(raw) + return self._parsed_of(raw) def core(model: Any) -> Response: async def _run() -> AsyncIterator[Any]: @@ -178,42 +224,14 @@ async def _run() -> AsyncIterator[Any]: 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() - - # Prefer the tool calls the model actually requested (captured before the - # tool ran); ``final`` is the tool-result message, which carries none. - if structured: - tool_calls: list[dict] = [] - usage: dict[str, Any] = {} - else: - tool_calls = self._requested_tool_calls or list(getattr(final, "tool_calls", None) or []) - usage = self._usage or self._usage_from_message(final) - - return AgentResponse( - content=content, - tool_calls=tool_calls, - usage=usage, - raw=result, - parsed=parsed, - tool_events=list(self._tool_events), - transcript=list(self._transcript), - ) + return self._parsed_of(raw) + + @staticmethod + def _parsed_of(result: Any) -> Any: + # A structured-output model returns {"parsed": ..., "raw": AIMessage}. + if isinstance(result, dict) and "parsed" in result: + return result.get("parsed") + return None async def stream( self, @@ -235,16 +253,20 @@ def core(m: Any) -> Response: async for chunk in pipeline(chat_model): yield chunk - # Expose the same structured turn a buffered prompt() records, so the + # Expose the same structured turn a buffered prompt() returns, so the # record-and-replay harness can persist streamed runs identically. - self.last_response = recording.to_response(self._transcript) - self.last_response.transcript = list(self._transcript) + self.last_response = {"messages": self._turn_messages(message)} async def _invoke(self, model: Runnable[Any, BaseMessage], messages: list) -> BaseMessage: started = time.perf_counter() response: AIMessage = await model.ainvoke(list(messages)) # type: ignore[assignment] latency_ms = (time.perf_counter() - started) * 1000 if isinstance(response, dict) and "parsed" in response: + # Structured output: record the raw model turn so it lands on the + # state's messages alongside the parsed object. + raw = response.get("raw") + if raw is not None: + self._record_ai_message(raw, latency_ms) return response # type: ignore[return-value] self._record_ai_message(response, latency_ms) if not response.tool_calls: @@ -252,35 +274,48 @@ async def _invoke(self, model: Runnable[Any, BaseMessage], messages: list) -> Ba return (await self._run_tools(response.tool_calls))[-1] async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list) -> AsyncIterator[dict]: + """Yield the same StandardStreamEvent dicts LangChain's astream_events + produces: on_chat_model_(start|stream|end) for the model turn, then + on_tool_(start|end) around each tool execution.""" + import uuid # noqa: PLC0415 + started = time.perf_counter() + run_id = str(uuid.uuid4()) + model_name = type(model).__name__ + payload = {"messages": [list(messages)]} + yield _stream_event("on_chat_model_start", model_name, run_id, {"input": payload}) + gathered: AIMessageChunk | None = None chunks: list[str] = [] async for chunk in model.astream(list(messages)): if chunk.content: - text = _as_text(chunk.content) - chunks.append(text) - yield {"type": "delta", "text": text} + chunks.append(_as_text(chunk.content)) + yield _stream_event("on_chat_model_stream", model_name, run_id, {"chunk": chunk}) gathered = chunk if gathered is None else gathered + chunk # type: ignore[operator] latency_ms = (time.perf_counter() - started) * 1000 if gathered is None: return + yield _stream_event("on_chat_model_end", model_name, run_id, {"output": gathered, "input": payload}) self._record_ai_message(gathered, latency_ms, chunks=chunks) if not gathered.tool_calls: return - for call, message in zip(gathered.tool_calls, await self._run_tools(gathered.tool_calls)): - yield {"type": "tool_response", "name": call["name"], "content": _as_text(message.content)} + for call in gathered.tool_calls: + tool_run_id = str(uuid.uuid4()) + yield _stream_event("on_tool_start", call["name"], tool_run_id, {"input": call.get("args", {})}) + message = await self._run_tool(call) + yield _stream_event("on_tool_end", call["name"], tool_run_id, {"output": message}) async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: + return [await self._run_tool(call) for call in tool_calls] + + async def _run_tool(self, call: ToolCall) -> 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 - started = time.perf_counter() - message = await selected.ainvoke(call) - self._record_tool_message(call, message, (time.perf_counter() - started) * 1000) - results.append(message) - return results + try: + selected = tools[call["name"]] + except KeyError: + raise ValueError(f"Agent has no tool named {call['name']!r}") from None + started = time.perf_counter() + message = await selected.ainvoke(call) + self._record_tool_message(call, message, (time.perf_counter() - started) * 1000) + return message diff --git a/fastapi_startkit/src/fastapi_startkit/ai/state.py b/fastapi_startkit/src/fastapi_startkit/ai/state.py new file mode 100644 index 00000000..5c02da06 --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/state.py @@ -0,0 +1,88 @@ +"""Helpers over the ``{"messages": [...]}`` state a prompt() returns. + +An agent turn is the same shape langchain's ``create_agent().invoke()`` yields: +``HumanMessage``, then the model's ``AIMessage`` (carrying ``tool_calls``), then +one ``ToolMessage`` per tool result — plus ``"structured_response"`` when the +agent has a schema. AI/tool messages carry their ``response_time`` (ms) in +``additional_kwargs``. +""" + +from __future__ import annotations + +import json +from typing import Any + + +def messages(state: dict | None) -> list: + return (state or {}).get("messages") or [] + + +def text(state: dict | None) -> str: + """The turn's final content — the last message's text (a tool-run turn ends + on the tool's result, a plain turn on the model's answer).""" + turn = messages(state) + if not turn: + return "" + content = turn[-1].content + return content if isinstance(content, str) else str(content) + + +def tool_calls(state: dict | None) -> list[dict]: + """Every tool call the model requested across the turn's AI messages.""" + return [call for m in messages(state) if m.type == "ai" for call in (getattr(m, "tool_calls", None) or [])] + + +def tool_events(state: dict | None) -> list[dict]: + """Each executed tool as {name, args, id, content, content_type, response_time} — + the requested call joined with its ToolMessage result (by call id, falling + back to order for replayed turns that don't store ids).""" + pending = list(tool_calls(state)) + events: list[dict] = [] + for m in messages(state): + if m.type != "tool": + continue + call_id = getattr(m, "tool_call_id", "") or "" + call = next((c for c in pending if c.get("id") == call_id), None) or (pending[0] if pending else {}) + if call in pending: + pending.remove(call) + content = m.content if isinstance(m.content, str) else str(m.content) + events.append( + { + "name": call.get("name") or m.name, + "args": call.get("args", {}), + "id": call.get("id") or call_id or None, + "content": content, + "content_type": _content_type(content), + "response_time": (m.additional_kwargs or {}).get("response_time", 0.0), + } + ) + return events + + +def usage(state: dict | None) -> dict: + """Summed token usage across the turn's AI messages: {"input", "output"}.""" + totals = {"input": 0, "output": 0} + for m in messages(state): + meta = getattr(m, "usage_metadata", None) or {} + totals["input"] += meta.get("input_tokens", 0) + totals["output"] += meta.get("output_tokens", 0) + return totals + + +def runtime(state: dict | None) -> float: + """Seconds spent on the turn's model and tool calls, summed from each + message's recorded ``response_time`` (ms).""" + total_ms = sum((m.additional_kwargs or {}).get("response_time", 0.0) for m in messages(state)) + return total_ms / 1000 + + +def structured(state: dict | None) -> Any: + return (state or {}).get("structured_response") + + +def _content_type(content: str) -> str: + try: + json.loads(content) + except (ValueError, TypeError): + return "text" + return "json" diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index cbd4a634..79f659da 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -11,7 +11,8 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Callable -from .response import AgentResponse +from . import recording +from . import state as agent_state if TYPE_CHECKING: from .agent import Agent @@ -23,16 +24,44 @@ def _joined(value: Any) -> str: def _frame_text(frame: Any) -> str: - """Text carried by a stream frame; cassettes keep storing plain strings.""" + """Text carried by a stream event; cassettes keep storing plain strings.""" if isinstance(frame, dict): - return frame.get("text") or frame.get("content") or "" + if frame.get("event") == "on_chat_model_stream": + chunk = (frame.get("data") or {}).get("chunk") + text = getattr(chunk, "content", "") or "" + return text if isinstance(text, str) else str(text) + if "event" in frame: # start/end events carry no streamable text + return "" + return frame.get("text") or frame.get("content") or "" # legacy frames return str(frame) +def _chunk_event(text: str) -> dict: + """A replayed on_chat_model_stream event, shaped like astream_events yields.""" + from langchain_core.messages import AIMessageChunk # noqa: PLC0415 + + return { + "event": "on_chat_model_stream", + "name": "replay", + "run_id": "", + "tags": [], + "metadata": {}, + "parent_ids": [], + "data": {"chunk": AIMessageChunk(content=text)}, + } + + def _new_token_totals() -> dict: return {"input": 0, "output": 0, "cache": 0, "total": 0} +def _text_state(content: str, tool_calls: list | None = None) -> dict: + """A minimal one-message state for legacy/fallback shapes.""" + from langchain_core.messages import AIMessage # noqa: PLC0415 + + return {"messages": [AIMessage(content=content, tool_calls=tool_calls or [])]} + + class TokenQuery: """Fluent predicate over accumulated token totals, e.g.:: @@ -78,7 +107,7 @@ def __init__(self, agent_cls: type[Agent], responses: list) -> None: 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_response: dict | None = None self.last_elapsed: float | None = None self._tokens: dict = _new_token_totals() self._response_time: float = 0.0 @@ -104,7 +133,7 @@ def __exit__(self, *_exc: Any) -> bool: async def prompt( self, message: str, *, attachments: list[Document] | None = None, config: dict | None = None - ) -> AgentResponse: + ) -> dict: start = time.monotonic() extra = {"config": config} if config is not None else {} self._last_response = await self._agent.prompt(message, attachments=attachments, **extra) @@ -119,27 +148,22 @@ async def stream(self, message: str, *, config: dict | None = None) -> AsyncIter # captures while streaming (content + tool_calls), not just joined text. runner = self._agent.runner() async for frame in runner.stream(message, **extra): - chunks.append(_frame_text(frame)) + if text := _frame_text(frame): + chunks.append(text) yield frame - self._last_response = getattr(runner, "last_response", None) or AgentResponse(content="".join(chunks)) + self._last_response = getattr(runner, "last_response", None) or _text_state("".join(chunks)) self._remember(message, self._last_response) - def _remember(self, message: str, response: AgentResponse) -> None: - self._accumulate_tokens(response) - self._response_time += response.runtime + def _remember(self, message: str, state: dict) -> None: + self._accumulate_tokens(state) + self._response_time += agent_state.runtime(state) self._records.append({"role": "user", "content": message}) - self._records.append({"role": "assistant", "content": response.content}) - - def _accumulate_tokens(self, response: AgentResponse) -> None: - """Add a turn's token usage to the running totals. Prefer the per-message - ``uses`` on the transcript; fall back to the response's summary usage.""" - from . import recording # noqa: PLC0415 + self._records.append({"role": "assistant", "content": agent_state.text(state)}) - if response.transcript: - recording.accumulate_uses(response.transcript, self._tokens) - else: - self._tokens["input"] += response.usage.get("input", 0) - self._tokens["output"] += response.usage.get("output", 0) + def _accumulate_tokens(self, state: dict) -> None: + """Add a turn's token usage — read from each AI message's usage_metadata — + to the running totals.""" + recording.accumulate_uses(agent_state.messages(state), self._tokens) def assert_tokens(self, predicate: Callable[[TokenQuery], Any]) -> None: """Assert on the tokens accumulated across every prompt()/stream() so far. @@ -162,12 +186,11 @@ def assert_prompt(self, expected: str | Callable[[str], bool]) -> None: ) 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}" + content = agent_state.text(self._require_response()) + assert expected in content, f"Expected response to contain {expected!r}, got {content!r}" def assert_tool_call(self, name: str) -> None: - response = self._require_response() - called = [tc.get("name") for tc in response.tool_calls] + called = [tc.get("name") for tc in agent_state.tool_calls(self._require_response())] assert name in called, f"Expected tool {name!r} to be called, but got: {called}" def assert_prompted(self, times: int | None = None) -> None: @@ -186,20 +209,20 @@ def reset(self) -> "AgentFake": self._response_time = 0.0 return self - def _require_response(self) -> AgentResponse: + def _require_response(self) -> dict: 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] + return [tc.get("name", "") for tc in agent_state.tool_calls(self._require_response())] def assert_text_response(self) -> None: - response = self._require_response() - assert response.content, "Expected a non-empty text response, but content was empty." + assert agent_state.text(self._require_response()), ( + "Expected a non-empty text response, but content was empty." + ) def assert_tool_called(self, name: str, predicate: Callable[[AssertToolCall], bool] | None = None) -> None: - response = self._require_response() - matches = [tc for tc in response.tool_calls if tc.get("name") == name] + matches = [tc for tc in agent_state.tool_calls(self._require_response()) 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(AssertToolCall(tc)) for tc in matches), ( @@ -222,11 +245,11 @@ def assert_response_time_lt(self, seconds: float) -> None: assert total < seconds, f"Expected total response time < {seconds}s, took {total:.3f}s" async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None: - response = self._require_response() - verdict = await self._judge(model, expectation, response.content, provider) + content = agent_state.text(self._require_response()) + verdict = await self._judge(model, expectation, 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}" + f"{verdict.get('reasoning', '')!r} — response was {content!r}" ) async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: @@ -276,7 +299,7 @@ def __init__(self, real: Agent, cassette: str | None = None, messages: list | No self._seed_messages: list = list(messages or []) self._records: list[dict] = [] self._real.messages = self._history # type: ignore[method-assign] - self._last_response: AgentResponse | None = None + self._last_response: dict | None = None self.last_elapsed: float | None = None self._tokens: dict = _new_token_totals() self._response_time: float = 0.0 @@ -312,100 +335,85 @@ 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)) - def _turn(self, message: str, response: AgentResponse, chunks: list[str] | None = None) -> list[dict]: + def _turn(self, message: str, state: dict, chunks: list[str] | None = None) -> list[dict]: """Build the ordered per-turn transcript persisted to the cassette: - the human input followed by the interaction the runner captured. When a - stream did not capture a structured transcript, ``chunks`` are attached - to the synthesized answer so replay can re-emit them.""" - from . import recording # noqa: PLC0415 - + the human input followed by the turn's AI/tool messages. Streamed + ``chunks`` are attached to the last model entry so replay can re-emit + them.""" entries: list[dict] = [recording.human(message)] - if response.transcript: - entries.extend(response.transcript) - else: - entries.append( - recording.ai( - content=response.content, - tool_calls=list(response.tool_calls), - uses=recording.uses_from_agent_usage(response.usage), - response_time=response.runtime * 1000, - chunks=chunks, - ) - ) + entries.extend(recording.entries_from_messages(agent_state.messages(state))) + if chunks: + for entry in reversed(entries): + if entry["type"] == "ai": + entry["chunks"] = chunks + break return entries @staticmethod - def _response_from_cache(value: Any) -> AgentResponse: - from . import recording # noqa: PLC0415 - + def _state_from_cache(value: Any) -> dict: if recording.is_transcript(value): - return recording.to_response(value) + return recording.to_state(value) if isinstance(value, dict) and "content" in value: # legacy flat shape - return AgentResponse( - content=_joined(value.get("content", "")), - tool_calls=value.get("tool_calls") or [], - usage=value.get("usage") or {}, - ) - return AgentResponse(content=_joined(value)) + return _text_state(_joined(value.get("content", "")), value.get("tool_calls") or []) + return _text_state(_joined(value)) - def _remember_turn(self, message: str, response: AgentResponse) -> None: - self._accumulate_tokens(response) - self._response_time += response.runtime + def _remember_turn(self, message: str, state: dict) -> None: + self._accumulate_tokens(state) + self._response_time += agent_state.runtime(state) 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 + turn: dict[str, Any] = {"role": "assistant", "content": agent_state.text(state)} + if agent_state.tool_calls(state): + turn["tool_calls"] = agent_state.tool_calls(state) self._records.append(turn) async def prompt( self, message: str, *, attachments: list[Document] | None = None, config: dict | None = None - ) -> AgentResponse: + ) -> dict: cassette, store = self._load() key = self._key(message, attachments) start = time.monotonic() if key in store: - response = self._response_from_cache(store[key]) + state = self._state_from_cache(store[key]) else: extra = {"config": config} if config is not None else {} - response = await self._real.prompt(message, attachments=attachments, **extra) - self._save(cassette, store, key, self._turn(message, response)) - response = self._real.runner()._apply_schema(response) + state = await self._real.prompt(message, attachments=attachments, **extra) + self._save(cassette, store, key, self._turn(message, state)) + state = self._real.runner()._apply_schema(state) self.last_elapsed = time.monotonic() - start - self._last_response = response - self._remember_turn(message, response) - return response + self._last_response = state + self._remember_turn(message, state) + return state async def stream(self, message: str, *, config: dict | None = None) -> AsyncIterator[dict]: - from . import recording # noqa: PLC0415 - cassette, store = self._load() key = self._key(message, None) if key in store: value = store[key] if recording.is_transcript(value): chunks = recording.chunks_from_transcript(value) - response = self._response_from_cache(value) + state = self._state_from_cache(value) elif isinstance(value, dict): # legacy flat shape with buffered chunks chunks = value.get("chunks", []) - response = self._response_from_cache(value) + state = self._state_from_cache(value) else: # legacy cassette: a bare list of text chunks chunks = value if isinstance(value, list) else [value] - response = AgentResponse(content=_joined(chunks)) + state = _text_state(_joined(chunks)) for chunk in chunks: - yield {"type": "delta", "text": chunk} + yield _chunk_event(chunk) else: extra = {"config": config} if config is not None else {} - # Drive the runner directly to capture the structured response - # (content + tool_calls + chunks) it records while streaming. + # Drive the runner directly to capture the structured state + # (messages + chunks) it records while streaming. runner = self._real.runner() chunks = [] async for frame in runner.stream(message, **extra): - chunks.append(_frame_text(frame)) + if text := _frame_text(frame): + chunks.append(text) yield frame - response = getattr(runner, "last_response", None) or AgentResponse(content=_joined(chunks)) - self._save(cassette, store, key, self._turn(message, response, chunks=chunks)) - self._last_response = response - self._remember_turn(message, response) + state = getattr(runner, "last_response", None) or _text_state(_joined(chunks)) + self._save(cassette, store, key, self._turn(message, state, chunks=chunks)) + self._last_response = state + self._remember_turn(message, state) async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: cassette, store = self._load() diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 9c77ea1a..0c82f21d 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -11,10 +11,10 @@ from langchain_core.tools import tool from fastapi_startkit.ai import AIConfig, Document +from fastapi_startkit.ai import state as ai_state 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 @@ -59,14 +59,15 @@ 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): + async def test_prompt_returns_a_messages_state(self): self.setup_agent([AIMessage(content="hello back")]) agent = Agent() result = await agent.prompt("hi there") - self.assertIsInstance(result, AgentResponse) - self.assertEqual(result.content, "hello back") + self.assertIsInstance(result, dict) + self.assertEqual([m.type for m in result["messages"]], ["human", "ai"]) + self.assertEqual(ai_state.text(result), "hello back") async def test_search_jobs_tool_returns_listing(self): self.setup_agent( @@ -81,8 +82,8 @@ async def test_search_jobs_tool_returns_listing(self): result = await JobAssistant().prompt("find me a python job") - self.assertEqual(result.content, "Python Developer at Shopify") - self.assertEqual([tc["name"] for tc in result.tool_calls], ["search_jobs"]) + self.assertEqual(ai_state.text(result), "Python Developer at Shopify") + self.assertEqual([tc["name"] for tc in ai_state.tool_calls(result)], ["search_jobs"]) async def test_prompt_maps_usage_metadata(self): self.setup_agent( @@ -91,7 +92,7 @@ async def test_prompt_maps_usage_metadata(self): result = await Agent().prompt("anything") - self.assertEqual(result.usage, {"input": 11, "output": 7}) + self.assertEqual(ai_state.usage(result), {"input": 11, "output": 7}) async def test_build_model_passes_langchain_provider_key(self): captured = {} @@ -113,13 +114,17 @@ class GoogleAgent(Agent): self.assertEqual(captured["provider"], "google_genai") self.assertEqual(captured["model"], "gemini-2.5-flash-lite") - async def test_stream_yields_tokens_from_the_model(self): + async def test_stream_yields_astream_events_shaped_events(self): self.setup_agent([AIMessage(content="streamed reply")]) - frames = [frame async for frame in Agent().stream("hello")] + events = [event async for event in Agent().stream("hello")] - self.assertTrue(all(frame["type"] == "delta" for frame in frames)) - self.assertEqual("".join(frame["text"] for frame in frames), "streamed reply") + kinds = [event["event"] for event in events] + self.assertEqual(kinds[0], "on_chat_model_start") + self.assertEqual(kinds[-1], "on_chat_model_end") + self.assertTrue(all({"event", "name", "run_id", "data", "parent_ids"} <= set(e) for e in events)) + text = "".join(e["data"]["chunk"].text for e in events if e["event"] == "on_chat_model_stream") + self.assertEqual(text, "streamed reply") async def test_middleware_streams_token_by_token_and_runs_after_hook(self): events: list = [] @@ -135,7 +140,8 @@ def middleware(self): self.setup_agent([AIMessage(content="one two three")], LoggedAgent) - chunks = [frame["text"] async for frame in LoggedAgent().stream("hi")] + stream = [event async for event in LoggedAgent().stream("hi")] + chunks = [e["data"]["chunk"].text for e in stream if e["event"] == "on_chat_model_stream" and e["data"]["chunk"].text] # Middleware must not buffer: the model's tokens arrive as separate chunks... self.assertEqual("".join(chunks), "one two three") @@ -159,7 +165,7 @@ def middleware(self): result = await LoggedAgent().prompt("hi") - self.assertEqual(result.content, "done") + self.assertEqual(ai_state.text(result), "done") self.assertEqual(events, ["before", "after"]) async def test_stream_yields_tool_result_without_calling_model_again(self): @@ -175,12 +181,15 @@ async def test_stream_yields_tool_result_without_calling_model_again(self): ) self.addCleanup(Ai.reset_fakes) - frames = [frame async for frame in JobAssistant().stream("find me a python job")] + events = [event async for event in JobAssistant().stream("find me a python job")] - self.assertEqual( - frames, - [{"type": "tool_response", "name": "search_jobs", "content": "Python Developer at Shopify"}], - ) + kinds = [event["event"] for event in events] + # One model turn only — the tool result is not fed back for a second call. + self.assertEqual(kinds.count("on_chat_model_start"), 1) + self.assertEqual(kinds[-2:], ["on_tool_start", "on_tool_end"]) + tool_end = events[-1] + self.assertEqual(tool_end["name"], "search_jobs") + self.assertEqual(tool_end["data"]["output"].content, "Python Developer at Shopify") def test_resolve_model_falls_back_to_lab_default(self): self.assertEqual(Ai()._resolve_model(Agent()), "gemini-2.5-flash-lite") diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 4ea8d6c5..3f1ed36d 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -4,9 +4,11 @@ import unittest from unittest import mock +from langchain_core.messages import AIMessage + +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.agent import Agent from fastapi_startkit.ai.ai import Ai -from fastapi_startkit.ai.response import AgentResponse class SimpleAgent(Agent): @@ -22,7 +24,7 @@ async def test_fake_replays_the_only_response(self): with SimpleAgent.fake(["Hello world!"]): result = await agent.prompt("anything") - self.assertEqual(result.content, "Hello world!") + self.assertEqual(ai_state.text(result), "Hello world!") async def test_fake_does_not_call_provider_build(self): import langchain.chat_models as chat_models @@ -44,8 +46,8 @@ async def test_fake_replays_responses_in_order(self): first = await agent.prompt("call one") second = await agent.prompt("call two") - self.assertEqual(first.content, "first reply") - self.assertEqual(second.content, "second reply") + self.assertEqual(ai_state.text(first), "first reply") + self.assertEqual(ai_state.text(second), "second reply") async def test_fake_raises_once_responses_are_exhausted(self): agent = SimpleAgent() @@ -68,7 +70,7 @@ async def run(): result = await run() - self.assertEqual(result.content, "decorated reply") + self.assertEqual(ai_state.text(result), "decorated reply") async def test_assert_prompted_passes_after_one_call(self): with SimpleAgent.fake(["ok"]) as agent: @@ -128,23 +130,25 @@ async def test_assert_not_prompted_passes_after_reset(self): async def test_fake_rebinding_overrides_previous(self): with SimpleAgent.fake(["first fake"]) as agent: - self.assertEqual((await agent.prompt("call")).content, "first fake") + self.assertEqual(ai_state.text(await agent.prompt("call")), "first fake") with SimpleAgent.fake(["second fake"]) as agent: - self.assertEqual((await agent.prompt("call again")).content, "second fake") + self.assertEqual(ai_state.text(await agent.prompt("call again")), "second fake") async def test_stream_returns_fake_response(self): with SimpleAgent.fake(["Faked stream!"]) as agent: - chunks = [frame["text"] async for frame in agent.stream("hello world")] + events = [e async for e in agent.stream("hello world")] + chunks = [e["data"]["chunk"].text for e in events if e["event"] == "on_chat_model_stream"] 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): with SimpleAgent.fake(["Hello there, friend"]) as agent: - chunks = [frame["text"] async for frame in agent.stream("hi")] + events = [e async for e in agent.stream("hi")] + chunks = [e["data"]["chunk"].text for e in events if e["event"] == "on_chat_model_stream"] self.assertEqual("".join(chunks), "Hello there, friend") async def test_stream_records_one_call_not_two(self): @@ -161,7 +165,7 @@ def setup_agent(self, content): async def fake_run(agent_self, message, **kwargs): calls.append(message) - return AgentResponse(content=content) + return {"messages": [AIMessage(content=content)]} patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() @@ -175,7 +179,7 @@ async def test_first_run_records_response_to_cassette(self): with SimpleAgent.record(cassette) as agent: result = await agent.prompt("hello") - self.assertEqual(result.content, "recorded reply") + self.assertEqual(ai_state.text(result), "recorded reply") self.assertEqual(calls, ["hello"]) self.assertTrue(os.path.exists(cassette)) with open(cassette) as f: @@ -195,15 +199,15 @@ async def test_second_run_replays_without_calling_run(self): with SimpleAgent.record(cassette) as agent: replayed = await agent.prompt("hello") - self.assertEqual(replayed.content, "recorded reply") + self.assertEqual(ai_state.text(replayed), "recorded reply") self.assertEqual(calls, ["hello"]) async def test_replay_prefers_cassette_over_live_response(self): async def first_run(s, m, **k): - return AgentResponse(content="from first record") + return {"messages": [AIMessage(content="from first record")]} async def changed_run(s, m, **k): - return AgentResponse(content="changed live value") + return {"messages": [AIMessage(content="changed live value")]} with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") @@ -214,7 +218,7 @@ async def changed_run(s, m, **k): with SimpleAgent.record(cassette) as agent: result = await agent.prompt("hello") - self.assertEqual(result.content, "from first record") + self.assertEqual(ai_state.text(result), "from first record") async def test_distinct_messages_are_recorded_separately(self): calls = self.setup_agent("reply") @@ -268,7 +272,7 @@ async def test_stream_second_run_replays_chunks_without_calling_stream(self): with SimpleAgent.record(cassette) as agent: [c async for c in agent.stream("hi")] with SimpleAgent.record(cassette) as agent: - replayed = [frame["text"] async for frame in agent.stream("hi")] + replayed = [e["data"]["chunk"].text async for e in agent.stream("hi")] self.assertEqual(replayed, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) # real stream invoked only on the first run @@ -282,4 +286,4 @@ async def test_prompt_reads_a_stream_recorded_cassette_as_joined_content(self): with SimpleAgent.record(cassette) as agent: response = await agent.prompt("hi") - self.assertEqual(response.content, "Hello!") + self.assertEqual(ai_state.text(response), "Hello!") diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py index 513a6aa5..b0371e7d 100644 --- a/fastapi_startkit/tests/ai/test_agent_record_fluent.py +++ b/fastapi_startkit/tests/ai/test_agent_record_fluent.py @@ -22,9 +22,9 @@ from langchain_core.messages import AIMessage, HumanMessage +from fastapi_startkit.ai import state as ai_state 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 @@ -36,6 +36,10 @@ def _tool_call(name: str, args: dict | None = None, call_id: str = "c1") -> dict return {"name": name, "args": args or {}, "id": call_id} +def _state(content: str, tool_calls: list | None = None) -> dict: + return {"messages": [AIMessage(content=content, tool_calls=tool_calls or [])]} + + class TestFluentPromptMechanics(unittest.IsolatedAsyncioTestCase): def setup_agent(self, responses: list): """responses: list of (content, tool_calls) tuples, consumed in call order.""" @@ -43,20 +47,20 @@ def setup_agent(self, responses: list): async def fake_run(agent_self, message, **kwargs): content, tool_calls = queue.pop(0) - return AgentResponse(content=content, tool_calls=tool_calls or []) + return _state(content, tool_calls) patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) - async def test_prompt_is_async_and_returns_agent_response(self): + async def test_prompt_is_async_and_returns_a_messages_state(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!") + self.assertIsInstance(response, dict) + self.assertEqual(ai_state.text(response), "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")])]) @@ -84,7 +88,7 @@ async def test_replaying_from_cassette_preserves_tool_calls(self): 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 []) + return _state(content, tool_calls) patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() @@ -115,7 +119,7 @@ async def test_fails_when_no_prompt_has_been_made(self): 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) + return _state("", tool_calls) patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() @@ -155,7 +159,7 @@ async def test_predicate_can_reject(self): 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) + return _state("Hello!", tool_calls) patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() @@ -180,7 +184,7 @@ async def test_fails_when_present(self): class TestAssertResponseTimeLt(unittest.IsolatedAsyncioTestCase): def setup_agent(self): async def fake_run(agent_self, message, **kwargs): - return AgentResponse(content="Hello!") + return _state("Hello!") patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() @@ -224,28 +228,28 @@ async def test_same_followup_text_with_different_seed_history_does_not_collide(s cassette = os.path.join(tmp, "shared.json") async def run_a(agent_self, message, **kwargs): - return AgentResponse(content="job list A") + return _state("job list A") with mock.patch.object(SimpleAgent, "prompt", 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") + return _state("job list B") seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] 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") - self.assertEqual(response_a.content, "job list A") - self.assertEqual(response_b.content, "job list B") + self.assertEqual(ai_state.text(response_a), "job list A") + self.assertEqual(ai_state.text(response_b), "job list B") class TestAssertResponseJudged(unittest.IsolatedAsyncioTestCase): def setup_agent(self, content): async def fake_run(agent_self, message, **kwargs): - return AgentResponse(content=content) + return _state(content) patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() @@ -336,7 +340,7 @@ class TestExistingRecordApiIsUnaffected(unittest.IsolatedAsyncioTestCase): async def test_bare_context_manager_prompt_still_works(self): async def fake_run(agent_self, message, **kwargs): - return AgentResponse(content="recorded reply") + return _state("recorded reply") with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") @@ -344,4 +348,4 @@ async def fake_run(agent_self, message, **kwargs): with SimpleAgent.record(cassette) as agent: result = await agent.prompt("hello") - self.assertEqual(result.content, "recorded reply") + self.assertEqual(ai_state.text(result), "recorded reply") diff --git a/fastapi_startkit/tests/ai/test_agent_response.py b/fastapi_startkit/tests/ai/test_agent_response.py deleted file mode 100644 index 0d203076..00000000 --- a/fastapi_startkit/tests/ai/test_agent_response.py +++ /dev/null @@ -1,104 +0,0 @@ -"""Tests for the AgentResponse dataclass.""" - -import unittest - -from fastapi_startkit.ai.response import AgentResponse - - -class TestAgentResponse(unittest.TestCase): - def test_text_returns_content(self): - response = AgentResponse(content="Hello, world!") - self.assertEqual(response.text(), "Hello, world!") - - def test_text_returns_empty_string_when_no_content(self): - response = AgentResponse() - self.assertEqual(response.text(), "") - - def test_text_returns_multiline_content(self): - content = "Line 1\nLine 2\nLine 3" - response = AgentResponse(content=content) - self.assertEqual(response.text(), content) - - def test_json_parses_content_as_json(self): - response = AgentResponse(content='{"key": "value", "number": 42}') - parsed = response.json() - self.assertEqual(parsed, {"key": "value", "number": 42}) - - def test_json_parses_list_content(self): - response = AgentResponse(content="[1, 2, 3]") - self.assertEqual(response.json(), [1, 2, 3]) - - def test_json_parses_nested_object(self): - response = AgentResponse(content='{"nested": {"a": 1}}') - self.assertEqual(response.json()["nested"]["a"], 1) - - def test_json_raises_on_invalid_content(self): - response = AgentResponse(content="not valid json") - with self.assertRaises(Exception): # json.JSONDecodeError - response.json() - - def test_json_raises_on_empty_content(self): - response = AgentResponse(content="") - with self.assertRaises(Exception): - response.json() - - def test_str_returns_content(self): - response = AgentResponse(content="My response text") - self.assertEqual(str(response), "My response text") - - def test_str_returns_empty_string_when_no_content(self): - response = AgentResponse() - self.assertEqual(str(response), "") - - def test_str_works_in_f_string(self): - response = AgentResponse(content="hello") - self.assertEqual(f"Result: {response}", "Result: hello") - - def test_bool_is_true_when_content_non_empty(self): - response = AgentResponse(content="some text") - self.assertIs(bool(response), True) - - def test_bool_is_false_when_content_empty(self): - response = AgentResponse(content="") - self.assertIs(bool(response), False) - - def test_bool_is_false_when_content_not_set(self): - response = AgentResponse() - self.assertIs(bool(response), False) - - def test_bool_is_true_with_whitespace_content(self): - """A response with only whitespace still evaluates as truthy (non-empty string).""" - response = AgentResponse(content=" ") - self.assertIs(bool(response), True) - - def test_bool_usable_in_conditional(self): - response = AgentResponse(content="text") - self.assertTrue(response) - - empty = AgentResponse(content="") - self.assertFalse(empty) - - def test_tool_calls_default_to_empty_list(self): - response = AgentResponse() - self.assertEqual(response.tool_calls, []) - - def test_usage_defaults_to_empty_dict(self): - response = AgentResponse() - self.assertEqual(response.usage, {}) - - def test_raw_defaults_to_none(self): - response = AgentResponse() - self.assertIsNone(response.raw) - - def test_all_fields_can_be_set(self): - raw_obj = object() - response = AgentResponse( - content="text", - tool_calls=[{"name": "search", "input": {"q": "test"}}], - usage={"input": 10, "output": 20}, - raw=raw_obj, - ) - self.assertEqual(response.content, "text") - self.assertEqual(response.tool_calls, [{"name": "search", "input": {"q": "test"}}]) - self.assertEqual(response.usage, {"input": 10, "output": 20}) - self.assertIs(response.raw, raw_obj) diff --git a/fastapi_startkit/tests/ai/test_agent_schema.py b/fastapi_startkit/tests/ai/test_agent_schema.py index 347abd7a..1c9e9461 100644 --- a/fastapi_startkit/tests/ai/test_agent_schema.py +++ b/fastapi_startkit/tests/ai/test_agent_schema.py @@ -5,9 +5,11 @@ from pydantic import BaseModel +from langchain_core.messages import AIMessage + +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.agent import Agent from fastapi_startkit.ai.ai import Ai -from fastapi_startkit.ai.response import AgentResponse class User(BaseModel): @@ -29,17 +31,18 @@ async def test_fake_json_is_built_into_the_schema(self): with UserAgent.fake(['{"id": "u-1", "name": "Alex"}']): response = await agent.prompt("get the user") - self.assertIsInstance(response.parsed, User) - self.assertEqual(response.parsed.id, "u-1") - self.assertEqual(response.parsed.name, "Alex") - self.assertEqual(response.content, '{"id": "u-1", "name": "Alex"}') + parsed = response["structured_response"] + self.assertIsInstance(parsed, User) + self.assertEqual(parsed.id, "u-1") + self.assertEqual(parsed.name, "Alex") + self.assertEqual(ai_state.text(response), '{"id": "u-1", "name": "Alex"}') async def test_no_schema_leaves_parsed_none(self): agent = Agent() with Agent.fake(['{"id": "u-1"}']): response = await agent.prompt("anything") - self.assertIsNone(response.parsed) + self.assertNotIn("structured_response", response) async def test_invalid_json_for_schema_raises(self): agent = UserAgent() @@ -49,7 +52,7 @@ async def test_invalid_json_for_schema_raises(self): async def test_record_stores_json_and_rebuilds_schema_on_replay(self): async def fake_run(self, message, **kwargs): - return AgentResponse(content='{"id": "u-9", "name": "Sam"}') + return {"messages": [AIMessage(content='{"id": "u-9", "name": "Sam"}')]} with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "user.json") @@ -59,5 +62,5 @@ async def fake_run(self, message, **kwargs): 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")) + self.assertEqual(recorded["structured_response"], User(id="u-9", name="Sam")) + self.assertEqual(replayed["structured_response"], User(id="u-9", name="Sam")) diff --git a/fastapi_startkit/tests/ai/test_agent_state.py b/fastapi_startkit/tests/ai/test_agent_state.py new file mode 100644 index 00000000..2cba5d21 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_agent_state.py @@ -0,0 +1,83 @@ +"""Tests for the state helpers over the ``{"messages": [...]}`` turn shape.""" + +import unittest + +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from fastapi_startkit.ai import state as ai_state + + +def _tool_turn() -> dict: + return { + "messages": [ + HumanMessage(content="find jobs"), + AIMessage( + content="", + tool_calls=[{"name": "job_search_tool", "args": {"q": "python"}, "id": "c1", "type": "tool_call"}], + additional_kwargs={"response_time": 12.0}, + usage_metadata={"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + ), + ToolMessage(content='[{"id": 2}]', tool_call_id="c1", additional_kwargs={"response_time": 5.0}), + ] + } + + +class TestText(unittest.TestCase): + def test_returns_the_last_messages_content(self): + self.assertEqual(ai_state.text(_tool_turn()), '[{"id": 2}]') + + def test_plain_turn_returns_the_ai_answer(self): + state = {"messages": [HumanMessage(content="hi"), AIMessage(content="Hello!")]} + self.assertEqual(ai_state.text(state), "Hello!") + + def test_empty_state_returns_empty_string(self): + self.assertEqual(ai_state.text({}), "") + self.assertEqual(ai_state.text(None), "") + + +class TestToolCalls(unittest.TestCase): + def test_collects_calls_from_ai_messages(self): + self.assertEqual([tc["name"] for tc in ai_state.tool_calls(_tool_turn())], ["job_search_tool"]) + + def test_plain_turn_has_no_calls(self): + self.assertEqual(ai_state.tool_calls({"messages": [AIMessage(content="hi")]}), []) + + +class TestToolEvents(unittest.TestCase): + def test_joins_the_call_with_its_tool_message(self): + (event,) = ai_state.tool_events(_tool_turn()) + self.assertEqual(event["name"], "job_search_tool") + self.assertEqual(event["args"], {"q": "python"}) + self.assertEqual(event["id"], "c1") + self.assertEqual(event["content"], '[{"id": 2}]') + self.assertEqual(event["content_type"], "json") + self.assertEqual(event["response_time"], 5.0) + + def test_falls_back_to_order_when_ids_are_missing(self): + state = { + "messages": [ + AIMessage(content="", tool_calls=[{"name": "echo", "args": {"x": "hi"}, "id": "c9"}]), + ToolMessage(content="echoed:hi", tool_call_id=""), + ] + } + (event,) = ai_state.tool_events(state) + self.assertEqual(event["name"], "echo") + self.assertEqual(event["args"], {"x": "hi"}) + self.assertEqual(event["content_type"], "text") + + +class TestUsageAndRuntime(unittest.TestCase): + def test_usage_sums_ai_messages(self): + self.assertEqual(ai_state.usage(_tool_turn()), {"input": 10, "output": 5}) + + def test_runtime_sums_response_times_in_seconds(self): + self.assertAlmostEqual(ai_state.runtime(_tool_turn()), 0.017) + + +class TestStructured(unittest.TestCase): + def test_returns_the_structured_response(self): + self.assertEqual(ai_state.structured({"structured_response": {"a": 1}}), {"a": 1}) + + def test_defaults_to_none(self): + self.assertIsNone(ai_state.structured({"messages": []})) + self.assertIsNone(ai_state.structured(None)) diff --git a/fastapi_startkit/tests/ai/test_ai_fake.py b/fastapi_startkit/tests/ai/test_ai_fake.py index 33a29247..2ada13f9 100644 --- a/fastapi_startkit/tests/ai/test_ai_fake.py +++ b/fastapi_startkit/tests/ai/test_ai_fake.py @@ -15,6 +15,7 @@ from langchain_core.messages import AIMessage, ToolCall from langchain_core.tools import tool +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.agent import Agent from fastapi_startkit.ai.ai import Ai from fastapi_startkit.application import app @@ -119,7 +120,7 @@ async def test_prompt_replays_the_registered_fake_model_reply(self): result = await SimpleAgent().prompt("hi there") - self.assertEqual(result.content, "faked via ai") + self.assertEqual(ai_state.text(result), "faked via ai") async def test_prompt_runs_a_faked_tool_call_end_to_end(self): Ai.fake( @@ -134,7 +135,7 @@ async def test_prompt_runs_a_faked_tool_call_end_to_end(self): result = await JobAssistant().prompt("find me a python job") - self.assertEqual(result.content, "Python Developer at Shopify") + self.assertEqual(ai_state.text(result), "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")]) diff --git a/fastapi_startkit/tests/ai/test_assert_response_time.py b/fastapi_startkit/tests/ai/test_assert_response_time.py index 5b988e95..c9fe80c0 100644 --- a/fastapi_startkit/tests/ai/test_assert_response_time.py +++ b/fastapi_startkit/tests/ai/test_assert_response_time.py @@ -16,7 +16,6 @@ from fastapi_startkit.ai import recording from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.response import AgentResponse class TimingAgent(Agent): @@ -62,8 +61,8 @@ async def _drive(): return _drive() -def _slow_response() -> AgentResponse: - return recording.to_response(_slow_turn_transcript()) +def _slow_response() -> dict: + return recording.to_state(_slow_turn_transcript()) class TestAssertResponseTime(unittest.IsolatedAsyncioTestCase): diff --git a/fastapi_startkit/tests/ai/test_assert_tokens.py b/fastapi_startkit/tests/ai/test_assert_tokens.py index 8cfe9447..dfa031ee 100644 --- a/fastapi_startkit/tests/ai/test_assert_tokens.py +++ b/fastapi_startkit/tests/ai/test_assert_tokens.py @@ -11,25 +11,25 @@ import unittest from unittest import mock -from fastapi_startkit.ai import recording +from langchain_core.messages import AIMessage + from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.response import AgentResponse class TokenAgent(Agent): pass -def _ai(input_tokens: int, output_tokens: int) -> dict: - return recording.ai( +def _ai(input_tokens: int, output_tokens: int, cache_tokens: int = 0) -> AIMessage: + return AIMessage( content="reply", - uses={ - "input_token": input_tokens, - "output_token": output_tokens, - "cache_token": 0, - "total_token": input_tokens + output_tokens, + additional_kwargs={"response_time": 1.0}, + usage_metadata={ + "input_tokens": input_tokens, + "output_tokens": output_tokens, + "total_tokens": input_tokens + output_tokens + cache_tokens, + "input_token_details": {"cache_read": cache_tokens}, }, - response_time=1.0, ) @@ -45,8 +45,8 @@ async def prompt(agent_self, message, **kwargs): class TestAssertTokens(unittest.IsolatedAsyncioTestCase): async def test_passes_when_accumulated_tokens_are_within_limits(self): responses = [ - AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=[_ai(100, 20)]), - AgentResponse(content="b", usage={"input": 200, "output": 30}, transcript=[_ai(200, 30)]), + {"messages": [_ai(100, 20)]}, + {"messages": [_ai(200, 30)]}, ] with tempfile.TemporaryDirectory() as tmp: with _fake_prompt(responses): @@ -57,7 +57,7 @@ async def test_passes_when_accumulated_tokens_are_within_limits(self): agent.assert_tokens(lambda x: x.where("input", "<=", 5000).where("output", "<=", 5000)) async def test_fails_when_a_limit_is_exceeded(self): - responses = [AgentResponse(content="a", usage={"input": 400, "output": 20}, transcript=[_ai(400, 20)])] + responses = [{"messages": [_ai(400, 20)]}] with tempfile.TemporaryDirectory() as tmp: with _fake_prompt(responses): with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: @@ -66,21 +66,15 @@ async def test_fails_when_a_limit_is_exceeded(self): agent.assert_tokens(lambda x: x.where("input", "<=", 300)) async def test_supports_cache_and_total_fields(self): - transcript = [ - recording.ai( - content="a", - uses={"input_token": 100, "output_token": 20, "cache_token": 40, "total_token": 160}, - response_time=1.0, - ) - ] + messages = [_ai(100, 20, cache_tokens=40)] with tempfile.TemporaryDirectory() as tmp: - with _fake_prompt([AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=transcript)]): + with _fake_prompt([{"messages": messages}]): with TokenAgent.record(os.path.join(tmp, "c.json")) as agent: await agent.prompt("a") agent.assert_tokens(lambda x: x.where("cache", "<=", 40).where("total", ">=", 160)) async def test_accumulates_from_the_cassette_on_replay(self): - responses = [AgentResponse(content="a", usage={"input": 100, "output": 20}, transcript=[_ai(100, 20)])] + responses = [{"messages": [_ai(100, 20)]}] with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with _fake_prompt(responses): diff --git a/fastapi_startkit/tests/ai/test_graph_agent.py b/fastapi_startkit/tests/ai/test_graph_agent.py index 236fd1f1..7791c225 100644 --- a/fastapi_startkit/tests/ai/test_graph_agent.py +++ b/fastapi_startkit/tests/ai/test_graph_agent.py @@ -11,6 +11,7 @@ from langchain_core.tools import tool from langgraph.graph import END, START, StateGraph +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.graph import GraphAgent, GraphRunner, AgentState @@ -53,23 +54,27 @@ 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)) + self.assertEqual(ai_state.text(response), "The answer is 7") + self.assertTrue(any(call["name"] == "add" for call in ai_state.tool_calls(response))) 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") + add_call = next(call for call in ai_state.tool_calls(response) 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 = [frame["text"] async for frame in MathAgent().stream("What is 3 plus 4?")] + events = [e async for e in MathAgent().stream("What is 3 plus 4?")] + chunks = [e["data"]["chunk"].text for e in events if e["event"] == "on_chat_model_stream" and e["data"]["chunk"].text] self.assertEqual("".join(chunks), "The answer is 7") self.assertGreater(len(chunks), 1) + # The graph emits real astream_events: a root chain wraps the model turn. + self.assertEqual(events[0]["event"], "on_chain_start") + self.assertEqual(events[-1]["event"], "on_chain_end") async def test_records_the_prompt_call(self): with MathAgent.fake(["7"]) as agent: @@ -86,7 +91,7 @@ def middleware(self): with MiddlewareAgent.fake(["The answer is 7"]): response = await MiddlewareAgent().prompt("What is 3 plus 4?") - self.assertEqual(response.content, "The answer is 7") + self.assertEqual(ai_state.text(response), "The answer is 7") self.assertEqual(recorder.before, 1, "middleware handle() was never invoked") self.assertEqual(recorder.after, 1, "middleware after-hook never fired") diff --git a/fastapi_startkit/tests/ai/test_graph_recording.py b/fastapi_startkit/tests/ai/test_graph_recording.py index fa3f59de..a3576ae8 100644 --- a/fastapi_startkit/tests/ai/test_graph_recording.py +++ b/fastapi_startkit/tests/ai/test_graph_recording.py @@ -11,6 +11,7 @@ from langchain_core.tools import tool from langgraph.graph import END, START, StateGraph +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.graph import AgentState, GraphAgent, GraphRunner @@ -45,21 +46,22 @@ class TestGraphRunnerCapture(unittest.IsolatedAsyncioTestCase): def tearDown(self): Ai.reset_fakes() - async def test_transcript_records_the_full_tool_loop(self): + async def test_messages_record_the_full_tool_loop(self): with MathAgent.fake(list(SCRIPT)): response = await MathAgent().prompt("Add 3 and 4.") - kinds = [entry["type"] for entry in response.transcript] - self.assertEqual(kinds, ["ai", "tool_response", "ai"]) + kinds = [m.type for m in response["messages"]] + self.assertEqual(kinds, ["human", "ai", "tool", "ai"]) async def test_tool_response_and_final_answer_are_captured(self): with MathAgent.fake(list(SCRIPT)): response = await MathAgent().prompt("Add 3 and 4.") - tool_entry = next(e for e in response.transcript if e["type"] == "tool_response") - self.assertEqual(tool_entry["content"], "7") - self.assertGreaterEqual(tool_entry["response_time"], 0) + tool_message = next(m for m in response["messages"] if m.type == "tool") + self.assertEqual(tool_message.content, "7") + self.assertGreaterEqual(tool_message.additional_kwargs["response_time"], 0) - self.assertEqual(response.transcript[-1].get("content"), "The answer is 7") - self.assertEqual(len(response.tool_events), 1) - self.assertEqual(response.tool_events[0]["name"], "add") + self.assertEqual(response["messages"][-1].content, "The answer is 7") + events = ai_state.tool_events(response) + self.assertEqual(len(events), 1) + self.assertEqual(events[0]["name"], "add") diff --git a/fastapi_startkit/tests/ai/test_judge_agent.py b/fastapi_startkit/tests/ai/test_judge_agent.py index bcb71569..98230d1b 100644 --- a/fastapi_startkit/tests/ai/test_judge_agent.py +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -2,7 +2,7 @@ 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. +path (``state["structured_response"]``) — no hand-rolled verdict parsing. """ import os @@ -12,9 +12,9 @@ from langchain_core.messages import AIMessage +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.judge import JudgeAgent, Verdict from fastapi_startkit.ai.runner import Runner -from fastapi_startkit.ai.response import AgentResponse class TestJudgeAgent(unittest.IsolatedAsyncioTestCase): @@ -68,7 +68,7 @@ def test_model_and_provider_default_to_agent_defaults(self): 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"}') + return {"messages": [AIMessage(content='{"passed": true, "reasoning": "ok"}')]} with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "judge.json") @@ -76,5 +76,5 @@ async def fake_run(agent_self, message, **kwargs): with JudgeAgent.record(cassette) as agent: response = await agent.prompt("grade this") - self.assertIn('"passed"', response.content) + self.assertIn('"passed"', ai_state.text(response)) self.assertTrue(os.path.exists(cassette)) diff --git a/fastapi_startkit/tests/ai/test_record_cassette_format.py b/fastapi_startkit/tests/ai/test_record_cassette_format.py index 92e0b4b9..97f26e3a 100644 --- a/fastapi_startkit/tests/ai/test_record_cassette_format.py +++ b/fastapi_startkit/tests/ai/test_record_cassette_format.py @@ -11,15 +11,35 @@ import unittest from unittest import mock +from langchain_core.messages import AIMessage, HumanMessage, ToolMessage + +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.response import AgentResponse class RecordAgent(Agent): pass -def _fake_prompt(response: AgentResponse): +def _tool_turn_messages() -> list: + return [ + HumanMessage(content="suggest me jobs"), + AIMessage( + content="", + tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1", "type": "tool_call"}], + additional_kwargs={"response_time": 12.0}, + usage_metadata={ + "input_tokens": 117, + "output_tokens": 22, + "total_tokens": 139, + "input_token_details": {"cache_read": 0}, + }, + ), + ToolMessage(content='[{"id": 2}]', tool_call_id="c1", additional_kwargs={"response_time": 5.0}), + ] + + +def _fake_prompt(response: dict): async def prompt(agent_self, message, **kwargs): return response @@ -28,19 +48,7 @@ async def prompt(agent_self, message, **kwargs): class TestNewCassetteFormat(unittest.IsolatedAsyncioTestCase): async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): - response = AgentResponse( - transcript=[ - { - "type": "ai", - "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - "uses": {"input_token": 117, "output_token": 22, "cache_token": 0, "total_token": 139}, - "response_time": 12.0, - }, - {"type": "tool_response", "content_type": "json", "content": '[{"id": 2}]', "response_time": 5.0}, - ], - content='[{"id": 2}]', - tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - ) + response = {"messages": _tool_turn_messages()} with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with _fake_prompt(response): @@ -58,19 +66,7 @@ async def test_recording_persists_an_ordered_transcript_keyed_by_input(self): self.assertEqual(turn[2]["content"], '[{"id": 2}]') async def test_replaying_a_new_format_cassette_reconstructs_tool_calls(self): - response = AgentResponse( - transcript=[ - { - "type": "ai", - "tool_calls": [{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - "uses": {"input_token": 117, "output_token": 22}, - "response_time": 12.0, - }, - {"type": "tool_response", "content_type": "json", "content": '[{"id": 2}]', "response_time": 5.0}, - ], - content='[{"id": 2}]', - tool_calls=[{"name": "job_search_tool", "args": {"query": "python"}, "id": "c1"}], - ) + response = {"messages": _tool_turn_messages()} with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with _fake_prompt(response): @@ -82,11 +78,18 @@ async def test_replaying_a_new_format_cassette_reconstructs_tool_calls(self): replayed = await agent.prompt("suggest me jobs") agent.assert_tool_called("job_search_tool") - self.assertEqual([tc["name"] for tc in replayed.tool_calls], ["job_search_tool"]) - self.assertEqual(replayed.content, '[{"id": 2}]') + self.assertEqual([tc["name"] for tc in ai_state.tool_calls(replayed)], ["job_search_tool"]) + self.assertEqual(ai_state.text(replayed), '[{"id": 2}]') async def test_synthesizes_a_transcript_for_a_plain_response(self): - response = AgentResponse(content="Hi there!", usage={"input": 3, "output": 2}) + response = { + "messages": [ + AIMessage( + content="Hi there!", + usage_metadata={"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, + ) + ] + } with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with _fake_prompt(response): @@ -122,5 +125,5 @@ async def test_old_flat_cassette_still_replays(self): replayed = await agent.prompt("suggest me jobs") agent.assert_tool_called("job_search_tool") - self.assertEqual(replayed.content, "I found a job.") - self.assertEqual([tc["name"] for tc in replayed.tool_calls], ["job_search_tool"]) + self.assertEqual(ai_state.text(replayed), "I found a job.") + self.assertEqual([tc["name"] for tc in ai_state.tool_calls(replayed)], ["job_search_tool"]) diff --git a/fastapi_startkit/tests/ai/test_recording.py b/fastapi_startkit/tests/ai/test_recording.py index 3c8cec80..a074da4a 100644 --- a/fastapi_startkit/tests/ai/test_recording.py +++ b/fastapi_startkit/tests/ai/test_recording.py @@ -12,6 +12,7 @@ """ from fastapi_startkit.ai import recording +from fastapi_startkit.ai import state as ai_state class TestEntryBuilders: @@ -75,7 +76,7 @@ def test_handles_missing_metadata(self): } -class TestReconstructResponse: +class TestReconstructState: def test_final_ai_answer_wins_over_the_tool_result_for_content(self): transcript = [ recording.human("suggest me jobs"), @@ -88,13 +89,14 @@ def test_final_ai_answer_wins_over_the_tool_result_for_content(self): recording.ai(content="Here is a job.", uses={"input_token": 5, "output_token": 9}, response_time=8.0), ] - response = recording.to_response(transcript) + state = recording.to_state(transcript) - assert response.content == "Here is a job." - assert [tc["name"] for tc in response.tool_calls] == ["job_search_tool"] - assert response.usage == {"input": 5, "output": 9} - assert len(response.tool_events) == 1 - assert response.tool_events[0]["content"] == '[{"id": 2}]' + assert ai_state.text(state) == "Here is a job." + assert [tc["name"] for tc in ai_state.tool_calls(state)] == ["job_search_tool"] + # usage sums every model call on the turn (117+5 in, 22+9 out) + assert ai_state.usage(state) == {"input": 122, "output": 31} + assert len(ai_state.tool_events(state)) == 1 + assert ai_state.tool_events(state)[0]["content"] == '[{"id": 2}]' def test_falls_back_to_tool_result_when_there_is_no_final_ai_answer(self): transcript = [ @@ -107,10 +109,10 @@ def test_falls_back_to_tool_result_when_there_is_no_final_ai_answer(self): recording.tool_response(content='[{"id": 2}]', response_time=5.0), ] - response = recording.to_response(transcript) + state = recording.to_state(transcript) - assert response.content == '[{"id": 2}]' - assert [tc["name"] for tc in response.tool_calls] == ["job_search_tool"] + assert ai_state.text(state) == '[{"id": 2}]' + assert [tc["name"] for tc in ai_state.tool_calls(state)] == ["job_search_tool"] def test_plain_text_turn_reconstructs_content_and_no_tool_calls(self): transcript = [ @@ -118,10 +120,10 @@ def test_plain_text_turn_reconstructs_content_and_no_tool_calls(self): recording.ai(content="Hi there!", uses={"input_token": 3, "output_token": 2}, response_time=4.0), ] - response = recording.to_response(transcript) + state = recording.to_state(transcript) - assert response.content == "Hi there!" - assert response.tool_calls == [] + assert ai_state.text(state) == "Hi there!" + assert ai_state.tool_calls(state) == [] class TestChunksRoundTrip: diff --git a/fastapi_startkit/tests/ai/test_runner_recording.py b/fastapi_startkit/tests/ai/test_runner_recording.py index 78ca132e..3c1185cf 100644 --- a/fastapi_startkit/tests/ai/test_runner_recording.py +++ b/fastapi_startkit/tests/ai/test_runner_recording.py @@ -10,6 +10,7 @@ from langchain_core.messages import AIMessage as LCAIMessage from langchain_core.tools import tool +from fastapi_startkit.ai import state as ai_state from fastapi_startkit.ai.agent import Agent from fastapi_startkit.ai.runner import Runner @@ -53,15 +54,16 @@ async def test_tool_calls_are_preserved_not_dropped(self): response = await runner.run("hi") - self.assertEqual([tc["name"] for tc in response.tool_calls], ["echo_tool"]) + self.assertEqual([tc["name"] for tc in ai_state.tool_calls(response)], ["echo_tool"]) async def test_tool_response_and_latency_are_captured(self): runner = _runner_with_model(self._tool_ai_message()) response = await runner.run("hi") - self.assertEqual(len(response.tool_events), 1) - event = response.tool_events[0] + events = ai_state.tool_events(response) + self.assertEqual(len(events), 1) + event = events[0] self.assertEqual(event["name"], "echo_tool") self.assertEqual(event["content"], "echoed:hi") self.assertGreaterEqual(event["response_time"], 0) @@ -71,17 +73,18 @@ async def test_usage_comes_from_the_requesting_ai_message(self): response = await runner.run("hi") - self.assertEqual(response.usage, {"input": 10, "output": 5}) + self.assertEqual(ai_state.usage(response), {"input": 10, "output": 5}) - async def test_transcript_is_an_ordered_ai_then_tool_response(self): + async def test_messages_are_an_ordered_human_ai_then_tool(self): runner = _runner_with_model(self._tool_ai_message()) response = await runner.run("hi") - kinds = [entry["type"] for entry in response.transcript] - self.assertEqual(kinds, ["ai", "tool_response"]) - self.assertEqual(response.transcript[0]["tool_calls"][0]["name"], "echo_tool") - self.assertEqual(response.transcript[1]["content"], "echoed:hi") + kinds = [m.type for m in response["messages"]] + self.assertEqual(kinds, ["human", "ai", "tool"]) + self.assertEqual(response["messages"][1].tool_calls[0]["name"], "echo_tool") + self.assertEqual(response["messages"][2].content, "echoed:hi") + self.assertGreaterEqual(response["messages"][1].additional_kwargs["response_time"], 0) class TestRunnerPlainTurn(unittest.IsolatedAsyncioTestCase): @@ -93,7 +96,7 @@ async def test_plain_answer_records_a_single_ai_message(self): response = await runner.run("hello") - self.assertEqual(response.content, "Hi there!") - self.assertEqual(response.tool_calls, []) - self.assertEqual([e["type"] for e in response.transcript], ["ai"]) - self.assertEqual(response.transcript[0]["content"], "Hi there!") + self.assertEqual(ai_state.text(response), "Hi there!") + self.assertEqual(ai_state.tool_calls(response), []) + self.assertEqual([m.type for m in response["messages"]], ["human", "ai"]) + self.assertEqual(response["messages"][1].content, "Hi there!") diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py index d0957c64..6c06ee53 100644 --- a/fastapi_startkit/tests/ai/test_structured_output.py +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -3,8 +3,8 @@ 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. +model yields {"raw", "parsed", "parsing_error"}; the Runner unwraps it into the +state's "structured_response" and records the raw model turn on "messages". The fake/record paths bypass build(), so they keep parsing the JSON-string content via schema() for deterministic replay. @@ -148,16 +148,17 @@ async def ainvoke(self, messages): class TestResponseMapping(unittest.TestCase): - def test_unwraps_include_raw_into_parsed_and_content(self): + def test_unwraps_include_raw_into_parsed(self): parsed = Movie(title="Inception", year=2010) - response = Runner(MovieAgent())._to_agent_response( + result = Runner(MovieAgent())._parsed_of( {"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, []) + self.assertIs(result, parsed) + + def test_a_plain_message_has_no_parsed(self): + self.assertIsNone(Runner(MovieAgent())._parsed_of(AIMessage(content="hi"))) class TestPromptEndToEnd(unittest.IsolatedAsyncioTestCase): @@ -186,5 +187,6 @@ def with_structured_output(self, schema, **kwargs): response = await MovieAgent().prompt("best nolan movie") - self.assertEqual(response.parsed, parsed) - self.assertEqual(response.content, parsed.model_dump_json()) + self.assertEqual(response["structured_response"], parsed) + # The raw model turn is recorded on the state's messages. + self.assertEqual([m.type for m in response["messages"]], ["human", "ai"])