diff --git a/web/agent/src/agent.py b/web/agent/src/agent.py index 0d39771..93cbd83 100644 --- a/web/agent/src/agent.py +++ b/web/agent/src/agent.py @@ -26,8 +26,10 @@ logfire.configure( logfire.instrument_pydantic_ai() logfire.instrument_httpx() +from typing import Any from pydantic_ai import Agent, ModelMessage, RunContext from pydantic_ai.settings import ModelSettings +from pydantic_ai.mcp import CallToolFunc CAVE_MCP_URL = os.getenv("CAVE_MCP_URL", "https://mcp.caving.dev/mcp") @@ -85,6 +87,29 @@ Rules: 7. Use tools sparingly—one search usually suffices.""" +def create_tool_call_limiter(max_calls: int = 3): + """Create a process_tool_call callback that limits tool calls.""" + call_count = [0] # Mutable container for closure + + async def process_tool_call( + ctx: RunContext, + call_tool: CallToolFunc, + name: str, + tool_args: dict[str, Any], + ): + call_count[0] += 1 + if call_count[0] > max_calls: + return ( + f"TOOL LIMIT REACHED: You have made {max_calls} tool calls. " + "Stop searching and provide your answer now using the information " + "you already have. If you don't have enough information, tell the " + "user what you found and ask them to rephrase their question." + ) + return await call_tool(name, tool_args) + + return process_tool_call + + def create_agent(user_roles: list[str] | None = None): """Create an agent with MCP tools configured for the given user roles.""" toolsets = [] @@ -104,6 +129,7 @@ def create_agent(user_roles: list[str] | None = None): url=CAVE_MCP_URL, headers={"x-user-roles": roles_header}, timeout=30.0, + process_tool_call=create_tool_call_limiter(max_calls=3), ) toolsets.append(mcp_server) logger.info(f"MCP server configured with roles: {user_roles}") diff --git a/web/agent/src/main.py b/web/agent/src/main.py index 3edbe44..e7e3d3b 100644 --- a/web/agent/src/main.py +++ b/web/agent/src/main.py @@ -70,13 +70,12 @@ async def handle_agent_request(request: Request) -> Response: # Create agent with the user's roles agent = create_agent(user_roles) - # Dispatch the request using AGUIAdapter with usage limits + # Dispatch the request - tool limits handled by ToolCallLimiter in agent.py return await AGUIAdapter.dispatch_request( request, agent=agent, usage_limits=UsageLimits( - request_limit=5, # Max 5 LLM requests per query - tool_calls_limit=3, # Max 3 tool calls per query + request_limit=10, # Safety net for runaway requests ), model_settings=ModelSettings(max_tokens=4096), )