feat(ai): Ai model registry + fluent record() testing DSL

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

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

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)
 

fastapi_startkit/src/fastapi_startkit/ai/model_builder.py → 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: