diff --git a/.vscode/settings.json b/.vscode/settings.json new file mode 100644 index 00000000..83ddff40 --- /dev/null +++ b/.vscode/settings.json @@ -0,0 +1,3 @@ +{ + "python.languageServer": "Default" +} \ No newline at end of file diff --git a/example/agents/.gitignore b/example/agents/.gitignore index b2db17eb..dc2d3589 100644 --- a/example/agents/.gitignore +++ b/example/agents/.gitignore @@ -7,3 +7,7 @@ storage node_modules /public/build /public/hot +.ruff_cache +.pytest_cache +.ai +.vite diff --git a/example/agents/app/agents/chat.py b/example/agents/app/agents/chat.py index 6d134452..4e02a9e2 100644 --- a/example/agents/app/agents/chat.py +++ b/example/agents/app/agents/chat.py @@ -1,12 +1,27 @@ from typing import Callable -from fastapi_startkit.ai import Agent, Middleware +from fastapi_startkit.ai import Agent, GraphAgent, Middleware from app.middleware.agent_logger import AgentLogger from app.tools.job_search_tool import job_search_tool -class RouterAgent(Agent): +class RouterAgent(GraphAgent): + def graph(self): + graph = StateGraph(MessagesState) + + # Add nodes + graph.add_node("llm_call", llm_call) + graph.add_node("tool_node", tool_node) + + # Add edges to connect nodes + graph.add_edge(START, "llm_call") + graph.add_conditional_edges("llm_call", should_continue, ["tool_node", END]) + graph.add_edge("tool_node", "llm_call") + + # Compile the agent + agent = graph.compile() + def middleware(self) -> list[Middleware]: return [AgentLogger()] diff --git a/example/agents/app/agents/graph_agent.py b/example/agents/app/agents/graph_agent.py new file mode 100644 index 00000000..c5d4dd03 --- /dev/null +++ b/example/agents/app/agents/graph_agent.py @@ -0,0 +1,9 @@ +from fastapi_startkit.ai import GraphAgent, GraphRunner +from langgraph.graph import StateGraph + + +class SalesAgent(GraphAgent): + def checkpointer(self): + pass + def graph(self, runner: GraphRunner) -> StateGraph: + return StateGraph() diff --git a/example/agents/app/providers/langchain_provider.py b/example/agents/app/providers/langchain_provider.py new file mode 100644 index 00000000..67fcaeb3 --- /dev/null +++ b/example/agents/app/providers/langchain_provider.py @@ -0,0 +1,44 @@ +from fastapi_startkit.support import Provider +from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from psycopg.rows import dict_row +from psycopg_pool import AsyncConnectionPool + + +class LazyCheckpointer: + """An awaitable, lazily-initialised async Postgres checkpointer. + + ``Provider.boot()`` runs synchronously while the application is being + constructed — before any serving event loop exists. Async Postgres + connections are bound to the event loop that opens them, so the pool + cannot be opened at boot time. Instead we build the pool closed and open + it (and run ``setup()``) on first ``await`` inside the request loop, + caching the ready saver for subsequent calls. + + Usage: ``checkpointer = await app().make("checkpointer")`` + """ + + def __init__(self, uri: str): + self._pool = AsyncConnectionPool( + conninfo=uri, + open=False, + kwargs={"autocommit": True, "row_factory": dict_row}, + ) + self._saver: AsyncPostgresSaver | None = None + + async def resolve(self) -> AsyncPostgresSaver: + if self._saver is None: + await self._pool.open() + saver = AsyncPostgresSaver(self._pool) + await saver.setup() + self._saver = saver + return self._saver + + def __await__(self): + return self.resolve().__await__() + + +class LangChainProvider(Provider): + DB_URI = "postgresql://postgres:postgres@localhost:5432/agents?sslmode=disable" + + def boot(self): + self.app.bind("checkpointer", LazyCheckpointer(self.DB_URI)) diff --git a/example/agents/app/tools/job_search_tool.py b/example/agents/app/tools/job_search_tool.py index 4d76fd83..63f96a53 100644 --- a/example/agents/app/tools/job_search_tool.py +++ b/example/agents/app/tools/job_search_tool.py @@ -9,7 +9,7 @@ ] -@tool +@tool(description="Use this tools if user wants to search for jobs") def job_search_tool(query: str) -> list: """Searches for jobs based on the given query. Supports wildcards (* and ?) in each term.""" import fnmatch @@ -17,9 +17,7 @@ def job_search_tool(query: str) -> list: patterns = [f"*{term}*" for term in query.lower().split()] return [ - job for job in jobs - if any( - fnmatch.fnmatch(" ".join(str(v) for v in job.values()).lower(), pattern) - for pattern in patterns - ) + job + for job in jobs + if any(fnmatch.fnmatch(" ".join(str(v) for v in job.values()).lower(), pattern) for pattern in patterns) ] diff --git a/example/agents/bootstrap/application.py b/example/agents/bootstrap/application.py index f198c0a0..47011cfe 100644 --- a/example/agents/bootstrap/application.py +++ b/example/agents/bootstrap/application.py @@ -11,11 +11,13 @@ from config.logging import LoggingConfig from config.vite import ViteConfig from app.providers.fastapi_provider import FastapiProvider +from app.providers.langchain_provider import LangChainProvider app: Application = Application( base_path=Path(__file__).resolve().parent.parent, providers=[ AISkillProvider, + LangChainProvider, (LogProvider,LoggingConfig), (FastapiProvider, FastAPIConfig), AIProvider, diff --git a/example/agents/pyproject.toml b/example/agents/pyproject.toml index aa074982..4b1d78d7 100644 --- a/example/agents/pyproject.toml +++ b/example/agents/pyproject.toml @@ -12,6 +12,7 @@ dependencies = [ "fastapi-startkit[ai,database,fastapi,sqlite]==0.44.0", "langchain>=1.3.10", "langchain-google-genai>=4.2.5", + "langgraph-checkpoint-postgres>=3.1.0", ] [tool.uv] diff --git a/example/agents/routes/api.py b/example/agents/routes/api.py index e1804813..d78f5073 100644 --- a/example/agents/routes/api.py +++ b/example/agents/routes/api.py @@ -1,6 +1,7 @@ from fastapi import APIRouter, Request from fastapi.responses import StreamingResponse from fastapi_startkit.inertia import Inertia +from langchain.agents import create_agent from app.agents.chat import RouterAgent from app.requests.chat import ChatRequest @@ -20,6 +21,10 @@ async def index(request: Request): @api.post("/chat") async def chat(request: ChatRequest): + from fastapi_startkit.application import app + + agent = create_agent(checkpointer=await app().make("checkpointer")) + response = await RouterAgent().prompt(request.message) return {"content": response.content} diff --git a/example/agents/tests/features/job_search.json b/example/agents/tests/features/job_search.json new file mode 100644 index 00000000..d102026f --- /dev/null +++ b/example/agents/tests/features/job_search.json @@ -0,0 +1,6 @@ +{ + "4fd116f6d1844cc5e3bfb521bf361bf66ac050cf5e97bc506d66cc8fbf7ade43": { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "tool_calls": [] + } +} \ No newline at end of file diff --git a/example/agents/tests/features/record_no_stream.json b/example/agents/tests/features/record_no_stream.json index 37cca398..369f84ea 100644 --- a/example/agents/tests/features/record_no_stream.json +++ b/example/agents/tests/features/record_no_stream.json @@ -1,3 +1,6 @@ { - "1c27b2b038cae60eab7297ec0a808bfc67246c77a56a9e7c5ce1d9e5fdee0ba3": "Hi Alex, it's a pleasure to meet you! How can I help you today?" + "09e2753147eb813860891b4dd62050bceb302309d3f6db70019be2fb32582a47": { + "content": "Hi Alex, thanks for reaching out! How can I help you today?", + "tool_calls": [] + } } \ No newline at end of file diff --git a/example/agents/tests/features/record_stream.json b/example/agents/tests/features/record_stream.json index 764667ec..124252bd 100644 --- a/example/agents/tests/features/record_stream.json +++ b/example/agents/tests/features/record_stream.json @@ -1,6 +1,6 @@ { - "32e9c85d324f4cadef79b130717a28d0e23a85a2ac349ad2b1b78a233b31dbc4": [ + "97a14ff94f83842c3bf7a706f945aa10c9c6799ae651a74191ffc744f933a59f": [ "Hi", - " Bedram, it's a pleasure to assist you today! How may I help you?" + " Bedram, thanks for reaching out! How can I help you today?" ] } \ No newline at end of file diff --git a/example/agents/tests/features/test_chat_controller.py b/example/agents/tests/features/test_chat_controller.py index c0228d72..bbd47eee 100644 --- a/example/agents/tests/features/test_chat_controller.py +++ b/example/agents/tests/features/test_chat_controller.py @@ -1,17 +1,19 @@ +from dumpdie import dd + from app.agents.chat import RouterAgent from tests.test_case import TestCase 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,20 +23,20 @@ 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"}) response.assert_ok() response.assert_stream_contains('Hello there, This is stream chat, Hope you are doing well.') - @RouterAgent.record("record_no_stream.json") + @RouterAgent.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") - @RouterAgent.record("record_stream.json") + @RouterAgent.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." @@ -42,3 +44,11 @@ async def test_chat_responds_for_other_greetings(self): response.assert_ok() response.assert_stream_contains("Bedram") + + @RouterAgent.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/record_stream.json b/example/agents/tests/units/agents/record_stream.json new file mode 100644 index 00000000..756fda70 --- /dev/null +++ b/example/agents/tests/units/agents/record_stream.json @@ -0,0 +1,10 @@ +{ + "034751ef1322a12f3406dad43313f9cdb31f4ea85d9452c22bf7529ad8e92614": { + "content": "Hi there! How can I help you today?", + "tool_calls": [] + }, + "aaa92579b146c573996a42ba725a88d59fa731fa932e869862db333d4bc70d02": { + "content": "[{\"id\": 2, \"title\": \"Frontend Developer\", \"location\": \"Remote\", \"company\": \"Startup Inc\", \"type\": \"Full-time\"}]", + "tool_calls": [] + } +} \ No newline at end of file diff --git a/example/agents/tests/units/agents/test_langchain_agent.py b/example/agents/tests/units/agents/test_langchain_agent.py new file mode 100644 index 00000000..5e79ce6e --- /dev/null +++ b/example/agents/tests/units/agents/test_langchain_agent.py @@ -0,0 +1,26 @@ +from dumpdie import dd +from langchain.agents import create_agent +from langchain_core.language_models import GenericFakeChatModel + +from tests.test_case import TestCase + + +class Agent: + async def prompt(self, input): + from fastapi_startkit.application import app + + agent = create_agent( + model=GenericFakeChatModel(messages=iter(["hello"])), + system_prompt="You are a helpful assistant", + checkpointer=await app().make("checkpointer"), + ) + + return await agent.ainvoke( + {"messages": [{"role": "user", "content": "hello"}]}, {"configurable": {"thread_id": "1"}} + ) + + +class TestLangchainAgent(TestCase): + async def test_the_router_agent(self): + response = await Agent().prompt("hello") + dd(response) diff --git a/example/agents/tests/units/agents/test_router_agent.py b/example/agents/tests/units/agents/test_router_agent.py new file mode 100644 index 00000000..a5b8f7e1 --- /dev/null +++ b/example/agents/tests/units/agents/test_router_agent.py @@ -0,0 +1,32 @@ +from fastapi_startkit.ai.tinker import ToolCall +from langchain_core.messages import AIMessage, HumanMessage + +from app.agents.chat import RouterAgent +from tests.test_case import TestCase + + +class TestRouterAgent(TestCase): + 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_time_lt(5) + + def assert_tool_calls(tool: ToolCall): + return tool.name == "job_search_tool" + + await agent.prompt("suggest python developer jobs") + agent.assert_tool_called("job_search_tool", assert_tool_calls) + + 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") diff --git a/example/agents/tinker.py b/example/agents/tinker.py new file mode 100644 index 00000000..d2fa09f4 --- /dev/null +++ b/example/agents/tinker.py @@ -0,0 +1,62 @@ +import asyncio +import operator +from collections.abc import Callable +from typing import Annotated, TypedDict + +from dumpdie import dump +from langchain.agents import create_agent +from langchain_core.messages import AnyMessage +from langgraph.checkpoint.memory import InMemorySaver + +from app.tools.job_search_tool import job_search_tool +from bootstrap.application import app # NOQA + + +class ChatState(TypedDict): + messages: Annotated[list[AnyMessage], operator.add] + llm_calls: int + + +agent = create_agent( + model="google_genai:gemini-3.1-flash-lite", + checkpointer=InMemorySaver(), + tools=[job_search_tool], +) + + +async def prompt(message: str): + config = {"configurable": {"thread_id": "1"}} + return await agent.ainvoke(input={"messages": [{"role": "user", "content": message}]}, config=config) + + +class Agent: + def __init__(self, prompt_handler=Callable): + self.prompt_handler = prompt_handler + + async def prompt(self, message): + pass + + async def ainvoke(self): + await prompt(message="suggest me frontend developer jobs") + + +async def main(): + response = await prompt(message="Hello, world!") + dump(response["messages"]) + response = await prompt(message="suggest me frontend developer jobs") + dump(response["messages"]) + + +asyncio.run(main()) + + +# def test_it_can_prompt(): +# with Agent(prompt_handler=prompt) as agent: +# agent.prompt("hi") # hit the real end point for the first time, records the responses +# agent.assert_prompted("hi") +# agent.assert_prompt_judged(model="", expectation="") +# +# agent.prompt("suggest me python developer jobs") # hit for the first time and second records will be recorded +# agent.assert_tool_called( +# lambda tool: tool.name == "job_search_tool" and tool.args == {"query": "python developer jobs"} +# ) diff --git a/example/agents/uv.lock b/example/agents/uv.lock index 1edf626f..302e78f4 100644 --- a/example/agents/uv.lock +++ b/example/agents/uv.lock @@ -417,7 +417,7 @@ wheels = [ [[package]] name = "fastapi-startkit" -version = "0.45.0" +version = "0.50.0" source = { editable = "../../fastapi_startkit" } dependencies = [ { name = "cleo" }, @@ -433,6 +433,7 @@ dependencies = [ ai = [ { name = "langchain" }, { name = "langchain-core" }, + { name = "langgraph" }, ] database = [ { name = "faker" }, @@ -455,13 +456,14 @@ requires-dist = [ { name = "dotenv", specifier = ">=0.9.9" }, { name = "dotty-dict", specifier = ">=1.3.1" }, { name = "faker", marker = "extra == 'database'", specifier = ">=40.13.0" }, - { name = "fastapi", extras = ["standard"], marker = "extra == 'fastapi'", specifier = ">=0.124.4,<0.125.0" }, + { name = "fastapi", extras = ["standard"], marker = "extra == 'fastapi'", specifier = ">=0.124.4" }, { name = "inflection", specifier = ">=0.5.1" }, { name = "itsdangerous", marker = "extra == 'fastapi'", specifier = ">=2.2.0" }, { name = "jinja2", marker = "extra == 'inertia'", specifier = ">=3.1" }, { name = "jinja2", marker = "extra == 'vite'", specifier = ">=3.1" }, { name = "langchain", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "langchain-core", marker = "extra == 'ai'", specifier = ">=1.0.0" }, + { name = "langgraph", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "markupsafe", marker = "extra == 'inertia'", specifier = ">=2.0" }, { name = "pendulum", specifier = ">=3.1.0,<4.0.0" }, { name = "pydantic", specifier = ">=2.12.5" }, @@ -481,6 +483,7 @@ dev = [ { name = "itsdangerous", specifier = ">=2.2.0" }, { name = "langchain", specifier = ">=1.0.0" }, { name = "langchain-core", specifier = ">=1.0.0" }, + { name = "pyright", specifier = ">=1.1.411" }, { name = "pytest", specifier = ">=9.0.3" }, { name = "pytest-asyncio", specifier = ">=1.3.0" }, { name = "pytest-cov", specifier = ">=6.0.0" }, @@ -498,6 +501,7 @@ dependencies = [ { name = "fastapi-startkit", extra = ["ai", "database", "fastapi", "sqlite"] }, { name = "langchain" }, { name = "langchain-google-genai" }, + { name = "langgraph-checkpoint-postgres" }, ] [package.dev-dependencies] @@ -514,6 +518,7 @@ requires-dist = [ { name = "fastapi-startkit", extras = ["ai", "database", "fastapi", "sqlite"], editable = "../../fastapi_startkit" }, { name = "langchain", specifier = ">=1.3.10" }, { name = "langchain-google-genai", specifier = ">=4.2.5" }, + { name = "langgraph-checkpoint-postgres", specifier = ">=3.1.0" }, ] [package.metadata.requires-dev] @@ -917,6 +922,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, ] +[[package]] +name = "langgraph-checkpoint-postgres" +version = "3.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "langgraph-checkpoint" }, + { name = "orjson" }, + { name = "psycopg" }, + { name = "psycopg-pool" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/51/5a2dc42e8b5d5942b933b5b7237eae5a4dbc92508a04727c263dd383ad8a/langgraph_checkpoint_postgres-3.1.0.tar.gz", hash = "sha256:02bff4ab63d9dae8eab3a9640fce1d479da8965c9fba7b0dc04cb1f7c56f0a55", size = 148473, upload-time = "2026-05-12T03:40:10.599Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2f/cd/eff9b82bc3b5f62d481b437099f44f3ef7b1d907f166fb4ee25e8f84a1e7/langgraph_checkpoint_postgres-3.1.0-py3-none-any.whl", hash = "sha256:814cce2ef35d792bf07b090a95eed004f1acac0724fe6605536b13f6d1e7032c", size = 48988, upload-time = "2026-05-12T03:40:08.925Z" }, +] + [[package]] name = "langgraph-prebuilt" version = "1.1.0" @@ -1204,6 +1224,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, ] +[[package]] +name = "psycopg" +version = "3.3.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/2f/cb91e5502ec9de1de6f1b76cfbf69531932725361168bb06963620c77e2e/psycopg-3.3.4.tar.gz", hash = "sha256:e21207764952cff81b6b8bdacad9a3939f2793367fdac2987b3aac36a651b5bc", size = 165799, upload-time = "2026-05-01T23:31:55.179Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5c/e0/7b3dee031daae7743609ce3c746565d4a3ed7c2c186479eb48e34e838c64/psycopg-3.3.4-py3-none-any.whl", hash = "sha256:b6bbc25ccf05c8fad3b061d9db2ef0909a555171b84b07f29458a447253d679a", size = 213001, upload-time = "2026-05-01T23:20:50.816Z" }, +] + +[[package]] +name = "psycopg-pool" +version = "3.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/90/82/7a23d26039827ecd4ebe93905651029ddd307c5182ad59296dfb6f67b528/psycopg_pool-3.3.1.tar.gz", hash = "sha256:b10b10b7a175d5cc1592147dc5b7eec8a9e0834eb3ed2c4a92c858e2f51eb63c", size = 31661, upload-time = "2026-05-01T23:31:59.809Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/ed/89c2c620af0e1660354cd8aabf9f5b21f911597ce22acb37c805d6c86bc8/psycopg_pool-3.3.1-py3-none-any.whl", hash = "sha256:2af5b432941c4c9ad5c87b3fa410aec910ec8f7c122855897983a06c45f2e4b5", size = 40023, upload-time = "2026-05-01T23:31:53.136Z" }, +] + [[package]] name = "pyasn1" version = "0.6.3" diff --git a/fastapi_startkit/pyproject.toml b/fastapi_startkit/pyproject.toml index b7f66198..0be5d27c 100644 --- a/fastapi_startkit/pyproject.toml +++ b/fastapi_startkit/pyproject.toml @@ -86,6 +86,7 @@ inertia = [ ai = [ "langchain>=1.0.0", "langchain-core>=1.0.0", + "langgraph>=1.0.0", ] [dependency-groups] @@ -129,7 +130,7 @@ exclude = [ "**/__pycache__", "**/.venv", ] -typeCheckingMode = "basic" +typeCheckingMode = "standard" pythonVersion = "3.12" [tool.pytest.ini_options] diff --git a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py index dbf4c85e..6096c894 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/__init__.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/__init__.py @@ -6,32 +6,37 @@ from .config.ai import AIConfig from .decorators import max_steps, max_tokens, model, provider, timeout, top_p from .document import Document -from .fakes import fake_chat_model +from .graph import GraphAgent, GraphRunner, GraphState from .image import Image, ImageResponse from .image_factory import ImageFactory +from .ai import Ai +from .judge import JudgeAgent from .providers.ai_provider import AIProvider from .response import AgentResponse, AgentSnapshot -from .testing import AgentBinding, FakeAgent, NoFakeResponse, RecordingAgent +from .testing import AgentFake, AgentRecordFake, ToolCallView __all__ = [ "Agent", + "Ai", "Middleware", - "AgentBinding", + "AgentFake", "AgentResponse", "AgentSnapshot", "AIConfig", "AIProvider", "AnthropicConfig", - "FakeAgent", - "NoFakeResponse", - "RecordingAgent", + "JudgeAgent", + "AgentRecordFake", + "ToolCallView", "Audio", "AudioResponse", "AudioFactory", "Document", "ElevenLabsConfig", - "fake_chat_model", "GoogleConfig", + "GraphAgent", + "GraphRunner", + "GraphState", "Image", "ImageFactory", "ImageResponse", diff --git a/fastapi_startkit/src/fastapi_startkit/ai/agent.py b/fastapi_startkit/src/fastapi_startkit/ai/agent.py index 17c947f3..5875325e 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/agent.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/agent.py @@ -1,15 +1,17 @@ from __future__ import annotations -import fnmatch -from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, Optional, Type +from typing import TYPE_CHECKING, AsyncIterator, Optional, Type from .document import Document -from .response import AgentResponse, AgentSnapshot -from .testing import AgentBinding +from .response import AgentResponse +from .runner import BaseRunner, Runner if TYPE_CHECKING: from langchain_core.tools import BaseTool + from .testing import AgentFake, AgentRecordFake + from .types import Middleware + class Agent: provider: str | None = None @@ -19,10 +21,6 @@ class Agent: timeout: float = 30.0 top_p: float = 1.0 - def __init__(self): - self._fakes: dict[str, AgentResponse | AgentSnapshot] = {} - self._call_log: list[dict] = [] - def messages(self) -> list[dict]: return [] @@ -35,7 +33,7 @@ def schema(self) -> Optional[Type]: def tools(self) -> list[BaseTool]: return [] - def middleware(self) -> list[Callable]: + def middleware(self) -> list[Middleware]: return [] def provider_options(self) -> dict: @@ -49,34 +47,10 @@ 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, + return await self.runner().run( + 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) - async def stream( self, message: str, @@ -84,216 +58,21 @@ async def stream( model: str | None = None, 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.runner().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)) + def runner(self) -> BaseRunner: + return Runner(self) @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 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: - assert len(calls) == times, f"Expected {times} prompt call(s), got {len(calls)}" - else: - assert len(calls) > 0, "Expected at least one prompt() or stream() call, but none were made" - - 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}) - - async def _run_pipeline(self, chat_model: Any, messages: list) -> AgentResponse: - from .pipeline import Response, build_pipeline # noqa: PLC0415 - from .runner import Runner # noqa: PLC0415 - - chain = list(self.middleware()) - if not chain: - return await self._invoke(chat_model, messages) - - def core(model: Any) -> Response: - async def _run(): - result = await Runner(self, model).run(messages) - yield result - - return Response(_run) - - pipeline = build_pipeline(chain, core) - raw = await pipeline(chat_model) - return self._to_agent_response(raw) - - async def _apply_middleware( - self, - chat_model: Any, - final: Callable[[Any], Any], - ) -> AgentResponse: - chain = list(self.middleware()) - - def build(mw_list: list, fn: Callable) -> Callable: - if not mw_list: - return fn - head, *tail = mw_list - next_fn = build(tail, fn) - mw = head() if isinstance(head, type) else head - return lambda model: mw(model, next_fn) - - return await build(chain, final)(chat_model) - - def _build_instruction(self) -> str | None: - return self.instructions() - - def _build_messages( - self, - message: str, - attachments: list[Document] | None = None, - ) -> list[dict]: - messages: list[dict] = [] - - instruction = self.instructions() - if instruction: - messages.append({"role": "system", "content": instruction}) - - messages.extend(self.messages() or []) - - if message: - messages.append({"role": "user", "content": message}) - - if attachments: - content: Any = [{"type": "text", "text": message}] - for doc in attachments: - content.append(doc.to_langchain_block()) - messages.append({"role": "user", "content": content}) - - return 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) - - def _to_agent_response(self, result: Any) -> AgentResponse: - messages = result.get("messages", []) if isinstance(result, dict) else [] - final = messages[-1] if messages else result - - content = getattr(final, "content", "") - if not isinstance(content, str): - content = str(content) - - tool_calls = 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) - - def _apply_schema(self, response: AgentResponse) -> AgentResponse: - schema = self.schema() - if schema is not None and response.parsed is None and response.content: - response.parsed = self._build_schema(schema, response.content) - return response - - @staticmethod - def _build_schema(schema: Any, content: str) -> Any: - import json # noqa: PLC0415 - - if hasattr(schema, "model_validate_json"): - return schema.model_validate_json(content) - if hasattr(schema, "model_validate"): - return schema.model_validate(json.loads(content)) - return schema(**json.loads(content)) - - async def _invoke(self, chat_model: Any, messages: list[dict]) -> AgentResponse: - from .runner import Runner # noqa: PLC0415 + def record(cls, cassette: str | None = None, messages: list | None = None) -> "AgentRecordFake": + from .testing import AgentRecordFake - result = await Runner(self, chat_model).run(messages) - return self._to_agent_response(result) + return AgentRecordFake(cls(), cassette, messages) - async def _run( - self, - message: str, - model: str | None = None, - attachments: list[Document] | None = None, - provider_options: dict | None = None, - ) -> AgentResponse: - messages = self._build_messages(message, attachments) - chat_model = self._build_model(model, provider_options) - return await self._invoke(chat_model, messages) - - async def _stream( - self, - message: str, - model: str | None = None, - provider_options: dict | None = None, - ) -> AsyncIterator[str]: - from .pipeline import Response, build_pipeline # noqa: PLC0415 - from .runner import StreamRunner # noqa: PLC0415 - - messages = self._build_messages(message) - chat_model = self._build_model(model, provider_options) - chain = list(self.middleware()) - - def core(m: Any) -> Response: - return Response(lambda: StreamRunner(self, m).run(messages)) - - pipeline = build_pipeline(chain, core) if chain else core - - async for chunk in pipeline(chat_model): - yield chunk diff --git a/fastapi_startkit/src/fastapi_startkit/ai/ai.py b/fastapi_startkit/src/fastapi_startkit/ai/ai.py new file mode 100644 index 00000000..28b76345 --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/ai.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from .lab import Lab + +if TYPE_CHECKING: + from langchain_core.language_models.fake_chat_models import GenericFakeChatModel + + from .agent import Agent + + +class Ai: + _fakes: dict[str, GenericFakeChatModel] = {} + + def __init__(self) -> None: + pass + + @staticmethod + def _key(agent: "Agent | str") -> str: + return agent if isinstance(agent, str) else type(agent).__name__ + + @classmethod + def fake(cls, agent: "Agent | str", messages: list) -> Any: + from langchain_core.language_models.fake_chat_models import GenericFakeChatModel + from langchain_core.messages import AIMessage + + turns = [message if hasattr(message, "content") else AIMessage(content=str(message)) for message in messages] + model = GenericFakeChatModel(messages=iter(turns)) + cls._fakes[cls._key(agent)] = model + return model + + @classmethod + def has_fake_model_for(cls, agent: "Agent | str") -> bool: + return cls._key(agent) in cls._fakes + + @classmethod + def get_fake_model_for(cls, agent: "Agent | str") -> Any: + return cls._fakes[cls._key(agent)] + + @classmethod + def forget(cls, agent: "Agent | str") -> None: + cls._fakes.pop(cls._key(agent), None) + + @classmethod + def reset_fakes(cls) -> None: + cls._fakes.clear() + + def get_model_for( + self, + agent: "Agent", + model: str | None = None, + provider_options: dict | None = None, + ) -> Any: + if self.has_fake_model_for(agent): + return self.get_fake_model_for(agent) + return self.build(agent, model, provider_options) + + 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(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 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(agent, provider_options)) + + chat_model = init_chat_model(self._resolve_model(agent, model), **kwargs) + + tools = list(agent.tools()) + chat_model = chat_model.bind_tools(tools) if tools else chat_model + + schema = agent.schema() + if schema is not None: + chat_model = chat_model.with_structured_output(schema, include_raw=True) + + return chat_model + + 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, agent: "Agent", override: dict | None = None) -> dict: + options = dict(agent.provider_options().get(agent.provider, {})) + if override: + provider_specific = override.get(agent.provider, override) + if isinstance(provider_specific, dict): + options.update(provider_specific) + return options diff --git a/fastapi_startkit/src/fastapi_startkit/ai/fakes.py b/fastapi_startkit/src/fastapi_startkit/ai/fakes.py deleted file mode 100644 index 870e624a..00000000 --- a/fastapi_startkit/src/fastapi_startkit/ai/fakes.py +++ /dev/null @@ -1,53 +0,0 @@ -from __future__ import annotations - -import json -import re -from typing import Any, Iterable - - -def _require_langchain(): - try: - from langchain_core.language_models.fake_chat_models import GenericFakeChatModel - from langchain_core.messages import AIMessage - except ImportError as exc: - raise ImportError( - "The agent test harness requires the 'ai' extra. Install it with: pip install \"fastapi-startkit[ai]\"" - ) from exc - return GenericFakeChatModel, AIMessage - - -def fake_chat_model(turns: Iterable[Any]): - generic_model, ai_message = _require_langchain() - - from langchain_core.messages import AIMessageChunk - from langchain_core.outputs import ChatGenerationChunk - - class _FakeChatModel(generic_model): - def bind_tools(self, tools, **kwargs): - return self - - def _stream(self, messages, stop=None, run_manager=None, **kwargs): - message = next(self.messages) - if not isinstance(message, ai_message): - message = ai_message(content=str(message)) - - content = message.content if isinstance(message.content, str) else str(message.content) - for token in re.split(r"(\s)", content): - if token: - yield ChatGenerationChunk(message=AIMessageChunk(content=token, id=message.id)) - - tool_calls = list(message.tool_calls or []) - if tool_calls: - chunks = [ - { - "name": call["name"], - "args": json.dumps(call.get("args", {})), - "id": call.get("id"), - "index": index, - } - for index, call in enumerate(tool_calls) - ] - yield ChatGenerationChunk(message=AIMessageChunk(content="", tool_call_chunks=chunks, id=message.id)) - - normalized = [t if isinstance(t, ai_message) else ai_message(content=str(t)) for t in turns] - return _FakeChatModel(messages=iter(normalized)) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/graph.py b/fastapi_startkit/src/fastapi_startkit/ai/graph.py new file mode 100644 index 00000000..20ab656f --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/graph.py @@ -0,0 +1,113 @@ + +from __future__ import annotations + +import time +from abc import ABC, abstractmethod +from collections.abc import AsyncIterator +from typing import TYPE_CHECKING, Annotated, Any, TypedDict + +from langgraph.graph import END, StateGraph +from langgraph.graph.message import add_messages + +from .agent import Agent +from .response import AgentResponse +from .runner import BaseRunner + +if TYPE_CHECKING: + from .document import Document + + +class GraphState(TypedDict): + + messages: Annotated[list, add_messages] + llm_calls: int + + +class GraphAgent(Agent, ABC): + + def runner(self) -> GraphRunner: + return GraphRunner(self) + + @abstractmethod + def graph(self, runner: GraphRunner) -> StateGraph: + ... + + +class GraphRunner(BaseRunner): + + agent: GraphAgent + + def __init__(self, agent: GraphAgent) -> None: + super().__init__(agent) + self._chat_model: Any = None + + + async def llm(self, state: GraphState) -> dict: + reply = await self._chat_model.ainvoke(state["messages"]) + return {"messages": [reply], "llm_calls": state.get("llm_calls", 0) + 1} + + async def call_tools(self, state: GraphState) -> 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)) + return {"messages": results} + + def route(self, state: GraphState) -> str: + return "tools" if getattr(state["messages"][-1], "tool_calls", None) else END + + + def _compile(self, model: str | None, provider_options: dict | None) -> Any: + self._chat_model = self._build_model(model, provider_options) + return self.agent.graph(self).compile() + + @property + def _config(self) -> dict: + return {"recursion_limit": self.agent.max_steps * 2 + 1} + + async def run( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: + started = time.perf_counter() + compiled = self._compile(model, provider_options) + state = {"messages": self._build_messages(message, attachments), "llm_calls": 0} + result = await compiled.ainvoke(state, config=self._config) + response = self._apply_schema(self._to_response(result)) + response.runtime = time.perf_counter() - started + return response + + async def stream( + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: + compiled = self._compile(model, provider_options) + state = {"messages": self._build_messages(message), "llm_calls": 0} + async for chunk, _meta in compiled.astream(state, stream_mode="messages", config=self._config): + text = chunk.content if isinstance(chunk.content, str) else str(chunk.content) + if text: + yield text + + @staticmethod + def _to_response(result: dict) -> AgentResponse: + messages = result.get("messages", []) + final = messages[-1] if messages else None + content = getattr(final, "content", "") or "" + if not isinstance(content, str): + content = str(content) + + tool_calls = [call for m in messages for call in (getattr(m, "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) diff --git a/fastapi_startkit/src/fastapi_startkit/ai/judge.py b/fastapi_startkit/src/fastapi_startkit/ai/judge.py new file mode 100644 index 00000000..6923bd3d --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/judge.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from pydantic import BaseModel + +from .agent import Agent + + +class Verdict(BaseModel): + passed: bool + reasoning: str = "" + + +class JudgeAgent(Agent): + """Grades a response against a natural-language expectation. + + A plain ``Agent`` whose ``schema()`` is a ``Verdict`` model, so the JSON + reply is parsed into a typed result through the standard structured-output + path — no hand-rolled verdict parsing. Set ``.model``/``.provider`` like + any other agent; it's fakeable via ``fake()`` and replayable via + ``record()`` for free. + """ + + def instructions(self) -> str: + return ( + "You are grading whether an AI agent's response satisfies an expectation. " + 'Reply with strict JSON only, no prose: {"passed": true|false, "reasoning": ""}' + ) + + def schema(self): + return Verdict + + async def judge(self, expectation: str, content: str) -> dict: + response = await self.prompt(f"Expectation: {expectation}\n\nResponse to grade:\n{content}") + return response.parsed.model_dump() diff --git a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py b/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py deleted file mode 100644 index 7f82254f..00000000 --- a/fastapi_startkit/src/fastapi_startkit/ai/model_builder.py +++ /dev/null @@ -1,47 +0,0 @@ -from __future__ import annotations - -from typing import TYPE_CHECKING, Any - -from .lab import Lab - -if TYPE_CHECKING: - from .agent import Agent - - -class ModelBuilder: - def __init__(self, agent: "Agent") -> None: - self._agent = agent - - def build(self, 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) - 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 - - kwargs.update(self._resolve_provider_options(provider_options)) - - chat_model = init_chat_model(self._resolve_model(model), **kwargs) - - tools = list(self._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_provider_options(self, override: dict | None = None) -> dict: - options = dict(self._agent.provider_options().get(self._agent.provider, {})) - if override: - provider_specific = override.get(self._agent.provider, override) - if isinstance(provider_specific, dict): - options.update(provider_specific) - return options diff --git a/fastapi_startkit/src/fastapi_startkit/ai/response.py b/fastapi_startkit/src/fastapi_startkit/ai/response.py index 46d02ebd..57e2adfa 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/response.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/response.py @@ -1,5 +1,3 @@ -"""AgentResponse and AgentSnapshot — response containers for AI agents.""" - from __future__ import annotations import json @@ -13,20 +11,21 @@ @dataclass class AgentResponse: - """Returned by Agent.prompt(). Wraps the LLM response.""" - content: str = "" tool_calls: list[dict] = field(default_factory=list) usage: dict = field(default_factory=dict) raw: Any = None parsed: Any = None + runtime: float = 0.0 + + @property + def runtime_ms(self) -> float: + return self.runtime * 1000 def text(self) -> str: - """Return the text content.""" return self.content def json(self) -> Any: - """Parse the content as JSON.""" return json.loads(self.content) def __str__(self) -> str: @@ -38,37 +37,23 @@ def __bool__(self) -> bool: @dataclass class AgentSnapshot: - """ - Record-and-replay snapshot for testing. - - - If the file at ``path`` **does not exist**: the agent calls the real API, - saves the response as JSON, then returns it. - - If the file **exists**: the saved response is loaded and returned without - hitting the API. - - Example:: - - agent.fake({"*analyze*": AgentSnapshot(path="tests/fixtures/analysis.json")}) - """ path: str def exists(self) -> bool: - """Return True if the snapshot file is already recorded.""" return os.path.exists(self.path) def load(self) -> AgentResponse: - """Load the recorded response from disk.""" with open(self.path) as f: data = json.load(f) return AgentResponse( content=data.get("content", ""), tool_calls=data.get("tool_calls", []), usage=data.get("usage", {}), + runtime=data.get("runtime", 0.0), ) def save(self, response: AgentResponse) -> None: - """Persist a real API response to disk for future replays.""" os.makedirs(os.path.dirname(self.path) or ".", exist_ok=True) with open(self.path, "w") as f: json.dump( @@ -76,18 +61,15 @@ def save(self, response: AgentResponse) -> None: "content": response.content, "tool_calls": response.tool_calls, "usage": response.usage, + "runtime": response.runtime, }, f, indent=2, ) async def resolve(self, agent: "Agent", message: str, **run_kwargs: Any) -> AgentResponse: - """ - Return the response — from disk if recorded, or from the real API - (which is then saved for future runs). - """ if self.exists(): return self.load() - response = await agent._run(message, **run_kwargs) + response = await agent.prompt(message, **run_kwargs) self.save(response) return response diff --git a/fastapi_startkit/src/fastapi_startkit/ai/runner.py b/fastapi_startkit/src/fastapi_startkit/ai/runner.py index 49fe5a44..ae35bfea 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/runner.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/runner.py @@ -1,14 +1,21 @@ from __future__ import annotations -from collections.abc import AsyncIterator, Sequence +import time +from abc import ABC, abstractmethod +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 .pipeline import Response, build_pipeline +from .response import AgentResponse + if TYPE_CHECKING: from .agent import Agent + from .document import Document Message = BaseMessage | dict[str, Any] @@ -17,43 +24,181 @@ def _as_text(content: Any) -> str: return content if isinstance(content, str) else str(content) -class Runner: - def __init__(self, agent: Agent, model: Runnable[Any, BaseMessage]) -> None: - self._tools: dict[str, BaseTool] = {tool.name: tool for tool in agent.tools()} - self.model: Runnable[Any, BaseMessage] = model - self.max_steps = agent.max_steps - - async def run(self, messages: Sequence[Message]) -> BaseMessage: - history: list[Message] = list(messages) - response: AIMessage = await self.model.ainvoke(history) # type: ignore[assignment] - +class BaseRunner(ABC): + def __init__(self, agent: Agent) -> None: + self.agent = agent + + def _build_messages(self, message: str, attachments: list[Document] | None = None) -> list[dict]: + agent = self.agent + messages: list[dict] = [] + + instruction = agent.instructions() + if instruction: + messages.append({"role": "system", "content": instruction}) + + messages.extend(agent.messages() or []) + + if message: + messages.append({"role": "user", "content": message}) + + if attachments: + content: Any = [{"type": "text", "text": message}] + for doc in attachments: + content.append(doc.to_langchain_block()) + messages.append({"role": "user", "content": content}) + + return messages + + def _build_model(self, model: str | None = None, provider_options: dict | None = None) -> Any: + from .ai import Ai # noqa: PLC0415 + + return Ai().get_model_for(self.agent, model, provider_options) + + def _apply_schema(self, response: AgentResponse) -> AgentResponse: + schema = self.agent.schema() + if schema is not None and response.parsed is None and response.content: + response.parsed = self._build_schema(schema, response.content) + return response + + @staticmethod + def _build_schema(schema: Any, content: str) -> Any: + import json # noqa: PLC0415 + + if hasattr(schema, "model_validate_json"): + return schema.model_validate_json(content) + if hasattr(schema, "model_validate"): + return schema.model_validate(json.loads(content)) + return schema(**json.loads(content)) + + @abstractmethod + async def run( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: ... + + @abstractmethod + def stream( + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: ... + + +class Runner(BaseRunner): + async def run( + self, + message: str, + *, + model: str | None = None, + attachments: list[Document] | None = None, + provider_options: dict | None = None, + ) -> AgentResponse: + started = time.perf_counter() + 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 + return response + + async def _run_pipeline(self, chat_model: Any, messages: list) -> AgentResponse: + chain = list(self.agent.middleware()) + if not chain: + raw = await self._invoke(chat_model, messages) + return self._to_agent_response(raw) + + def core(model: Any) -> Response: + async def _run() -> AsyncIterator[Any]: + yield await self._invoke(model, messages) + + return Response(_run) + + pipeline = build_pipeline(chain, core) + raw = await pipeline(chat_model) + return self._to_agent_response(raw) + + def _to_agent_response(self, result: Any) -> AgentResponse: + parsed = None + structured = isinstance(result, dict) and "parsed" in result and "raw" in result + if structured: + parsed = result.get("parsed") + result = result.get("raw") + + messages = result.get("messages", []) if isinstance(result, dict) else [] + final = messages[-1] if messages else result + + content = getattr(final, "content", "") + if not isinstance(content, str): + content = str(content) + 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) + + async def stream( + self, + message: str, + *, + model: str | None = None, + provider_options: dict | None = None, + ) -> AsyncIterator[str]: + messages = self._build_messages(message) + chat_model = self._build_model(model, provider_options) + chain = list(self.agent.middleware()) + + def core(m: Any) -> Response: + return Response(lambda: self._invoke_stream(m, messages)) + + pipeline = build_pipeline(chain, core) if chain else core + + async for chunk in pipeline(chat_model): + yield chunk + + async def _invoke(self, model: Runnable[Any, BaseMessage], messages: list) -> BaseMessage: + response: AIMessage = await model.ainvoke(list(messages)) # type: ignore[assignment] + if isinstance(response, dict) and "parsed" in response: + return response # type: ignore[return-value] if not response.tool_calls: return response - return (await self._run_tools(response.tool_calls))[-1] - async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: - return [await self._resolve_tool(call["name"]).ainvoke(call) for call in tool_calls] - - def _resolve_tool(self, name: str) -> BaseTool: - try: - return self._tools[name] - except KeyError: - raise ValueError(f"Agent has no tool named {name!r}") from None - - -class StreamRunner(Runner): - async def run(self, messages: Sequence[Message]) -> AsyncIterator[str]: # type: ignore[override] - history: list[Message] = list(messages) - + async def _invoke_stream(self, model: Runnable[Any, BaseMessage], messages: list) -> AsyncIterator[str]: gathered: AIMessageChunk | None = None - async for chunk in self.model.astream(history): + async for chunk in model.astream(list(messages)): if chunk.content: yield _as_text(chunk.content) gathered = chunk if gathered is None else gathered + chunk # type: ignore[operator] if gathered is None or not gathered.tool_calls: return - for message in await self._run_tools(gathered.tool_calls): yield _as_text(message.content) + + async def _run_tools(self, tool_calls: list[ToolCall]) -> list[BaseMessage]: + tools: dict[str, BaseTool] = {tool.name: tool for tool in self.agent.tools()} + results: list[BaseMessage] = [] + for call in tool_calls: + try: + 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)) + return results diff --git a/fastapi_startkit/src/fastapi_startkit/ai/testing.py b/fastapi_startkit/src/fastapi_startkit/ai/testing.py index e3e8533b..32f00f13 100644 --- a/fastapi_startkit/src/fastapi_startkit/ai/testing.py +++ b/fastapi_startkit/src/fastapi_startkit/ai/testing.py @@ -1,12 +1,11 @@ from __future__ import annotations -import fnmatch import functools import hashlib import inspect import json -import re import sys +import time from collections.abc import AsyncIterator from pathlib import Path from typing import TYPE_CHECKING, Any, Callable @@ -18,97 +17,204 @@ from .document import Document -class NoFakeResponse(LookupError): - pass +def _joined(value: Any) -> str: + return "".join(value) if isinstance(value, list) else value -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 +class AgentFake: + def __init__(self, agent_cls: type[Agent], responses: list) -> None: + self._agent_cls = agent_cls + self._responses = list(responses) + self._agent = agent_cls() + self._agent.messages = self._history # type: ignore[method-assign] + self._records: list[dict] = [] + self._last_response: AgentResponse | None = None + self.last_elapsed: float | None = None -def _reply_text(reply: Any) -> str: - if isinstance(reply, AgentResponse): - return reply.content - return getattr(reply, "content", None) or str(reply) + def _history(self) -> list: + return self._records + @property + def _prompts(self) -> list[str]: + return [r["content"] for r in self._records if r.get("role") == "user"] -class _Recorder: - def __init__(self) -> None: - self.calls: list[str] = [] - self.attachments: list[list[Document]] = [] + def __enter__(self) -> "AgentFake": + from .ai import Ai - def _record_call(self, message: str, attachments: list[Document] | None) -> None: - self.calls.append(message) - self.attachments.append(list(attachments or [])) + Ai.fake(self._agent_cls.__name__, self._responses) + return self - @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 __exit__(self, *_exc: Any) -> bool: + from .ai import Ai + + Ai.forget(self._agent_cls.__name__) + return False + + async def prompt(self, message: str, *, attachments: list[Document] | None = None) -> AgentResponse: + start = time.monotonic() + self._last_response = await self._agent.prompt(message, attachments=attachments) + self.last_elapsed = time.monotonic() - start + self._remember(message, self._last_response) + return self._last_response + + async def stream(self, message: str) -> AsyncIterator[str]: + chunks: list[str] = [] + async for chunk in self._agent.stream(message): + chunks.append(chunk) + yield chunk + self._last_response = AgentResponse(content="".join(chunks)) + self._remember(message, self._last_response) + + 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), ( + f"No recorded prompt satisfied the predicate. Got: {self._prompts!r}" + ) + else: + assert any(expected in p for p in self._prompts), ( + f"Expected a prompt containing {expected!r}, but none did. Got: {self._prompts!r}" + ) + + def assert_response(self, expected: str) -> None: + response = self._require_response() + assert expected in response.content, f"Expected response to contain {expected!r}, got {response.content!r}" + + def assert_tool_call(self, name: str) -> None: + response = self._require_response() + called = [tc.get("name") for tc in response.tool_calls] + assert name in called, f"Expected tool {name!r} to be called, but got: {called}" + + def assert_prompted(self, times: int | None = None) -> None: + if times is not None: + assert len(self._prompts) == times, f"Expected {times} prompt call(s), got {len(self._prompts)}" + else: + assert self._prompts, "Expected at least one prompt() or stream() call, but none were made" def assert_not_prompted(self) -> None: - assert not self.calls, f"Expected no prompts, but got: {self.calls!r}" + self.assert_prompted(times=0) + + def reset(self) -> "AgentFake": + self._records.clear() + self._last_response = None + return self + + def _require_response(self) -> AgentResponse: + assert self._last_response is not None, "No prompt() call has been made yet." + return self._last_response + + def _tool_call_names(self) -> list[str]: + return [tc.get("name", "") for tc in self._require_response().tool_calls] + + 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: + 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), ( + f"Tool {name!r} was called, but no call satisfied the given predicate." + ) + + def assert_tool_not_called(self, names: list[str]) -> None: + unexpected = set(self._tool_call_names()) & set(names) + assert not unexpected, f"Expected tools {sorted(names)} not to be called, but got: {sorted(unexpected)}" + + def assert_response_time_lt(self, seconds: float) -> None: + assert self.last_elapsed is not None, "No prompt() call has been made yet." + assert self.last_elapsed < seconds, f"Expected response time < {seconds}s, took {self.last_elapsed:.3f}s" + + async def assert_response_judged(self, *, model: str, expectation: str, provider: str | None = None) -> None: + response = self._require_response() + verdict = await self._judge(model, expectation, response.content, provider) + assert verdict.get("passed"), ( + f"Judge ({model}) rejected the response for expectation {expectation!r}: " + f"{verdict.get('reasoning', '')!r} — response was {response.content!r}" + ) + async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: + return await self._judge_live(model, expectation, content, provider) -def _joined(value: Any) -> str: - """A cassette value is a buffered string or a list of stream chunks.""" - return "".join(value) if isinstance(value, list) else value + async def _judge_live(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: + from .judge import JudgeAgent # noqa: PLC0415 + + judge = JudgeAgent() + judge.model = model + judge.provider = provider + return await judge.judge(expectation, content) + + 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 -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] + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + with self: + return func(*args, **kwargs) + return wrapper -class FakeAgent(_Recorder): - def __init__(self, responses: dict[str, Any] | None = None) -> None: - super().__init__() - self.responses = responses or {} - 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}") +class ToolCallView: - async def prompt(self, message: str, attachments: list[Document] | None = None) -> AgentResponse: - self._record_call(message, attachments) - return AgentResponse(content=self._resolve(message)) + def __init__(self, data: dict) -> None: + self.name = data.get("name", "") + self.args = data.get("args") or {} + self.id = data.get("id") + self._data = data - async def stream(self, message: str) -> AsyncIterator[str]: - self._record_call(message, None) - for chunk in _word_chunks(self._resolve(message)): - yield chunk + def __repr__(self) -> str: + return f"ToolCallView(name={self.name!r}, args={self.args!r})" -class RecordingAgent(_Recorder): - def __init__(self, real: Agent, cassette: str | None = None) -> None: - super().__init__() +class AgentRecordFake(AgentFake): + def __init__(self, real: Agent, cassette: str | None = None, messages: list | None = None) -> None: self._real = real self.cassette: Path | None = Path(cassette) if cassette else None + self._seed_messages: list = list(messages or []) + self._records: list[dict] = [] + self._real.messages = self._history # type: ignore[method-assign] + self._last_response: AgentResponse | None = None + self.last_elapsed: float | None = None + + def _history(self) -> list: + return self._seed_messages + self._records @staticmethod - def _key(message: str, attachments: list[Document] | None) -> str: + def _serialize(value: Any) -> Any: + if isinstance(value, dict): + return value + return {"type": type(value).__name__, "content": getattr(value, "content", str(value))} + + def _key(self, 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) + payload = json.dumps( + { + "history": [self._serialize(m) for m in self._history()], + "message": message, + "attachments": names, + }, + sort_keys=True, + ) 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" + assert cassette is not None, "AgentRecordFake has no cassette resolved" return cassette, (json.loads(cassette.read_text()) if cassette.exists() else {}) def _save(self, cassette: Path, store: dict, key: str, value: Any) -> None: @@ -116,58 +222,81 @@ 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)) - async def prompt(self, message: str, attachments: list[Document] | None = None) -> AgentResponse: - self._record_call(message, attachments) + @staticmethod + def _cache_prompt_value(response: AgentResponse) -> dict: + return {"content": response.content, "tool_calls": response.tool_calls} + + @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 []) + return AgentResponse(content=_joined(value)) + + def _remember_turn(self, message: str, response: AgentResponse) -> None: + self._records.append({"role": "user", "content": message}) + turn: dict[str, Any] = {"role": "assistant", "content": response.content} + if response.tool_calls: + turn["tool_calls"] = response.tool_calls + self._records.append(turn) + + async def prompt(self, message: str, *, attachments: list[Document] | None = None) -> AgentResponse: cassette, store = self._load() key = self._key(message, attachments) + start = time.monotonic() 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) + response = self._response_from_cache(store[key]) + else: + response = await self._real.prompt(message, attachments=attachments) + self._save(cassette, store, key, self._cache_prompt_value(response)) + response = self._real.runner()._apply_schema(response) + self.last_elapsed = time.monotonic() - start + self._last_response = response + self._remember_turn(message, response) return response async def stream(self, message: str) -> AsyncIterator[str]: - self._record_call(message, None) cassette, store = self._load() key = self._key(message, None) if key in store: value = store[key] - for chunk in value if isinstance(value, list) else [value]: - yield chunk - return - chunks = [chunk async for chunk in self._real._stream(message)] - self._save(cassette, store, key, chunks) + chunks = value if isinstance(value, list) else [value] + else: + chunks = [chunk async for chunk in self._real.stream(message)] + self._save(cassette, store, key, chunks) for chunk in chunks: yield chunk + self._remember_turn(message, AgentResponse(content=_joined(chunks))) + async def _judge(self, model: str, expectation: str, content: str, provider: str | None = None) -> dict: + cassette, store = self._load() + key = self._judge_key(model, expectation, content, provider) + if key in store: + return store[key] + verdict = await self._judge_live(model, expectation, content, provider) + self._save(cassette, store, key, verdict) + return verdict -class AgentBinding: - def __init__(self, agent_cls: type[Agent], stand_in: Any) -> None: - self._agent_cls = agent_cls - self._stand_in = stand_in + @staticmethod + def _judge_key(model: str, expectation: str, content: str, provider: str | None = None) -> str: + payload = json.dumps( + {"judge_model": model, "judge_provider": provider, "expectation": expectation, "content": content}, + sort_keys=True, + ) + return "judge:" + hashlib.sha256(payload.encode()).hexdigest() 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 - - def __enter__(self) -> Any: - from fastapi_startkit.application import app + 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 __enter__(self) -> "AgentRecordFake": 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 + return self def __exit__(self, *_exc: Any) -> bool: - from fastapi_startkit.application import app - - app().unbind(self._agent_cls.__name__) return False def __call__(self, func: Callable) -> Callable: diff --git a/fastapi_startkit/src/fastapi_startkit/ai/tinker.py b/fastapi_startkit/src/fastapi_startkit/ai/tinker.py new file mode 100644 index 00000000..e69de29b diff --git a/fastapi_startkit/src/fastapi_startkit/ai/types.py b/fastapi_startkit/src/fastapi_startkit/ai/types.py new file mode 100644 index 00000000..62419dbc --- /dev/null +++ b/fastapi_startkit/src/fastapi_startkit/ai/types.py @@ -0,0 +1,5 @@ +from collections.abc import AsyncIterator +from typing import Callable + +Handler = Callable[[list], AsyncIterator] +Middleware = Callable[[list, Handler], AsyncIterator] diff --git a/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py b/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py index 44fc49cc..5f32852b 100644 --- a/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py +++ b/fastapi_startkit/src/fastapi_startkit/configuration/Configuration.py @@ -1,5 +1,5 @@ from fastapi_startkit.loader import Loader -from ..utils.structures import data +from ..support.structures import data from ..exceptions import InvalidConfigurationSetup diff --git a/fastapi_startkit/src/fastapi_startkit/loader/Loader.py b/fastapi_startkit/src/fastapi_startkit/loader/Loader.py index 07b01cde..4b34dfbb 100644 --- a/fastapi_startkit/src/fastapi_startkit/loader/Loader.py +++ b/fastapi_startkit/src/fastapi_startkit/loader/Loader.py @@ -4,7 +4,7 @@ import pkgutil from ..exceptions import LoaderNotFound -from ..utils.structures import load +from ..support.structures import load def parameters_filter(obj_name, obj): diff --git a/fastapi_startkit/src/fastapi_startkit/utils/structures.py b/fastapi_startkit/src/fastapi_startkit/support/structures.py similarity index 100% rename from fastapi_startkit/src/fastapi_startkit/utils/structures.py rename to fastapi_startkit/src/fastapi_startkit/support/structures.py diff --git a/fastapi_startkit/tests/ai/test_agent.py b/fastapi_startkit/tests/ai/test_agent.py index 91a2203a..83c8029f 100644 --- a/fastapi_startkit/tests/ai/test_agent.py +++ b/fastapi_startkit/tests/ai/test_agent.py @@ -1,17 +1,40 @@ +import json +import re import unittest +from typing import cast from unittest import mock import langchain.chat_models as chat_models -from langchain_core.messages import AIMessage, ToolCall +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, AIMessageChunk, ToolCall +from langchain_core.outputs import ChatGenerationChunk from langchain_core.tools import tool -from fastapi_startkit.ai import AIConfig, Document, fake_chat_model +from fastapi_startkit.ai import AIConfig, Document from fastapi_startkit.ai.agent import Agent -from fastapi_startkit.ai.model_builder import ModelBuilder +from fastapi_startkit.ai.runner import Runner +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse from fastapi_startkit.application import app +class StreamingToolFake(GenericFakeChatModel): + # GenericFakeChatModel can't stream tool calls; this emits tool-call chunks + # so streamed-tool-result behaviour stays testable without a shared fakes module. + def _stream(self, messages, stop=None, run_manager=None, **kwargs): + message = cast(AIMessage, next(self.messages)) + content = message.content if isinstance(message.content, str) else str(message.content) + for token in re.split(r"(\s)", content): + if token: + yield ChatGenerationChunk(message=AIMessageChunk(content=token, id=message.id)) + if message.tool_calls: + chunks = [ + {"name": c["name"], "args": json.dumps(c.get("args", {})), "id": c.get("id"), "index": i} + for i, c in enumerate(message.tool_calls) + ] + yield ChatGenerationChunk(message=AIMessageChunk(content="", tool_call_chunks=chunks, id=message.id)) + + @tool def search_jobs(query: str) -> str: """Search the job board for roles matching the query.""" @@ -32,12 +55,9 @@ def setUp(self): container.bind("ai", AIConfig()) container.make("config").set("ai", AIConfig()) - def setup_agent(self, turns: list[AIMessage]): - model = fake_chat_model(turns) - patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: model) - patcher.start() - self.addCleanup(patcher.stop) - return model + def setup_agent(self, turns: list, agent_cls: type[Agent] = Agent): + Ai.fake(agent_cls.__name__, turns) + self.addCleanup(Ai.reset_fakes) async def test_prompt_returns_agent_response(self): self.setup_agent([AIMessage(content="hello back")]) @@ -47,7 +67,6 @@ async def test_prompt_returns_agent_response(self): self.assertIsInstance(result, AgentResponse) self.assertEqual(result.content, "hello back") - agent.assert_prompted() async def test_search_jobs_tool_returns_listing(self): self.setup_agent( @@ -56,7 +75,8 @@ async def test_search_jobs_tool_returns_listing(self): content="", tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], ), - ] + ], + JobAssistant, ) result = await JobAssistant().prompt("find me a python job") @@ -79,7 +99,7 @@ async def test_build_model_passes_langchain_provider_key(self): def fake_init(model, **kwargs): captured["model"] = model captured["provider"] = kwargs.get("model_provider") - return fake_chat_model([AIMessage(content="ok")]) + return GenericFakeChatModel(messages=iter([AIMessage(content="ok")])) patcher = mock.patch.object(chat_models, "init_chat_model", fake_init) patcher.start() @@ -101,8 +121,6 @@ async def test_stream_yields_tokens_from_the_model(self): self.assertEqual("".join(chunks), "streamed reply") async def test_middleware_streams_token_by_token_and_runs_after_hook(self): - self.setup_agent([AIMessage(content="one two three")]) - events: list = [] class Logger: @@ -114,6 +132,8 @@ class LoggedAgent(Agent): def middleware(self): return [Logger()] + self.setup_agent([AIMessage(content="one two three")], LoggedAgent) + chunks = [chunk async for chunk in LoggedAgent().stream("hi")] # Middleware must not buffer: the model's tokens arrive as separate chunks... @@ -123,8 +143,6 @@ def middleware(self): self.assertEqual(events, ["before", "after"]) async def test_middleware_after_hook_runs_on_prompt(self): - self.setup_agent([AIMessage(content="done")]) - events: list = [] class Logger: @@ -136,38 +154,43 @@ class LoggedAgent(Agent): def middleware(self): return [Logger()] + self.setup_agent([AIMessage(content="done")], LoggedAgent) + result = await LoggedAgent().prompt("hi") self.assertEqual(result.content, "done") self.assertEqual(events, ["before", "after"]) async def test_stream_yields_tool_result_without_calling_model_again(self): - self.setup_agent( - [ - AIMessage( - content="", - tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], - ), - ] + Ai._fakes["JobAssistant"] = StreamingToolFake( + messages=iter( + [ + AIMessage( + content="", + tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], + ), + ] + ) ) + self.addCleanup(Ai.reset_fakes) chunks = [chunk async for chunk in JobAssistant().stream("find me a python job")] self.assertEqual(chunks, ["Python Developer at Shopify"]) def test_resolve_model_falls_back_to_lab_default(self): - self.assertEqual(ModelBuilder(Agent())._resolve_model(), "gemini-2.5-flash-lite") + self.assertEqual(Ai()._resolve_model(Agent()), "gemini-2.5-flash-lite") class AnthropicAgent(Agent): provider = "anthropic" - self.assertEqual(ModelBuilder(AnthropicAgent())._resolve_model(), "claude-sonnet-4-6") + self.assertEqual(Ai()._resolve_model(AnthropicAgent()), "claude-sonnet-4-6") def test_resolve_model_prefers_explicit_override(self): - self.assertEqual(ModelBuilder(Agent())._resolve_model("my-model"), "my-model") + self.assertEqual(Ai()._resolve_model(Agent(), "my-model"), "my-model") def test_instructions_lead_the_message_list(self): - messages = JobAssistant()._build_messages("find me a job") + messages = Runner(JobAssistant())._build_messages("find me a job") self.assertEqual(messages[0], {"role": "system", "content": "You help users find jobs."}) self.assertEqual(sum(m.get("role") == "system" for m in messages), 1) @@ -177,19 +200,19 @@ class DynamicAgent(Agent): def instructions(self) -> str: return "Computed identity." - messages = DynamicAgent()._build_messages("hi") + messages = Runner(DynamicAgent())._build_messages("hi") self.assertEqual(messages[0], {"role": "system", "content": "Computed identity."}) def test_no_instructions_prepends_no_system_message(self): - messages = Agent()._build_messages("hi") + messages = Runner(Agent())._build_messages("hi") self.assertTrue(all(m.get("role") != "system" for m in messages)) def test_build_messages_inlines_text_attachment(self): doc = Document(content="Q3 revenue was $1.2M.", name="q3-report.txt") - messages = Agent()._build_messages("Summarise this report.", attachments=[doc]) + messages = Runner(Agent())._build_messages("Summarise this report.", attachments=[doc]) user_content = messages[-1]["content"] self.assertEqual(user_content[0], {"type": "text", "text": "Summarise this report."}) @@ -199,7 +222,7 @@ def test_build_messages_inlines_text_attachment(self): def test_build_messages_encodes_binary_attachment_as_file_block(self): doc = Document(content=b"%PDF-1.7 ...", name="q3.pdf", media_type="application/pdf") - messages = Agent()._build_messages("Summarise", attachments=[doc]) + messages = Runner(Agent())._build_messages("Summarise", attachments=[doc]) block = messages[-1]["content"][1] self.assertEqual(block["type"], "file") diff --git a/fastapi_startkit/tests/ai/test_agent_fake.py b/fastapi_startkit/tests/ai/test_agent_fake.py index 3191a8b3..6bcd8187 100644 --- a/fastapi_startkit/tests/ai/test_agent_fake.py +++ b/fastapi_startkit/tests/ai/test_agent_fake.py @@ -1,11 +1,3 @@ -"""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. -""" - import json import os import tempfile @@ -13,6 +5,7 @@ from unittest import mock from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse @@ -21,175 +14,145 @@ class SimpleAgent(Agent): class TestAgentFake(unittest.IsolatedAsyncioTestCase): - async def test_fake_with_agent_response_returns_it(self): + def tearDown(self): + Ai.reset_fakes() + + async def test_fake_replays_the_only_response(self): agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="Hello world!")}): + with SimpleAgent.fake(["Hello world!"]): result = await agent.prompt("anything") 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 test_fake_does_not_call_provider_build(self): + import langchain.chat_models as chat_models - async def patched_run(*args, **kwargs): - called.append(True) - return await original_run(*args, **kwargs) + def fail_if_called(*args, **kwargs): + raise AssertionError("init_chat_model must not be called when a fake is registered") - 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") + patcher = mock.patch.object(chat_models, "init_chat_model", fail_if_called) + patcher.start() + self.addCleanup(patcher.stop) - async def test_fake_with_exact_pattern(self): agent = SimpleAgent() - with SimpleAgent.fake({"hello": AgentResponse(content="matched hello")}): - result = await agent.prompt("hello") - - self.assertEqual(result.content, "matched hello") + with SimpleAgent.fake(["faked"]): + await agent.prompt("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 reply", "second reply"]): + first = await agent.prompt("call one") + second = await agent.prompt("call two") - self.assertEqual(result.content, "hi there") + self.assertEqual(first.content, "first reply") + self.assertEqual(second.content, "second reply") - async def test_fake_glob_analyze_wildcard(self): + async def test_fake_raises_once_responses_are_exhausted(self): agent = SimpleAgent() - with SimpleAgent.fake({"*analyze*": AgentResponse(content="analysis done")}): - result = await agent.prompt("please analyze this report") + with SimpleAgent.fake(["only reply"]): + await agent.prompt("call one") - self.assertEqual(result.content, "analysis done") + with self.assertRaises(RuntimeError): + await agent.prompt("call two") - 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_unregisters_after_the_block_exits(self): + with SimpleAgent.fake(["faked"]): + self.assertTrue(Ai.has_fake_model_for("SimpleAgent")) - 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.assertFalse(Ai.has_fake_model_for("SimpleAgent")) - self.assertEqual(result.content, "case insensitive") + async def test_fake_as_decorator(self): + @SimpleAgent.fake(["decorated reply"]) + async def run(): + return await SimpleAgent().prompt("call") - 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") + result = await run() - self.assertEqual(result.content, "first match") + self.assertEqual(result.content, "decorated reply") async def test_assert_prompted_passes_after_one_call(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]) as agent: 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(["ok", "ok"]) as agent: 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"]) as agent: await agent.prompt("only once") with self.assertRaises(AssertionError): agent.assert_prompted(times=2) def test_assert_prompted_fails_when_never_called(self): - agent = SimpleAgent() - - with self.assertRaises(AssertionError): - agent.assert_prompted() + with SimpleAgent.fake([]) as agent: + with self.assertRaises(AssertionError): + agent.assert_prompted() def test_assert_prompted_times_zero_passes_when_never_called(self): - agent = SimpleAgent() - agent.assert_prompted(times=0) + with SimpleAgent.fake([]) as agent: + agent.assert_prompted(times=0) def test_assert_not_prompted_passes_when_no_calls_made(self): - agent = SimpleAgent() - agent.assert_not_prompted() + with SimpleAgent.fake([]) as agent: + agent.assert_not_prompted() async def test_assert_not_prompted_fails_after_one_call(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("a prompt") with self.assertRaises(AssertionError): agent.assert_not_prompted() async def test_reset_clears_call_log(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("first") - self.assertEqual(len(agent._call_log), 1) + self.assertEqual(len(agent._prompts), 1) - agent.reset() - self.assertEqual(agent._call_log, []) + agent.reset() + self.assertEqual(agent._prompts, []) - def test_reset_returns_agent_for_chaining(self): - agent = SimpleAgent() - result = agent.reset() - self.assertIs(result, agent) + def test_reset_returns_the_handle_for_chaining(self): + with SimpleAgent.fake([]) as agent: + self.assertIs(agent.reset(), agent) async def test_assert_not_prompted_passes_after_reset(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*": AgentResponse(content="ok")}): + with SimpleAgent.fake(["ok"]) as agent: await agent.prompt("call before reset") - - agent.reset() - agent.assert_not_prompted() + agent.reset() + agent.assert_not_prompted() async def test_fake_rebinding_overrides_previous(self): - agent = SimpleAgent() - - with SimpleAgent.fake({"*": AgentResponse(content="first fake")}): + with SimpleAgent.fake(["first fake"]) as agent: self.assertEqual((await agent.prompt("call")).content, "first fake") - with SimpleAgent.fake({"*": AgentResponse(content="second fake")}): + with SimpleAgent.fake(["second fake"]) as agent: self.assertEqual((await agent.prompt("call again")).content, "second fake") async def test_stream_returns_fake_response(self): - agent = SimpleAgent() - with SimpleAgent.fake({"*hello*": AgentResponse(content="Faked stream!")}): + with SimpleAgent.fake(["Faked stream!"]) as agent: 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) + 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"}): + async def test_stream_replays_the_registered_text_exactly(self): + with SimpleAgent.fake(["Hello there, friend"]) as agent: chunks = [chunk async for chunk in agent.stream("hi")] - self.assertEqual(chunks, ["Hello ", "there, ", "friend"]) - self.assertEqual("".join(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"]) as agent: [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) + # Streaming must record exactly one call — not one for stream + one for prompt. + agent.assert_prompted(times=1) class TestAgentRecord(unittest.IsolatedAsyncioTestCase): @@ -200,7 +163,7 @@ async def fake_run(agent_self, message, **kwargs): calls.append(message) return AgentResponse(content=content) - patcher = mock.patch.object(SimpleAgent, "_run", fake_run) + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) patcher.start() self.addCleanup(patcher.stop) return calls @@ -209,23 +172,24 @@ async def test_first_run_records_response_to_cassette(self): calls = self.setup_agent("recorded reply") with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with SimpleAgent.record(cassette): - result = await SimpleAgent().prompt("hello") + with SimpleAgent.record(cassette) as agent: + result = await agent.prompt("hello") self.assertEqual(result.content, "recorded reply") self.assertEqual(calls, ["hello"]) self.assertTrue(os.path.exists(cassette)) with open(cassette) as f: - self.assertIn("recorded reply", json.load(f).values()) + store = json.load(f) + self.assertTrue(any(v.get("content") == "recorded reply" for v in store.values())) async def test_second_run_replays_without_calling_run(self): calls = self.setup_agent("recorded reply") with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "c.json") - with SimpleAgent.record(cassette): - await SimpleAgent().prompt("hello") - with SimpleAgent.record(cassette): - replayed = await SimpleAgent().prompt("hello") + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + with SimpleAgent.record(cassette) as agent: + replayed = await agent.prompt("hello") self.assertEqual(replayed.content, "recorded reply") self.assertEqual(calls, ["hello"]) @@ -239,12 +203,12 @@ async def changed_run(s, m, **k): 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 mock.patch.object(SimpleAgent, "prompt", first_run): + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + with mock.patch.object(SimpleAgent, "prompt", changed_run): + with SimpleAgent.record(cassette) as agent: + result = await agent.prompt("hello") self.assertEqual(result.content, "from first record") @@ -252,9 +216,9 @@ async def test_distinct_messages_are_recorded_separately(self): calls = self.setup_agent("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") + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + await agent.prompt("goodbye") self.assertEqual(calls, ["hello", "goodbye"]) with open(cassette) as f: @@ -268,7 +232,7 @@ async def fake_stream(agent_self, message, **kwargs): for chunk in chunks: yield chunk - patcher = mock.patch.object(SimpleAgent, "_stream", fake_stream) + patcher = mock.patch.object(SimpleAgent, "stream", fake_stream) patcher.start() self.addCleanup(patcher.stop) return calls @@ -277,8 +241,8 @@ async def test_stream_first_run_records_chunk_list_to_cassette(self): calls = self.setup_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")] + with SimpleAgent.record(cassette) as agent: + chunks = [c async for c in agent.stream("hi")] self.assertEqual(chunks, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) @@ -289,10 +253,10 @@ 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")] - with SimpleAgent.record(cassette): - replayed = [c async for c in SimpleAgent().stream("hi")] + with SimpleAgent.record(cassette) as agent: + [c async for c in agent.stream("hi")] + with SimpleAgent.record(cassette) as agent: + replayed = [c async for c in agent.stream("hi")] self.assertEqual(replayed, ["Hel", "lo!"]) self.assertEqual(calls, ["hi"]) # real stream invoked only on the first run @@ -301,9 +265,9 @@ async def test_prompt_reads_a_stream_recorded_cassette_as_joined_content(self): 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")] - with SimpleAgent.record(cassette): - response = await SimpleAgent().prompt("hi") + with SimpleAgent.record(cassette) as agent: + [c async for c in agent.stream("hi")] + with SimpleAgent.record(cassette) as agent: + response = await agent.prompt("hi") self.assertEqual(response.content, "Hello!") diff --git a/fastapi_startkit/tests/ai/test_agent_record_fluent.py b/fastapi_startkit/tests/ai/test_agent_record_fluent.py new file mode 100644 index 00000000..1c0d08fa --- /dev/null +++ b/fastapi_startkit/tests/ai/test_agent_record_fluent.py @@ -0,0 +1,347 @@ +"""Tests for the fluent Agent.record() testing DSL. + +``with Agent.record(cassette) as agent:`` binds an ``AgentRecordFake`` handle +whose async ``prompt()`` and assertion methods judge the most recent turn — +mirroring how a browser-testing ``page`` object exposes assertions against +current page state: + + with RouterAgent.record("cassette.json") as agent: + await agent.prompt("hello") + agent.assert_text_response() + agent.assert_tool_not_called(["job_search_tool"]) + 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") +""" + +import os +import tempfile +import unittest +from unittest import mock + +from langchain_core.messages import AIMessage, HumanMessage + +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.runner import Runner +from fastapi_startkit.ai.response import AgentResponse +from fastapi_startkit.ai.testing import AgentRecordFake + + +class SimpleAgent(Agent): + pass + + +def _tool_call(name: str, args: dict | None = None, call_id: str = "c1") -> dict: + return {"name": name, "args": args or {}, "id": call_id} + + +class TestFluentPromptMechanics(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, responses: list): + """responses: list of (content, tool_calls) tuples, consumed in call order.""" + queue = list(responses) + + async def fake_run(agent_self, message, **kwargs): + content, tool_calls = queue.pop(0) + return AgentResponse(content=content, tool_calls=tool_calls or []) + + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_prompt_is_async_and_returns_agent_response(self): + self.setup_agent([("Hello there!", [])]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + response = await agent.prompt("hi") + + self.assertIsInstance(response, AgentResponse) + self.assertEqual(response.content, "Hello there!") + + async def test_second_prompt_continues_the_same_session(self): + self.setup_agent([("Hi!", []), ("here are some jobs", [_tool_call("job_search_tool")])]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + agent.assert_text_response() + + await agent.prompt("suggest python developer jobs") + agent.assert_tool_called("job_search_tool") + + async def test_replaying_from_cassette_preserves_tool_calls(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + self.setup_agent([("", [_tool_call("job_search_tool", {"q": "python"})])]) + with SimpleAgent.record(cassette) as agent: + await agent.prompt("find jobs") + + # No queued responses left — this must replay from cassette, not call _run again. + with SimpleAgent.record(cassette) as agent: + await agent.prompt("find jobs") + agent.assert_tool_called("job_search_tool") + + +class TestAssertTextResponse(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, content, tool_calls=None): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content=content, tool_calls=tool_calls or []) + + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_content_present(self): + self.setup_agent("Hi!") + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + agent.assert_text_response() + + async def test_fails_on_empty_content(self): + self.setup_agent("", tool_calls=[_tool_call("search")]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + with self.assertRaises(AssertionError): + agent.assert_text_response() + + async def test_fails_when_no_prompt_has_been_made(self): + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + with self.assertRaises(AssertionError): + agent.assert_text_response() + + +class TestAssertToolCalled(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, tool_calls): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="", tool_calls=tool_calls) + + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_tool_present(self): + self.setup_agent([_tool_call("job_search_tool", {"q": "python"})]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + agent.assert_tool_called("job_search_tool") + + async def test_fails_when_tool_absent(self): + self.setup_agent([]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + with self.assertRaises(AssertionError): + agent.assert_tool_called("job_search_tool") + + async def test_predicate_can_accept_via_attribute_access(self): + self.setup_agent([_tool_call("job_search_tool", {"q": "python"})]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + agent.assert_tool_called("job_search_tool", lambda tool: tool.name == "job_search_tool") + + async def test_predicate_can_reject(self): + self.setup_agent([_tool_call("job_search_tool", {"q": "python"})]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + with self.assertRaises(AssertionError): + agent.assert_tool_called("job_search_tool", lambda tool: tool.args.get("q") == "java") + + +class TestAssertToolNotCalled(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, tool_calls): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="Hello!", tool_calls=tool_calls) + + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_absent(self): + self.setup_agent([]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + agent.assert_tool_not_called(["job_search_tool"]) + + async def test_fails_when_present(self): + self.setup_agent([_tool_call("job_search_tool")]) + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("find jobs") + with self.assertRaises(AssertionError): + agent.assert_tool_not_called(["job_search_tool"]) + + +class TestAssertResponseTimeLt(unittest.IsolatedAsyncioTestCase): + def setup_agent(self): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="Hello!") + + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_for_a_fast_call(self): + self.setup_agent() + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + agent.assert_response_time_lt(5) + + async def test_fails_when_exceeded(self): + self.setup_agent() + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hi") + with self.assertRaises(AssertionError): + agent.assert_response_time_lt(0) + + async def test_fails_when_no_prompt_has_been_made(self): + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + with self.assertRaises(AssertionError): + agent.assert_response_time_lt(5) + + +class TestRecordMessagesSeed(unittest.IsolatedAsyncioTestCase): + async def test_seed_messages_are_included_when_building_the_real_agents_messages(self): + seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json"), messages=seed) as agent: + built = Runner(agent._real)._build_messages("suggest python developer jobs") + + self.assertEqual(built[0], seed[0]) + self.assertEqual(built[1], seed[1]) + self.assertEqual(built[-1], {"role": "user", "content": "suggest python developer jobs"}) + + async def test_same_followup_text_with_different_seed_history_does_not_collide(self): + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "shared.json") + + async def run_a(agent_self, message, **kwargs): + return AgentResponse(content="job list A") + + with mock.patch.object(SimpleAgent, "prompt", run_a): + with SimpleAgent.record(cassette) as agent: + response_a = await agent.prompt("suggest python developer jobs") + + async def run_b(agent_self, message, **kwargs): + return AgentResponse(content="job list B") + + seed = [HumanMessage(content="Hi"), AIMessage(content="Hello, how can I help?")] + with mock.patch.object(SimpleAgent, "prompt", run_b): + with SimpleAgent.record(cassette, messages=seed) as agent: + response_b = await agent.prompt("suggest python developer jobs") + + self.assertEqual(response_a.content, "job list A") + self.assertEqual(response_b.content, "job list B") + + +class TestAssertResponseJudged(unittest.IsolatedAsyncioTestCase): + def setup_agent(self, content): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content=content) + + patcher = mock.patch.object(SimpleAgent, "prompt", fake_run) + patcher.start() + self.addCleanup(patcher.stop) + + async def test_passes_when_judge_approves(self): + self.setup_agent("Hello there, welcome!") + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object( + AgentRecordFake, + "_judge_live", + mock.AsyncMock(return_value={"passed": True, "reasoning": "greets the user"}), + ): + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + await agent.assert_response_judged( + model="gpt-3.5-turbo", expectation="The llm should respond with greetings" + ) + + async def test_fails_when_judge_rejects(self): + self.setup_agent("Completely unrelated content") + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object( + AgentRecordFake, + "_judge_live", + mock.AsyncMock(return_value={"passed": False, "reasoning": "not a greeting"}), + ): + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + with self.assertRaises(AssertionError): + await agent.assert_response_judged( + model="gpt-3.5-turbo", expectation="The llm should respond with greetings" + ) + + async def test_verdict_is_cached_in_the_cassette_and_not_re_judged(self): + self.setup_agent("Hello there!") + judge = mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"}) + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with mock.patch.object(AgentRecordFake, "_judge_live", judge): + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") + await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") + + judge.assert_called_once() + + async def test_verdict_persists_to_disk_for_a_later_replay(self): + self.setup_agent("Hello there!") + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with mock.patch.object( + AgentRecordFake, "_judge_live", mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"}) + ): + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") + + judge = mock.AsyncMock(side_effect=AssertionError("must not be called on replay")) + with mock.patch.object(AgentRecordFake, "_judge_live", judge): + with SimpleAgent.record(cassette) as agent: + await agent.prompt("hello") + await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") + + judge.assert_not_called() + + async def test_fails_when_no_prompt_has_been_made(self): + with tempfile.TemporaryDirectory() as tmp: + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + with self.assertRaises(AssertionError): + await agent.assert_response_judged(model="gpt-3.5-turbo", expectation="greet") + + async def test_provider_is_forwarded_to_the_judge(self): + self.setup_agent("Hello there!") + judge = mock.AsyncMock(return_value={"passed": True, "reasoning": "ok"}) + with tempfile.TemporaryDirectory() as tmp: + with mock.patch.object(AgentRecordFake, "_judge_live", judge): + with SimpleAgent.record(os.path.join(tmp, "c.json")) as agent: + await agent.prompt("hello") + await agent.assert_response_judged(model="gpt-3.5-turbo", provider="openai", expectation="greet") + + judge.assert_called_once_with("gpt-3.5-turbo", "greet", "Hello there!", "openai") + + +class TestExistingRecordApiIsUnaffected(unittest.IsolatedAsyncioTestCase): + """The pre-existing bare-context-manager Agent.record() usage (task #327) + must keep working unchanged alongside the new fluent handle.""" + + async def test_bare_context_manager_prompt_still_works(self): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content="recorded reply") + + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "c.json") + with mock.patch.object(SimpleAgent, "prompt", fake_run): + with SimpleAgent.record(cassette) as agent: + result = await agent.prompt("hello") + + self.assertEqual(result.content, "recorded reply") diff --git a/fastapi_startkit/tests/ai/test_agent_schema.py b/fastapi_startkit/tests/ai/test_agent_schema.py index 0e7d705b..347abd7a 100644 --- a/fastapi_startkit/tests/ai/test_agent_schema.py +++ b/fastapi_startkit/tests/ai/test_agent_schema.py @@ -6,6 +6,7 @@ from pydantic import BaseModel from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.ai import Ai from fastapi_startkit.ai.response import AgentResponse @@ -20,9 +21,12 @@ def schema(self): class TestAgentSchema(unittest.IsolatedAsyncioTestCase): + def tearDown(self): + Ai.reset_fakes() + 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,14 +36,14 @@ 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") @@ -49,11 +53,11 @@ async def fake_run(self, message, **kwargs): with tempfile.TemporaryDirectory() as tmp: cassette = os.path.join(tmp, "user.json") - with mock.patch.object(UserAgent, "_run", fake_run): - with UserAgent.record(cassette): - recorded = await UserAgent().prompt("get the user") - with UserAgent.record(cassette): - replayed = await UserAgent().prompt("get the user") + with mock.patch.object(UserAgent, "prompt", fake_run): + with UserAgent.record(cassette) as agent: + recorded = await agent.prompt("get the user") + with UserAgent.record(cassette) as agent: + replayed = await agent.prompt("get the user") self.assertEqual(recorded.parsed, User(id="u-9", name="Sam")) self.assertEqual(replayed.parsed, User(id="u-9", name="Sam")) diff --git a/fastapi_startkit/tests/ai/test_ai_fake.py b/fastapi_startkit/tests/ai/test_ai_fake.py new file mode 100644 index 00000000..33a29247 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_ai_fake.py @@ -0,0 +1,142 @@ +"""Tests for Ai's fake-model registry. + +Ai.fake() swaps the chat model a given agent (by class name or instance) +resolves to for a deterministic GenericFakeChatModel that replays a fixed +list of message turns — no live LLM call, no network access. +Ai().get_model_for(agent) is what Agent._build_model() calls: it returns +the registered fake when one exists, otherwise it builds a real provider +model exactly as Ai.build() always has. +""" + +import unittest +from unittest import mock + +import langchain.chat_models as chat_models +from langchain_core.messages import AIMessage, ToolCall +from langchain_core.tools import tool + +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.ai import Ai +from fastapi_startkit.application import app + + +@tool +def search_jobs(query: str) -> str: + """Search the job board for roles matching the query.""" + return "Python Developer at Shopify" + + +class JobAssistant(Agent): + def tools(self): + return [search_jobs] + + +class SimpleAgent(Agent): + pass + + +class TestAiFakeBase(unittest.IsolatedAsyncioTestCase): + def setUp(self): + from fastapi_startkit.ai import AIConfig + + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + self.addCleanup(Ai.reset_fakes) + + +class TestAiFakeRegistration(TestAiFakeBase): + def test_fake_registers_a_model_for_an_agent_class_name(self): + Ai.fake("SimpleAgent", [AIMessage(content="hi")]) + + self.assertTrue(Ai.has_fake_model_for("SimpleAgent")) + + def test_fake_accepts_an_agent_instance_keyed_by_its_class_name(self): + Ai.fake(SimpleAgent(), [AIMessage(content="hi")]) + + self.assertTrue(Ai.has_fake_model_for(SimpleAgent())) + self.assertTrue(Ai.has_fake_model_for("SimpleAgent")) + + def test_has_fake_model_for_is_false_when_nothing_registered(self): + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) + + def test_fake_coerces_plain_strings_into_ai_messages(self): + model = Ai.fake("SimpleAgent", ["plain text reply"]) + + result = model.invoke([]) + + self.assertEqual(result.content, "plain text reply") + + def test_fake_returns_the_registered_chat_model(self): + model = Ai.fake("SimpleAgent", [AIMessage(content="hi")]) + + self.assertIs(Ai.get_fake_model_for("SimpleAgent"), model) + + def test_forget_removes_a_single_registration(self): + Ai.fake("SimpleAgent", [AIMessage(content="hi")]) + Ai.fake("JobAssistant", [AIMessage(content="hi")]) + + Ai.forget("SimpleAgent") + + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) + self.assertTrue(Ai.has_fake_model_for("JobAssistant")) + + def test_forget_is_a_no_op_when_nothing_registered(self): + Ai.forget("SimpleAgent") + + self.assertFalse(Ai.has_fake_model_for("SimpleAgent")) + + +class TestAiGetModelFor(TestAiFakeBase): + def test_returns_registered_fake_without_building_a_real_model(self): + fake_model = Ai.fake("SimpleAgent", [AIMessage(content="faked")]) + + def fail_if_called(*args, **kwargs): + raise AssertionError("init_chat_model must not be called when a fake is registered") + + patcher = mock.patch.object(chat_models, "init_chat_model", fail_if_called) + patcher.start() + self.addCleanup(patcher.stop) + + resolved = Ai().get_model_for(SimpleAgent()) + + self.assertIs(resolved, fake_model) + + def test_falls_back_to_build_when_no_fake_is_registered(self): + sentinel = object() + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: sentinel) + patcher.start() + self.addCleanup(patcher.stop) + + resolved = Ai().get_model_for(SimpleAgent()) + + self.assertIs(resolved, sentinel) + + +class TestAgentPromptUsesFakeModelEndToEnd(TestAiFakeBase): + async def test_prompt_replays_the_registered_fake_model_reply(self): + Ai.fake("SimpleAgent", [AIMessage(content="faked via ai")]) + + result = await SimpleAgent().prompt("hi there") + + self.assertEqual(result.content, "faked via ai") + + async def test_prompt_runs_a_faked_tool_call_end_to_end(self): + Ai.fake( + "JobAssistant", + [ + AIMessage( + content="", + tool_calls=[ToolCall(name="search_jobs", args={"query": "python"}, id="c1", type="tool_call")], + ) + ], + ) + + result = await JobAssistant().prompt("find me a python job") + + self.assertEqual(result.content, "Python Developer at Shopify") + + async def test_registering_a_fake_does_not_affect_other_agent_classes(self): + Ai.fake("SimpleAgent", [AIMessage(content="only for SimpleAgent")]) + + self.assertFalse(Ai.has_fake_model_for("JobAssistant")) diff --git a/fastapi_startkit/tests/ai/test_graph_agent.py b/fastapi_startkit/tests/ai/test_graph_agent.py new file mode 100644 index 00000000..b5d0d858 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_graph_agent.py @@ -0,0 +1,77 @@ +"""Tests for GraphAgent — a LangGraph tool-calling loop layered on Agent. + +Uses the same ``Agent.fake()`` model-swap as the other agent tests: scripted +turns drive the real graph (message assembly -> llm node -> tool node -> loop), +only the model at the bottom is a fake. +""" + +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 GraphAgent, GraphRunner, GraphState + + +@tool +def add(a: int, b: int) -> int: + """Add ``a`` and ``b``.""" + return a + b + + +class MathAgent(GraphAgent): + def instructions(self): + return "You are a helpful assistant tasked with performing arithmetic." + + def tools(self): + return [add] + + def graph(self, runner: GraphRunner) -> StateGraph: + graph = StateGraph(GraphState) + 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 + + +# A tool-calling turn (add 3 + 4), then the model's final natural-language answer. +SCRIPT = [ + AIMessage(content="", tool_calls=[{"name": "add", "args": {"a": 3, "b": 4}, "id": "call_0"}]), + "The answer is 7", +] + + +class TestGraphAgent(unittest.IsolatedAsyncioTestCase): + def tearDown(self): + Ai.reset_fakes() + + async def test_prompt_runs_the_tool_loop(self): + with MathAgent.fake(list(SCRIPT)): + response = await MathAgent().prompt("Add 3 and 4.") + + self.assertEqual(response.content, "The answer is 7") + self.assertTrue(any(call["name"] == "add" for call in response.tool_calls)) + + async def test_prompt_executes_the_tool_and_feeds_it_back(self): + with MathAgent.fake(list(SCRIPT)): + response = await MathAgent().prompt("Add 3 and 4.") + + # The add tool ran (3 + 4) and its result was fed back for the final turn. + add_call = next(call for call in response.tool_calls if call["name"] == "add") + self.assertEqual(add_call["args"], {"a": 3, "b": 4}) + + async def test_stream_yields_the_final_answer_token_by_token(self): + with MathAgent.fake(["The answer is 7"]): + chunks = [chunk async for chunk in MathAgent().stream("What is 3 plus 4?")] + + self.assertEqual("".join(chunks), "The answer is 7") + self.assertGreater(len(chunks), 1) + + async def test_records_the_prompt_call(self): + with MathAgent.fake(["7"]) as agent: + await agent.prompt("Add 3 and 4.") + agent.assert_prompted(times=1) diff --git a/fastapi_startkit/tests/ai/test_judge_agent.py b/fastapi_startkit/tests/ai/test_judge_agent.py new file mode 100644 index 00000000..bcb71569 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_judge_agent.py @@ -0,0 +1,80 @@ +"""Tests for JudgeAgent — grades a response against an expectation. + +It's a plain Agent whose ``schema()`` is a ``Verdict`` model, so the model's +JSON reply is turned into a typed result through the standard structured-output +path (``response.parsed``) — no hand-rolled verdict parsing. +""" + +import os +import tempfile +import unittest +from unittest import mock + +from langchain_core.messages import AIMessage + +from fastapi_startkit.ai.judge import JudgeAgent, Verdict +from fastapi_startkit.ai.runner import Runner +from fastapi_startkit.ai.response import AgentResponse + + +class TestJudgeAgent(unittest.IsolatedAsyncioTestCase): + async def test_judge_returns_a_passing_verdict_dict(self): + with JudgeAgent.fake(['{"passed": true, "reasoning": "Greets the user politely."}']): + result = await JudgeAgent().judge("The llm should respond with greetings", "Hello there!") + + self.assertEqual(result, {"passed": True, "reasoning": "Greets the user politely."}) + + async def test_judge_returns_a_failing_verdict(self): + with JudgeAgent.fake(['{"passed": false, "reasoning": "Not a greeting."}']): + result = await JudgeAgent().judge("greet", "Completely unrelated") + + self.assertFalse(result["passed"]) + + def test_schema_is_the_verdict_model(self): + self.assertIs(JudgeAgent().schema(), Verdict) + + async def test_judge_feeds_expectation_and_response_to_the_model(self): + seen = {} + + class Capturing: + def bind_tools(self, tools, **kwargs): + return self + + async def ainvoke(self, messages): + seen["messages"] = messages + return AIMessage(content='{"passed": true, "reasoning": "ok"}') + + with mock.patch.object(Runner, "_build_model", lambda self, *a, **k: Capturing()): + await JudgeAgent().judge("The llm should respond with greetings", "Hello there!") + + blob = " ".join(str(getattr(m, "content", m)) for m in seen["messages"]) + self.assertIn("The llm should respond with greetings", blob) + self.assertIn("Hello there!", blob) + + def test_model_and_provider_are_plain_agent_attributes(self): + """No custom constructor — set like any other Agent's model/provider.""" + judge = JudgeAgent() + judge.model = "gpt-4o-mini" + judge.provider = "openai" + + self.assertEqual(judge.model, "gpt-4o-mini") + self.assertEqual(judge.provider, "openai") + + def test_model_and_provider_default_to_agent_defaults(self): + judge = JudgeAgent() + + self.assertIsNone(judge.model) + self.assertIsNone(judge.provider) + + async def test_judge_is_usable_via_the_record_fluent_dsl(self): + async def fake_run(agent_self, message, **kwargs): + return AgentResponse(content='{"passed": true, "reasoning": "ok"}') + + with tempfile.TemporaryDirectory() as tmp: + cassette = os.path.join(tmp, "judge.json") + with mock.patch.object(JudgeAgent, "prompt", fake_run): + with JudgeAgent.record(cassette) as agent: + response = await agent.prompt("grade this") + + self.assertIn('"passed"', response.content) + self.assertTrue(os.path.exists(cassette)) diff --git a/fastapi_startkit/tests/ai/test_structured_output.py b/fastapi_startkit/tests/ai/test_structured_output.py new file mode 100644 index 00000000..d0957c64 --- /dev/null +++ b/fastapi_startkit/tests/ai/test_structured_output.py @@ -0,0 +1,190 @@ +"""Structured output. + +When an Agent declares a schema(), the built model is wrapped with +model.with_structured_output(schema, include_raw=True) so the provider returns +the parsed object. Tools are bound as usual and left untouched. The wrapped +model yields {"raw", "parsed", "parsing_error"}; the Runner passes that through +and _to_agent_response unwraps it into response.parsed. + +The fake/record paths bypass build(), so they keep parsing the JSON-string +content via schema() for deterministic replay. +""" + +import unittest +from unittest import mock + +import langchain.chat_models as chat_models +from langchain_core.messages import AIMessage +from langchain_core.tools import tool +from pydantic import BaseModel + +from fastapi_startkit.ai import AIConfig +from fastapi_startkit.ai.agent import Agent +from fastapi_startkit.ai.ai import Ai +from fastapi_startkit.ai.runner import Runner +from fastapi_startkit.application import app + + +class Movie(BaseModel): + title: str + year: int + + +class MovieAgent(Agent): + def schema(self): + return Movie + + +@tool +def noop(query: str) -> str: + """A no-op tool that echoes its query.""" + return query + + +class ToolMovieAgent(Agent): + def schema(self): + return Movie + + def tools(self): + return [noop] + + +def _real_tool_call(**args) -> dict: + return {"name": "noop", "args": args, "id": "1", "type": "tool_call"} + + +class _FakeModel: + """Records the build calls made against it.""" + + def __init__(self): + self.calls = [] + + def bind_tools(self, tools, **kwargs): + self.calls.append(("bind_tools", list(tools))) + return self + + def with_structured_output(self, schema, **kwargs): + self.calls.append(("with_structured_output", schema, kwargs)) + return "STRUCTURED" + + +class TestBuild(unittest.TestCase): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + + def tearDown(self): + Ai.reset_fakes() + + def _patch(self, fake): + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: fake) + patcher.start() + self.addCleanup(patcher.stop) + + def test_schema_wraps_model_with_structured_output(self): + fake = _FakeModel() + self._patch(fake) + + result = Ai().build(MovieAgent()) + + self.assertEqual(result, "STRUCTURED") + self.assertEqual(fake.calls, [("with_structured_output", Movie, {"include_raw": True})]) + + def test_tools_are_bound_then_structured_output_is_applied(self): + fake = _FakeModel() + self._patch(fake) + + result = Ai().build(ToolMovieAgent()) + + self.assertEqual(result, "STRUCTURED") + self.assertEqual( + fake.calls, + [("bind_tools", [noop]), ("with_structured_output", Movie, {"include_raw": True})], + ) + + def test_tools_only_binds_tools_without_structured_output(self): + fake = _FakeModel() + self._patch(fake) + + class ToolAgent(Agent): + def tools(self): + return [noop] + + result = Ai().build(ToolAgent()) + + self.assertIs(result, fake) + self.assertEqual(fake.calls, [("bind_tools", [noop])]) + + def test_no_schema_no_tools_returns_the_plain_model(self): + fake = _FakeModel() + self._patch(fake) + + self.assertIs(Ai().build(Agent()), fake) + self.assertEqual(fake.calls, []) + + +class TestRunner(unittest.IsolatedAsyncioTestCase): + async def test_passes_structured_output_dict_through(self): + parsed = Movie(title="Inception", year=2010) + payload = {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + + class Model: + async def ainvoke(self, messages): + return payload + + result = await Runner(MovieAgent())._invoke(Model(), ["hi"]) + + self.assertEqual(result, payload) + + async def test_runs_the_tool_when_model_calls_a_real_tool(self): + class Model: + async def ainvoke(self, messages): + return AIMessage(content="", tool_calls=[_real_tool_call(query="hello")]) + + result = await Runner(ToolMovieAgent())._invoke(Model(), ["hi"]) + + self.assertEqual(result.content, "hello") + + +class TestResponseMapping(unittest.TestCase): + def test_unwraps_include_raw_into_parsed_and_content(self): + parsed = Movie(title="Inception", year=2010) + + response = Runner(MovieAgent())._to_agent_response( + {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + ) + + self.assertIs(response.parsed, parsed) + self.assertEqual(response.content, parsed.model_dump_json()) + self.assertEqual(response.tool_calls, []) + + +class TestPromptEndToEnd(unittest.IsolatedAsyncioTestCase): + def setUp(self): + container = app() + container.bind("ai", AIConfig()) + container.make("config").set("ai", AIConfig()) + + def tearDown(self): + Ai.reset_fakes() + + async def test_prompt_populates_parsed_via_structured_output(self): + parsed = Movie(title="Inception", year=2010) + + class Structured: + async def ainvoke(self, messages): + return {"raw": AIMessage(content=""), "parsed": parsed, "parsing_error": None} + + class FakeModel: + def with_structured_output(self, schema, **kwargs): + return Structured() + + patcher = mock.patch.object(chat_models, "init_chat_model", lambda *a, **k: FakeModel()) + patcher.start() + self.addCleanup(patcher.stop) + + response = await MovieAgent().prompt("best nolan movie") + + self.assertEqual(response.parsed, parsed) + self.assertEqual(response.content, parsed.model_dump_json()) diff --git a/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py b/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py index 675cf79a..f86fc1ef 100644 --- a/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py +++ b/fastapi_startkit/tests/masoniteorm/sqlite/relationships/test_sqlite_polymorphic.py @@ -9,10 +9,10 @@ class TestRelationships(TestCase): async def test_can_get_polymorphic_relation(self): likes = await Like.get() for like in likes: - record = await like.record + record = await like.log assert isinstance(record, (Articles, Product)) async def test_can_get_eager_load_polymorphic_relation(self): likes = await Like.with_("record").get() for like in likes: - assert isinstance(like.record, (Articles, Product)) + assert isinstance(like.log, (Articles, Product)) diff --git a/fastapi_startkit/uv.lock b/fastapi_startkit/uv.lock index 5808c35b..032c1a36 100644 --- a/fastapi_startkit/uv.lock +++ b/fastapi_startkit/uv.lock @@ -527,7 +527,7 @@ wheels = [ [[package]] name = "fastapi-startkit" -version = "0.48.0" +version = "0.50.0" source = { editable = "." } dependencies = [ { name = "cleo" }, @@ -543,6 +543,7 @@ dependencies = [ ai = [ { name = "langchain" }, { name = "langchain-core" }, + { name = "langgraph" }, ] database = [ { name = "faker" }, @@ -605,6 +606,7 @@ requires-dist = [ { name = "jinja2", marker = "extra == 'vite'", specifier = ">=3.1" }, { name = "langchain", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "langchain-core", marker = "extra == 'ai'", specifier = ">=1.0.0" }, + { name = "langgraph", marker = "extra == 'ai'", specifier = ">=1.0.0" }, { name = "markupsafe", marker = "extra == 'inertia'", specifier = ">=2.0" }, { name = "pendulum", specifier = ">=3.1.0,<4.0.0" }, { name = "pydantic", specifier = ">=2.12.5" },