refactor(ai): fake/record swap the model, not the whole agent

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",

fastapi_startkit/src/fastapi_startkit/ai/agent.py

@@ -1,15 +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 .response import AgentResponse
 
 if TYPE_CHECKING:
     from langchain_core.tools import BaseTool
 
+    from .testing import AgentFake, AgentRecordFake
+
 
 class Agent:
     provider: str | None = None
@@ -20,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]:
@@ -49,30 +48,8 @@ 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(
-            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,55 +62,20 @@ 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):
             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 record(cls, cassette: str | None = None, messages: list | None = None) -> "AgentRecordFake":
+        from .testing import AgentRecordFake
 
-    def _faked(self) -> Any:
-        binding = type(self)._binding()
-        return binding if binding is not self else None
+        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")]
@@ -146,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})
 
@@ -225,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 []

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