From da512fd47d51e288df08d64bb7e464768746bba6 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 11:25:14 -0700 Subject: [PATCH 1/2] refactor(ai): route Agent.prompt()/stream() through a Transport strategy The fake, record, inline-fake, and real branches each duplicated the _log_call() + _apply_schema() tail. Resolve the source once via a Transport (live / stand-in / inline fake) selected in _transport(), so each public method has a single call path and logs/applies its schema in one place. Pure internal refactor: the public prompt/stream/fake/record API and the fake/record mechanism are unchanged. --- .../src/fastapi_startkit/ai/agent.py | 60 +++----- .../src/fastapi_startkit/ai/transport.py | 142 ++++++++++++++++++ 2 files changed, 161 insertions(+), 41 deletions(-) create mode 100644 fastapi_startkit/src/fastapi_startkit/ai/transport.py diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 17c947f3..4d76e6a1 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -6,6 +6,7 @@ from .document import Document from .response import AgentResponse, AgentSnapshot from .testing import AgentBinding +from .transport import InlineFakeTransport, LiveTransport, StandInTransport, Transport if TYPE_CHECKING: from langchain_core.tools import BaseTool @@ -49,31 +50,12 @@ 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) - - _run_kwargs = dict( + response = await self._transport(message).prompt( + message, 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) - - response = await self._run_pipeline(chat_model, messages) self._log_call("prompt", message) return self._apply_schema(response) @@ -85,26 +67,11 @@ async def stream( provider_options: dict | None = None, ) -> AsyncIterator[str]: self._log_call("stream", message) - - 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 - 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): + async for chunk in self._transport(message).stream( + message, + model=model, + provider_options=provider_options, + ): yield chunk @classmethod @@ -135,6 +102,17 @@ def _faked(self) -> Any: binding = type(self)._binding() return binding if binding is not self else None + def _transport(self, message: str) -> Transport: + stand_in = self._faked() + if stand_in is not None: + return StandInTransport(self, stand_in) + + match = self._match_fake(message) + if match is not None: + return InlineFakeTransport(self, match) + + return LiveTransport(self) + 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: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/transport.py b/fastapi_startkit/src/fastapi_startkit/ai/transport.py new file mode 100644 index 00000000..58071624 --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/transport.py @@ -0,0 +1,142 @@ +"""Transport strategies for Agent.prompt()/stream(). + +A single call resolves exactly one transport — the real model pipeline, a +stand-in bound via ``Agent.fake()``/``Agent.record()``, or an inline fake matched +from ``Agent._fakes``. Concentrating the branch here lets the public methods keep +one call path and log / apply their schema in one place instead of once per branch. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, AsyncIterator + +from .response import AgentResponse, AgentSnapshot + +if TYPE_CHECKING: + from .agent import Agent + from .document import Document + + +class Transport: + def __init__(self, agent: Agent) -> None: + self._agent = agent + + async def prompt( + self, + message: str, + *, + model: str | None, + attachments: list[Document] | None, + provider_options: dict | None, + ) -> AgentResponse: + raise NotImplementedError + + def stream( + self, + message: str, + *, + model: str | None, + provider_options: dict | None, + ) -> AsyncIterator[str]: + raise NotImplementedError + + +class LiveTransport(Transport): + """Runs the real model pipeline.""" + + async def prompt( + self, + message: str, + *, + model: str | None, + attachments: list[Document] | None, + provider_options: dict | None, + ) -> AgentResponse: + agent = self._agent + messages = agent._build_messages(message, attachments) + chat_model = agent._build_model(model, provider_options) + return await agent._run_pipeline(chat_model, messages) + + async def stream( + self, + message: str, + *, + model: str | None, + provider_options: dict | None, + ) -> AsyncIterator[str]: + async for chunk in self._agent._stream(message, model=model, provider_options=provider_options): + yield chunk + + +class StandInTransport(Transport): + """Delegates to a stand-in bound via Agent.fake()/record().""" + + def __init__(self, agent: Agent, stand_in: Any) -> None: + super().__init__(agent) + self._stand_in = stand_in + + async def prompt( + self, + message: str, + *, + model: str | None, + attachments: list[Document] | None, + provider_options: dict | None, + ) -> AgentResponse: + return await self._stand_in.prompt(message, attachments=attachments) + + async def stream( + self, + message: str, + *, + model: str | None, + provider_options: dict | None, + ) -> AsyncIterator[str]: + stand_in = self._stand_in + if hasattr(stand_in, "stream"): + async for chunk in stand_in.stream(message): + yield chunk + else: + response = await stand_in.prompt(message) + yield response.content + + +class InlineFakeTransport(Transport): + """Serves a canned response matched from Agent._fakes.""" + + def __init__(self, agent: Agent, match: AgentResponse | AgentSnapshot) -> None: + super().__init__(agent) + self._match = match + + async def prompt( + self, + message: str, + *, + model: str | None, + attachments: list[Document] | None, + provider_options: dict | None, + ) -> AgentResponse: + match = self._match + if isinstance(match, AgentSnapshot): + return await match.resolve( + self._agent, + message, + model=model, + attachments=attachments, + provider_options=provider_options, + ) + return match + + async def stream( + self, + message: str, + *, + model: str | None, + provider_options: dict | None, + ) -> AsyncIterator[str]: + match = self._match + if isinstance(match, AgentSnapshot): + response = await match.resolve(self._agent, message) + else: + response = match + yield response.content From 27bc67f9866ddd9a0082a707f63ec05bf45b6818 Mon Sep 17 00:00:00 2001 From: Bedram Tamang Date: Fri, 17 Jul 2026 12:00:10 -0700 Subject: [PATCH 2/2] refactor(ai): fake/record swap the model, not the whole agent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Agent.fake()/record() previously bound an entire stand-in Agent into the container, so a faked agent bypassed its own pipeline — instructions, middleware, tools, and schema never ran. Move the seam to the model: - ModelBuilder.get_model_for() returns a registered double if one exists for the agent class, else builds the real model. Agent._build_model() uses it. - fake(responses: list) swaps in fake_chat_model(responses); record(cassette, messages) swaps in a RecordingModel that replays a cassette or, on a miss, calls the real model once and records it. - prompt()/stream() now have a single call path — the agent always runs its real pipeline, only the model is doubled. Removes the whole-agent binding machinery (make/_binding/_faked, the inline _fakes/_match_fake path, FakeAgent/RecordingAgent/AgentBinding, and the Transport indirection) now that no branch remains to dispatch. --- .../src/fastapi_startkit/ai/__init__.py | 8 +- .../src/fastapi_startkit/ai/agent.py | 71 ++---- .../src/fastapi_startkit/ai/model_builder.py | 19 ++ .../src/fastapi_startkit/ai/testing.py | 224 ++++++++---------- .../src/fastapi_startkit/ai/transport.py | 142 ----------- fastapi_startkit/tests/ai/test_agent_fake.py | 216 +++++++---------- .../tests/ai/test_agent_schema.py | 22 +- 7 files changed, 238 insertions(+), 464 deletions(-) delete mode 100644 fastapi_startkit/src/fastapi_startkit/ai/transport.py diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index dbf4c85e..16c07c01 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -11,20 +11,18 @@ from .image_factory import ImageFactory from .providers.ai_provider import AIProvider from .response import AgentResponse, AgentSnapshot -from .testing import AgentBinding, FakeAgent, NoFakeResponse, RecordingAgent +from .testing import AgentFake, AgentRecordFake __all__ = [ "Agent", "Middleware", - "AgentBinding", + "AgentFake", + "AgentRecordFake", "AgentResponse", "AgentSnapshot", "AIConfig", "AIProvider", "AnthropicConfig", - "FakeAgent", - "NoFakeResponse", - "RecordingAgent", "Audio", "AudioResponse", "AudioFactory", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 4d76e6a1..3180f00b 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -1,16 +1,15 @@ 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 .testing import AgentBinding -from .transport import InlineFakeTransport, LiveTransport, StandInTransport, Transport +from .response import AgentResponse if TYPE_CHECKING: from langchain_core.tools import BaseTool + from .testing import AgentFake, AgentRecordFake + class Agent: provider: str | None = None @@ -21,7 +20,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]: @@ -50,12 +48,9 @@ async def prompt( attachments: list[Document] | None = None, provider_options: dict | None = None, ) -> AgentResponse: - response = await self._transport(message).prompt( - message, - model=model, - attachments=attachments, - provider_options=provider_options, - ) + 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) @@ -67,51 +62,20 @@ async def stream( provider_options: dict | None = None, ) -> AsyncIterator[str]: self._log_call("stream", message) - async for chunk in self._transport(message).stream( - message, - model=model, - provider_options=provider_options, - ): + 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 - - return AgentBinding(cls, FakeAgent(responses)) - - @classmethod - def record(cls, cassette: str | None = None) -> "AgentBinding": - from .testing import AgentBinding, RecordingAgent - - return AgentBinding(cls, RecordingAgent(cls(), cassette)) - - @classmethod - def _binding(cls) -> Any: - from fastapi_startkit.application import app + def fake(cls, responses: list) -> "AgentFake": + from .testing import AgentFake - container = app() - return container.make(cls.__name__) if container.has(cls.__name__) else None + return AgentFake(cls, responses) @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 _transport(self, message: str) -> Transport: - stand_in = self._faked() - if stand_in is not None: - return StandInTransport(self, stand_in) + def record(cls, cassette: str | None = None, messages: list | None = None) -> "AgentRecordFake": + from .testing import AgentRecordFake - match = self._match_fake(message) - if match is not None: - return InlineFakeTransport(self, match) - - return LiveTransport(self) + return AgentRecordFake(cls(), cassette, messages) def assert_prompted(self, times: int | None = None) -> None: calls = [c for c in self._call_log if c["method"] in ("prompt", "stream")] @@ -124,16 +88,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}) @@ -203,7 +160,7 @@ def _build_messages( def _build_model(self, model: str | None = None, provider_options: dict | None = None) -> Any: from .model_builder import ModelBuilder # noqa: PLC0415 - return ModelBuilder(agent=self).build(model, provider_options) + return ModelBuilder(agent=self).get_model_for(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/model_builder.py b/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py index 7f82254f..1b2b47cf 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py @@ -9,9 +9,28 @@ class ModelBuilder: + _fakes: dict[type, Any] = {} + def __init__(self, agent: "Agent") -> None: self._agent = agent + @classmethod + def register_fake(cls, agent_cls: type, model: Any) -> None: + cls._fakes[agent_cls] = model + + @classmethod + def clear_fake(cls, agent_cls: type) -> None: + cls._fakes.pop(agent_cls, None) + + def has_fake(self) -> bool: + return type(self._agent) in self._fakes + + def get_model_for(self, model: str | None = None, provider_options: dict | None = None) -> Any: + double = self._fakes.get(type(self._agent)) + if double is not None: + return double + return self.build(model, provider_options) + def build(self, model: str | None = None, provider_options: dict | None = None) -> Any: from langchain.chat_models import init_chat_model # noqa: PLC0415 diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index e3e8533b..ac1e4052 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -1,63 +1,17 @@ from __future__ import annotations -import fnmatch import functools import hashlib import inspect import json -import re import sys -from collections.abc import AsyncIterator from pathlib import Path from typing import TYPE_CHECKING, Any, Callable -from .response import AgentResponse +from .model_builder import ModelBuilder if TYPE_CHECKING: from .agent import Agent - 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 "*?["): - return fnmatch.fnmatch(message, pattern) - 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] = [] - 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: @@ -65,114 +19,93 @@ 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] +def _text(content: Any) -> str: + return content if isinstance(content, str) else str(content) -class FakeAgent(_Recorder): - def __init__(self, responses: dict[str, Any] | None = None) -> None: - super().__init__() - self.responses = responses or {} +class RecordingModel: + """A chat-model stand-in that replays a cassette or, on a miss, calls the + real model once and records the response for future runs.""" - 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 __init__(self, agent: Agent, cassette: Path, initial_messages: list | None = None) -> None: + self._agent = agent + self._cassette = cassette + self._initial = list(initial_messages or []) - async def prompt(self, message: str, attachments: list[Document] | None = None) -> AgentResponse: - self._record_call(message, attachments) - return AgentResponse(content=self._resolve(message)) - - async def stream(self, message: str) -> AsyncIterator[str]: - self._record_call(message, None) - for chunk in _word_chunks(self._resolve(message)): - yield chunk + def bind_tools(self, tools: Any, **kwargs: Any) -> RecordingModel: + return self + def _real(self) -> Any: + return ModelBuilder(agent=self._agent).build() -class RecordingAgent(_Recorder): - def __init__(self, real: Agent, cassette: str | None = None) -> None: - super().__init__() - self._real = real - self.cassette: Path | None = Path(cassette) if cassette else None + def _prepared(self, messages: Any) -> list: + return self._initial + list(messages) @staticmethod - def _key(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) + def _key(messages: list) -> str: + payload = json.dumps(messages, sort_keys=True, default=str) return hashlib.sha256(payload.encode()).hexdigest() - def _load(self) -> tuple[Path, dict]: - cassette = self.cassette - assert cassette is not None, "RecordingAgent has no cassette resolved" - return cassette, (json.loads(cassette.read_text()) if cassette.exists() else {}) + def _load(self) -> dict: + return json.loads(self._cassette.read_text()) if self._cassette.exists() else {} - def _save(self, cassette: Path, store: dict, key: str, value: Any) -> None: + def _save(self, store: dict, key: str, value: Any) -> None: store[key] = value - cassette.parent.mkdir(parents=True, exist_ok=True) - cassette.write_text(json.dumps(store, indent=2, sort_keys=True)) + self._cassette.parent.mkdir(parents=True, exist_ok=True) + self._cassette.write_text(json.dumps(store, indent=2, sort_keys=True)) + + async def ainvoke(self, messages: Any, **kwargs: Any) -> Any: + from langchain_core.messages import AIMessage # noqa: PLC0415 - async def prompt(self, message: str, attachments: list[Document] | None = None) -> AgentResponse: - self._record_call(message, attachments) - cassette, store = self._load() - key = self._key(message, attachments) + prepared = self._prepared(messages) + store = self._load() + key = self._key(prepared) 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) + return AIMessage(content=_joined(store[key])) + + response = await self._real().ainvoke(prepared) + self._save(store, key, _text(response.content)) return response - async def stream(self, message: str) -> AsyncIterator[str]: - self._record_call(message, None) - cassette, store = self._load() - key = self._key(message, None) + async def astream(self, messages: Any, **kwargs: Any) -> Any: + from langchain_core.messages import AIMessageChunk # noqa: PLC0415 + + prepared = self._prepared(messages) + store = self._load() + key = self._key(prepared) if key in store: value = store[key] for chunk in value if isinstance(value, list) else [value]: - yield chunk + yield AIMessageChunk(content=chunk) return - chunks = [chunk async for chunk in self._real._stream(message)] - self._save(cassette, store, key, chunks) - for chunk in chunks: + + chunks: list[str] = [] + async for chunk in self._real().astream(prepared): + chunks.append(_text(chunk.content)) yield chunk + self._save(store, key, chunks) -class AgentBinding: - def __init__(self, agent_cls: type[Agent], stand_in: Any) -> None: - self._agent_cls = agent_cls - self._stand_in = stand_in +class _ModelSwap: + """Registers a model stand-in for an agent class for the duration of a + ``with`` block (or a decorated function) — the agent itself is untouched, so + its real pipeline runs and only the model is swapped.""" - 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 + _agent_cls: type[Agent] - def __enter__(self) -> Any: - from fastapi_startkit.application import app + def _build_model(self) -> Any: + raise NotImplementedError - 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 + def __enter__(self) -> Any: + model = self._build_model() + ModelBuilder.register_fake(self._agent_cls, model) + return model def __exit__(self, *_exc: Any) -> bool: - from fastapi_startkit.application import app - - app().unbind(self._agent_cls.__name__) + ModelBuilder.clear_fake(self._agent_cls) return False def __call__(self, func: Callable) -> Callable: - self._resolve_cassette(inspect.getfile(func), func.__qualname__) - if inspect.iscoroutinefunction(func): @functools.wraps(func) @@ -188,3 +121,46 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: return func(*args, **kwargs) return wrapper + + +class AgentFake(_ModelSwap): + """Swaps the agent's model for one that replays ``responses`` in order.""" + + def __init__(self, agent_cls: type[Agent], responses: list) -> None: + self._agent_cls = agent_cls + self._responses = responses + + def _build_model(self) -> Any: + from .fakes import fake_chat_model # noqa: PLC0415 + + return fake_chat_model(self._responses) + + +class AgentRecordFake(_ModelSwap): + """Swaps the agent's model for a record-and-replay model backed by a cassette.""" + + def __init__(self, agent: Agent, cassette: str | None = None, messages: list | None = None) -> None: + self._agent = agent + self._agent_cls = type(agent) + self._cassette: Path | None = Path(cassette) if cassette else None + self._messages = messages + + def _resolve_cassette(self, filename: str, qualname: str) -> None: + here = Path(filename).parent + 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 _build_model(self) -> Any: + assert self._cassette is not None, "RecordingModel has no cassette resolved" + return RecordingModel(self._agent, self._cassette, self._messages) + + def __enter__(self) -> Any: + caller = sys._getframe(1).f_code + self._resolve_cassette(caller.co_filename, caller.co_qualname) + return super().__enter__() + + def __call__(self, func: Callable) -> Callable: + self._resolve_cassette(inspect.getfile(func), func.__qualname__) + return super().__call__(func) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/transport.py b/fastapi_startkit/src/fastapi_startkit/ai/transport.py deleted file mode 100644 index 58071624..00000000 --- a/fastapi_startkit/src/fastapi_startkit/ai/transport.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Transport strategies for Agent.prompt()/stream(). - -A single call resolves exactly one transport — the real model pipeline, a -stand-in bound via ``Agent.fake()``/``Agent.record()``, or an inline fake matched -from ``Agent._fakes``. Concentrating the branch here lets the public methods keep -one call path and log / apply their schema in one place instead of once per branch. -""" - -from __future__ import annotations - -from typing import TYPE_CHECKING, Any, AsyncIterator - -from .response import AgentResponse, AgentSnapshot - -if TYPE_CHECKING: - from .agent import Agent - from .document import Document - - -class Transport: - def __init__(self, agent: Agent) -> None: - self._agent = agent - - async def prompt( - self, - message: str, - *, - model: str | None, - attachments: list[Document] | None, - provider_options: dict | None, - ) -> AgentResponse: - raise NotImplementedError - - def stream( - self, - message: str, - *, - model: str | None, - provider_options: dict | None, - ) -> AsyncIterator[str]: - raise NotImplementedError - - -class LiveTransport(Transport): - """Runs the real model pipeline.""" - - async def prompt( - self, - message: str, - *, - model: str | None, - attachments: list[Document] | None, - provider_options: dict | None, - ) -> AgentResponse: - agent = self._agent - messages = agent._build_messages(message, attachments) - chat_model = agent._build_model(model, provider_options) - return await agent._run_pipeline(chat_model, messages) - - async def stream( - self, - message: str, - *, - model: str | None, - provider_options: dict | None, - ) -> AsyncIterator[str]: - async for chunk in self._agent._stream(message, model=model, provider_options=provider_options): - yield chunk - - -class StandInTransport(Transport): - """Delegates to a stand-in bound via Agent.fake()/record().""" - - def __init__(self, agent: Agent, stand_in: Any) -> None: - super().__init__(agent) - self._stand_in = stand_in - - async def prompt( - self, - message: str, - *, - model: str | None, - attachments: list[Document] | None, - provider_options: dict | None, - ) -> AgentResponse: - return await self._stand_in.prompt(message, attachments=attachments) - - async def stream( - self, - message: str, - *, - model: str | None, - provider_options: dict | None, - ) -> AsyncIterator[str]: - stand_in = self._stand_in - if hasattr(stand_in, "stream"): - async for chunk in stand_in.stream(message): - yield chunk - else: - response = await stand_in.prompt(message) - yield response.content - - -class InlineFakeTransport(Transport): - """Serves a canned response matched from Agent._fakes.""" - - def __init__(self, agent: Agent, match: AgentResponse | AgentSnapshot) -> None: - super().__init__(agent) - self._match = match - - async def prompt( - self, - message: str, - *, - model: str | None, - attachments: list[Document] | None, - provider_options: dict | None, - ) -> AgentResponse: - match = self._match - if isinstance(match, AgentSnapshot): - return await match.resolve( - self._agent, - message, - model=model, - attachments=attachments, - provider_options=provider_options, - ) - return match - - async def stream( - self, - message: str, - *, - model: str | None, - provider_options: dict | None, - ) -> AsyncIterator[str]: - match = self._match - if isinstance(match, AgentSnapshot): - response = await match.resolve(self._agent, message) - else: - response = match - yield response.content diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 3191a8b3..9fbdc201 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -1,9 +1,10 @@ """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(responses)`` swaps the agent's *model* for a fake that replays the +given responses in order — the agent itself is untouched, so its real pipeline +(instructions, middleware, schema) still runs. ``Agent.record(cassette)`` swaps in +a record-and-replay model: on a cassette miss it calls the real model once and +caches the response to disk; on a hit it replays without calling the model again. """ import json @@ -12,8 +13,13 @@ import unittest from unittest import mock +import langchain.chat_models as chat_models +from langchain_core.messages import AIMessage + +from fastapi_startkit.ai import AIConfig, fake_chat_model from fastapi_startkit.ai.agent import Agent from fastapi_startkit.ai.response import AgentResponse +from fastapi_startkit.application import app class SimpleAgent(Agent): @@ -21,92 +27,56 @@ class SimpleAgent(Agent): class TestAgentFake(unittest.IsolatedAsyncioTestCase): - async def test_fake_with_agent_response_returns_it(self): + async def test_fake_returns_the_response(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="Hello world!")}): + with SimpleAgent.fake([AIMessage(content="Hello world!")]): result = await agent.prompt("anything") + self.assertIsInstance(result, AgentResponse) 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) - - agent._run = patched_run - - with SimpleAgent.fake({"*": AgentResponse(content="faked")}): - await agent.prompt("hello") - - self.assertEqual(called, [], "_run() must not be called when a fake matches") - - async def test_fake_with_exact_pattern(self): + async def test_fake_accepts_plain_strings(self): agent = SimpleAgent() - with SimpleAgent.fake({"hello": AgentResponse(content="matched hello")}): + with SimpleAgent.fake(["matched hello"]): result = await agent.prompt("hello") self.assertEqual(result.content, "matched 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", "second"]): + first = await agent.prompt("one") + second = await agent.prompt("two") - self.assertEqual(result.content, "hi there") - - async def test_fake_glob_analyze_wildcard(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*analyze*": AgentResponse(content="analysis done")}): - result = await agent.prompt("please analyze this report") + self.assertEqual(first.content, "first") + self.assertEqual(second.content, "second") - self.assertEqual(result.content, "analysis done") + async def test_fake_does_not_build_the_real_model(self): + def boom(*_a, **_k): + raise AssertionError("the real model must not be built when faked") - 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_glob_case_insensitive(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*HELLO*": AgentResponse(content="case insensitive")}): - result = await agent.prompt("say hello please") - - self.assertEqual(result.content, "case insensitive") - - 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") + with mock.patch.object(chat_models, "init_chat_model", boom): + with SimpleAgent.fake(["faked"]): + result = await SimpleAgent().prompt("hello") - self.assertEqual(result.content, "first match") + self.assertEqual(result.content, "faked") 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(["a", "b"]): 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 +98,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 +106,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,76 +120,84 @@ 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() agent.assert_not_prompted() - async def test_fake_rebinding_overrides_previous(self): - agent = SimpleAgent() + async def test_fake_is_scoped_to_the_with_block(self): + def boom(*_a, **_k): + raise AssertionError("model should only be faked inside the with block") - with SimpleAgent.fake({"*": AgentResponse(content="first fake")}): + agent = SimpleAgent() + 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): + async def test_fake_as_decorator(self): + @SimpleAgent.fake(["decorated"]) + async def run(): + return await SimpleAgent().prompt("hi") + + result = await run() + self.assertEqual(result.content, "decorated") + + async def test_stream_returns_fake_response_in_word_chunks(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): - agent = SimpleAgent() - 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. agent.assert_prompted(times=1) class TestAgentRecord(unittest.IsolatedAsyncioTestCase): - def setup_agent(self, content): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + + def patch_real(self, contents): + """Patch the real model so a cassette miss returns the next content and + counts how many times the real model was built (i.e. cassette misses).""" calls = [] + remaining = iter(contents) - async def fake_run(agent_self, message, **kwargs): - calls.append(message) - return AgentResponse(content=content) + def init(*_a, **_k): + calls.append(True) + return fake_chat_model([AIMessage(content=next(remaining))]) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(chat_models, "init_chat_model", init) patcher.start() self.addCleanup(patcher.stop) return calls async def test_first_run_records_response_to_cassette(self): - calls = self.setup_agent("recorded reply") + calls = self.patch_real(["recorded reply"]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with SimpleAgent.record(cassette): result = await SimpleAgent().prompt("hello") self.assertEqual(result.content, "recorded reply") - self.assertEqual(calls, ["hello"]) + self.assertEqual(len(calls), 1) self.assertTrue(os.path.exists(cassette)) with open(cassette) as f: self.assertIn("recorded reply", json.load(f).values()) - async def test_second_run_replays_without_calling_run(self): - calls = self.setup_agent("recorded reply") + async def test_second_run_replays_without_calling_the_model(self): + calls = self.patch_real(["recorded reply"]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") with SimpleAgent.record(cassette): @@ -228,77 +206,57 @@ async def test_second_run_replays_without_calling_run(self): replayed = await SimpleAgent().prompt("hello") self.assertEqual(replayed.content, "recorded reply") - self.assertEqual(calls, ["hello"]) + self.assertEqual(len(calls), 1) # real model built only on the first run async def test_replay_prefers_cassette_over_live_response(self): - async def first_run(s, m, **k): - return AgentResponse(content="from first record") - - async def changed_run(s, m, **k): - return AgentResponse(content="changed live value") - + self.patch_real(["from first record", "changed live value"]) 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 SimpleAgent.record(cassette): + await SimpleAgent().prompt("hello") + with SimpleAgent.record(cassette): + result = await SimpleAgent().prompt("hello") self.assertEqual(result.content, "from first record") async def test_distinct_messages_are_recorded_separately(self): - calls = self.setup_agent("reply") + self.patch_real(["hello reply", "goodbye 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") - self.assertEqual(calls, ["hello", "goodbye"]) with open(cassette) as f: self.assertEqual(len(json.load(f)), 2) - def setup_stream(self, chunks): + def patch_real_stream(self, chunks): calls = [] - async def fake_stream(agent_self, message, **kwargs): - calls.append(message) - for chunk in chunks: - yield chunk + def init(*_a, **_k): + calls.append(True) + return fake_chat_model([AIMessage(content="".join(chunks))]) - patcher = mock.patch.object(SimpleAgent, "_stream", fake_stream) + patcher = mock.patch.object(chat_models, "init_chat_model", init) patcher.start() self.addCleanup(patcher.stop) return calls - async def test_stream_first_run_records_chunk_list_to_cassette(self): - calls = self.setup_stream(["Hel", "lo!"]) + async def test_stream_records_chunks_then_replays_without_calling_the_model(self): + calls = self.patch_real_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")] - - self.assertEqual(chunks, ["Hel", "lo!"]) - self.assertEqual(calls, ["hi"]) - with open(cassette) as f: - self.assertEqual(list(json.load(f).values()), [["Hel", "lo!"]]) - - 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")] + recorded = [c async for c in SimpleAgent().stream("hi")] with SimpleAgent.record(cassette): replayed = [c async for c in SimpleAgent().stream("hi")] - self.assertEqual(replayed, ["Hel", "lo!"]) - self.assertEqual(calls, ["hi"]) # real stream invoked only on the first run + self.assertEqual("".join(recorded), "Hello!") + self.assertEqual("".join(replayed), "Hello!") + self.assertEqual(len(calls), 1) # real stream invoked only on the first run async def test_prompt_reads_a_stream_recorded_cassette_as_joined_content(self): - self.setup_stream(["Hel", "lo!"]) + self.patch_real_stream(["Hel", "lo!"]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "s.json") with SimpleAgent.record(cassette): diff --git a/fastapi_startkit/tests/ai/test_agent_schema.py b/fastapi_startkit/tests/ai/test_agent_schema.py index 0e7d705b..6c7a8f29 100644 --- a/fastapi_startkit/tests/ai/test_agent_schema.py +++ b/fastapi_startkit/tests/ai/test_agent_schema.py @@ -3,10 +3,13 @@ import unittest from unittest import mock +import langchain.chat_models as chat_models +from langchain_core.messages import AIMessage from pydantic import BaseModel +from fastapi_startkit.ai import AIConfig, fake_chat_model from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.response import AgentResponse +from fastapi_startkit.application import app class User(BaseModel): @@ -20,9 +23,14 @@ def schema(self): class TestAgentSchema(unittest.IsolatedAsyncioTestCase): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + 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,24 +40,24 @@ 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") 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"}') + def init(*_a, **_k): + return fake_chat_model([AIMessage(content='{"id": "u-9", "name": "Sam"}')]) with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "user.json") - with mock.patch.object(UserAgent, "_run", fake_run): + with mock.patch.object(chat_models, "init_chat_model", init): with UserAgent.record(cassette): recorded = await UserAgent().prompt("get the user") with UserAgent.record(cassette):