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.

Author: tmgbedu

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"})
 

fastapi_startkit/src/fastapi_startkit/ai/__init__.py

@@ -10,21 +10,22 @@
 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",

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).get_model_for(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 []

fastapi_startkit/src/fastapi_startkit/ai/model_builder.py

@@ -8,19 +8,19 @@
     from .agent import Agent
 
 
-class ModelBuilder:
+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_models: dict[str, Any] = {}
-    # Reserved for response-level fakes (mirroring fake_models, but for
+    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_responses: dict[str, Any] = {}
+    fake_agent_responses: dict[str, Any] = {}
 
-    def __init__(self, agent: "Agent") -> None:
-        self._agent = agent
+    def __init__(self) -> None:
+        pass
 
     @staticmethod
     def _key(agent: "Agent | str") -> str:
@@ -38,59 +38,63 @@ 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_models[cls._key(agent)] = model
+        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_models
+        return cls._key(agent) in cls.fake_agent_models
 
     @classmethod
     def get_fake_model_for(cls, agent: "Agent | str") -> Any:
-        return cls.fake_models[cls._key(agent)]
+        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_models.clear()
-        cls.fake_responses.clear()
+        cls.fake_agent_models.clear()
+        cls.fake_agent_responses.clear()
 
-    def get_model_for(self, model: str | None = None, provider_options: dict | None = None) -> Any:
+    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
-        this builder's agent, otherwise a freshly-built provider model."""
-        if self.has_fake_model_for(self._agent):
-            return self.get_fake_model_for(self._agent)
-        return self.build(model, provider_options)
+        ``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, model: str | None = None, provider_options: dict | None = None) -> Any:
+    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

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,23 +17,13 @@
     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] = []
@@ -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):