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