""" Claude Agent SDK (Python) -- sales assistant with weather, HITL, and generative UI. Implements the AG-UI protocol directly using the Anthropic Python SDK. All demo routes share this single agent instance served by agent_server.py. """ from __future__ import annotations import json import os import traceback from collections.abc import AsyncIterator from textwrap import dedent from typing import Any import anthropic from ag_ui.core import ( EventType, Message, RunAgentInput, RunFinishedEvent, RunStartedEvent, StateSnapshotEvent, TextMessageContentEvent, TextMessageEndEvent, TextMessageStartEvent, ToolCallArgsEvent, ToolCallEndEvent, ToolCallResultEvent, ToolCallStartEvent, ) from ag_ui.encoder import EventEncoder from dotenv import load_dotenv from fastapi import FastAPI, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from pydantic import BaseModel from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import JSONResponse # Serve /health via middleware so it short-circuits BEFORE route resolution. # Any later catch-all mount at "/" (whether added here or by a downstream # adapter) would shadow a plain `@app.get("/health")` decorator. Middleware # runs above routing so the health endpoint stays reachable regardless. class HealthMiddleware(BaseHTTPMiddleware): async def dispatch(self, request, call_next): if request.url.path == "/health" and request.method == "GET": return JSONResponse({"status": "ok"}) return await call_next(request) load_dotenv() # Import shared tool implementations (via tools symlink -> ../../shared/python/tools) from tools import ( get_weather_impl, query_data_impl, manage_sales_todos_impl, get_sales_todos_impl, schedule_meeting_impl, search_flights_impl, build_a2ui_operations_from_tool_call, RENDER_A2UI_TOOL_SCHEMA, ) from tools.types import Flight # ============ # Tool schemas # ============ TOOLS: list[dict[str, Any]] = [ { "name": "get_weather", "description": ( "Get current weather for a location. " "Use this to render the frontend weather card." ), "input_schema": { "type": "object", "properties": { "location": { "type": "string", "description": "The city or region to get weather for.", }, }, "required": ["location"], }, }, { "name": "query_data", "description": ( "Query the financial database for chart data. " "Always call before showing a chart or graph." ), "input_schema": { "type": "object", "properties": { "query": { "type": "string", "description": "Natural language query for financial data.", }, }, "required": ["query"], }, }, { "name": "manage_sales_todos", "description": ( "Replace the entire list of sales todos with the provided values. " "Always include every todo you want to keep." ), "input_schema": { "type": "object", "properties": { "todos": { "type": "array", "items": { "type": "object", "properties": { "id": {"type": "string"}, "title": {"type": "string"}, "stage": { "type": "string", "enum": [ "prospect", "qualified", "proposal", "negotiation", "closed-won", "closed-lost", ], }, "value": {"type": "number"}, "dueDate": {"type": "string"}, "assignee": {"type": "string"}, "completed": {"type": "boolean"}, }, "required": [ "title", "stage", "value", "dueDate", "assignee", "completed", ], }, "description": "The complete list of sales todos.", }, }, "required": ["todos"], }, }, { "name": "get_sales_todos", "description": "Get the current sales pipeline todos.", "input_schema": { "type": "object", "properties": {}, }, }, { "name": "schedule_meeting", "description": ( "Schedule a meeting with the user. Requires human approval. " "Call this when the user wants to schedule or book a meeting." ), "input_schema": { "type": "object", "properties": { "reason": { "type": "string", "description": "Reason for the meeting.", }, }, "required": ["reason"], }, }, { "name": "generate_task_steps", "description": ( "Propose a list of steps for the user to review and approve. " "Used for human-in-the-loop task planning. " "Always call this tool when the user asks you to plan something." ), "input_schema": { "type": "object", "properties": { "steps": { "type": "array", "items": { "type": "object", "properties": { "description": {"type": "string"}, "status": { "type": "string", "enum": ["enabled", "disabled", "executing"], }, }, "required": ["description", "status"], }, "description": "The ordered list of steps for the user to review.", } }, "required": ["steps"], }, }, { "name": "change_background", "description": ( "Change the background color or gradient of the chat UI. " "ONLY call this when the user explicitly asks to change the background." ), "input_schema": { "type": "object", "properties": { "background": { "type": "string", "description": "CSS background value. Prefer gradients.", } }, "required": ["background"], }, }, { "name": "search_flights", "description": ( "Search for flights and display the results as rich A2UI cards. " "Return exactly 2 flights. Each flight must have: airline, airlineLogo, " "flightNumber, origin, destination, date, departureTime, arrivalTime, " "duration, status, statusColor, price, currency. " "For airlineLogo use: https://www.google.com/s2/favicons?domain={airline_domain}&sz=128" ), "input_schema": { "type": "object", "properties": { "flights": { "type": "array", "items": { "type": "object", "properties": { "airline": {"type": "string"}, "airlineLogo": {"type": "string"}, "flightNumber": {"type": "string"}, "origin": {"type": "string"}, "destination": {"type": "string"}, "date": {"type": "string"}, "departureTime": {"type": "string"}, "arrivalTime": {"type": "string"}, "duration": {"type": "string"}, "status": {"type": "string"}, "statusColor": {"type": "string"}, "price": {"type": "string"}, "currency": {"type": "string"}, }, }, "description": "List of flight objects to display.", }, }, "required": ["flights"], }, }, { "name": "generate_a2ui", "description": ( "Generate dynamic A2UI components based on the conversation. " "A secondary LLM designs the UI schema and data." ), "input_schema": { "type": "object", "properties": { "context": { "type": "string", "description": "Conversation context to generate UI for.", }, }, "required": ["context"], }, }, ] SYSTEM_PROMPT = dedent(""" You are a helpful sales assistant that manages a sales pipeline, discusses weather, queries financial data, schedules meetings, and helps with planning. Sales pipeline management: - The current list of sales todos is provided in the conversation context. - When you add, remove, or update todos, call `manage_sales_todos` with the FULL list. - CRITICAL: When asked to "add" a todo, include ALL existing todos + the new one. - When asked to "remove" a todo, include everything EXCEPT the removed one. Tool usage: - `get_weather`: only call when the user explicitly asks about weather. - `query_data`: call when the user asks about financial data, charts, or graphs. - `manage_sales_todos`: call to update the sales pipeline. - `get_sales_todos`: call to retrieve current sales pipeline. - `schedule_meeting`: call when the user wants to schedule a meeting. - `generate_task_steps`: call when the user asks you to plan something step-by-step. Wait for approval/rejection before continuing with the plan. - `change_background`: only call when user explicitly asks to change the background. - `search_flights`: call when the user asks about flights. Generate 2 realistic flights. - `generate_a2ui`: call when the user asks for a dashboard or dynamic UI. After executing tools, provide a brief summary of what changed. Keep responses concise and friendly. """).strip() # =========== # AG-UI runner # =========== class AgentState(BaseModel): todos: list[dict] = [] def _execute_tool( name: str, tool_input: dict[str, Any], state: AgentState, conversation_messages: list[dict[str, Any]] | None = None, ) -> tuple[str, AgentState | None]: """Execute backend tools and return (result_text, new_state_or_None).""" if name == "get_weather": return json.dumps(get_weather_impl(tool_input["location"])), None if name == "query_data": return json.dumps(query_data_impl(tool_input["query"])), None if name == "manage_sales_todos": result = manage_sales_todos_impl(tool_input["todos"]) state.todos = [dict(t) for t in result] return json.dumps({"status": "updated", "count": len(result)}), state if name == "get_sales_todos": return json.dumps( get_sales_todos_impl(state.todos if state.todos else None) ), None if name == "schedule_meeting": return json.dumps(schedule_meeting_impl(tool_input["reason"])), None if name == "generate_task_steps": # Frontend HITL tool -- backend just acknowledges; UI handles the interaction steps = tool_input.get("steps", []) return f"Presented {len(steps)} steps for review.", None if name == "change_background": # Frontend tool -- backend just acknowledges return f"Background change requested: {tool_input.get('background', '')}", None if name == "search_flights": flights_data = tool_input.get("flights", []) typed_flights = [Flight(**f) for f in flights_data] result = search_flights_impl(typed_flights) return json.dumps(result), None if name == "generate_a2ui": context = tool_input.get("context", "") import openai client = openai.OpenAI() llm_messages: list[dict[str, Any]] = [ {"role": "system", "content": context or "Generate a useful dashboard UI."}, ] # Pass conversation messages to the secondary LLM for context if conversation_messages: llm_messages.extend(conversation_messages) else: llm_messages.append( { "role": "user", "content": "Generate a dynamic A2UI dashboard based on the conversation.", } ) response = client.chat.completions.create( model="gpt-4.1", messages=llm_messages, tools=[{"type": "function", "function": RENDER_A2UI_TOOL_SCHEMA}], tool_choice={"type": "function", "function": {"name": "render_a2ui"}}, ) choice = response.choices[0] if choice.message.tool_calls: args = json.loads(choice.message.tool_calls[0].function.arguments) a2ui_result = build_a2ui_operations_from_tool_call(args) return json.dumps(a2ui_result), None return json.dumps({"error": "LLM did not call render_a2ui"}), None return f"Unknown tool: {name}", None def _build_frontend_tools(input_data: RunAgentInput) -> list[dict[str, Any]]: """Extract frontend-defined tools from the AG-UI request. The CopilotKit runtime forwards frontend tool definitions (registered via ``useFrontendTool``, ``useHumanInTheLoop``, etc.) in ``input_data.tools``. We convert them to the Anthropic ``tools`` schema so the LLM can call them. The runtime intercepts the resulting tool-call events and routes them to the frontend for resolution. """ out: list[dict[str, Any]] = [] for t in input_data.tools or []: name = getattr(t, "name", None) or ( t.get("name") if isinstance(t, dict) else None ) description = getattr(t, "description", None) or ( t.get("description", "") if isinstance(t, dict) else "" ) parameters = getattr(t, "parameters", None) or ( t.get("parameters", {}) if isinstance(t, dict) else {} ) if not name: continue out.append( { "name": name, "description": description or "", "input_schema": parameters or {"type": "object", "properties": {}}, } ) return out async def run_agent( input_data: RunAgentInput, *, system_prompt_override: str | None = None, disable_tools: bool = False, preprocess_user_parts: Any = None, ) -> AsyncIterator[str]: """Run the Claude agent and yield AG-UI SSE events. Keyword arguments let dedicated demo endpoints reuse this streaming loop with targeted overrides: - ``system_prompt_override`` — replace the shared ``SYSTEM_PROMPT`` (e.g. BYOC demos emit a JSON envelope, so the sales-assistant prompt is irrelevant). - ``disable_tools`` — run the model with no tool schemas. Useful for BYOC / pure-text demos where tool calls would derail the output. - ``preprocess_user_parts`` — a ``callable(part) -> part`` applied to each content part of every user message before they are sent to Claude. Used by the multimodal demo to convert AG-UI ``image``/``document`` parts into Claude's Messages API shape (``{"type": "image", "source": {...}}``) and to flatten PDFs to text via ``pypdf``. """ encoder = EventEncoder() client = anthropic.AsyncAnthropic(api_key=os.getenv("ANTHROPIC_API_KEY", "")) # Extract state state = AgentState() if input_data.state and isinstance(input_data.state, dict): state = AgentState(**input_data.state) # Convert AG-UI messages to Anthropic format. When a preprocessor is # supplied we preserve the structured content list (image blocks, # document text, etc.) — otherwise we collapse to a flat string for # the text-only happy path used by most demos. # # AG-UI delivers three message roles: # - "user" → plain user text # - "assistant" → assistant text + optional tool_use blocks # - "tool" → tool result from a resolved frontend tool # # Anthropic's Messages API represents tool results as a "user" role # message with content blocks of type "tool_result". We must convert # AG-UI "tool" messages into that shape so the LLM sees the resolved # result and aimock's ``hasToolResult`` matcher fires correctly. messages: list[dict[str, Any]] = [] for msg in input_data.messages or []: role = msg.role.value if hasattr(msg.role, "value") else str(msg.role) # Handle tool result messages from AG-UI (resolved frontend tools). # Convert to Anthropic's format: role="user" with tool_result blocks. if role == "tool": tool_call_id = getattr(msg, "tool_call_id", None) or ( getattr(msg, "toolCallId", None) ) raw_content = getattr(msg, "content", None) result_text = "" if isinstance(raw_content, str): result_text = raw_content elif isinstance(raw_content, list): parts = [] for part in raw_content: if hasattr(part, "text"): parts.append(part.text) elif isinstance(part, dict) and "text" in part: parts.append(part["text"]) parts_text = "".join(parts) if parts_text: result_text = parts_text else: result_text = json.dumps(raw_content) if tool_call_id: # Anthropic expects the assistant message containing the # tool_use to precede this tool_result message. The runtime # ensures message ordering, so we just need to emit the # tool_result in the right shape. messages.append( { "role": "user", "content": [ { "type": "tool_result", "tool_use_id": tool_call_id, "content": result_text, } ], } ) continue if role not in ("user", "assistant"): continue raw_content = getattr(msg, "content", None) if ( preprocess_user_parts is not None and role == "user" and isinstance(raw_content, list) ): converted_parts: list[Any] = [] for part in raw_content: # AG-UI emits pydantic models; normalise to a plain dict # before handing to the converter so the demo-specific # code can rely on ``.get(...)`` semantics. if hasattr(part, "model_dump"): part_dict = part.model_dump() elif isinstance(part, dict): part_dict = part else: part_dict = part converted = preprocess_user_parts(part_dict) if converted is not None: converted_parts.append(converted) if converted_parts: messages.append({"role": role, "content": converted_parts}) continue # For assistant messages, check if there are tool calls (AG-UI's # AssistantMessage stores them in `tool_calls`, not in `content`). # Anthropic requires tool_use blocks in the assistant content so # the subsequent tool_result can pair with them. if role == "assistant": msg_tool_calls = getattr(msg, "tool_calls", None) text_content = "" if isinstance(raw_content, str): text_content = raw_content elif isinstance(raw_content, list): for part in raw_content: if hasattr(part, "text"): text_content += part.text elif isinstance(part, dict) and "text" in part: text_content += part["text"] if msg_tool_calls: content_blocks: list[dict[str, Any]] = [] if text_content: content_blocks.append({"type": "text", "text": text_content}) for tc in msg_tool_calls: # AG-UI ToolCall: {id, function: {name, arguments}} tc_id = getattr(tc, "id", None) or ( tc.get("id") if isinstance(tc, dict) else None ) func = getattr(tc, "function", None) or ( tc.get("function") if isinstance(tc, dict) else None ) if func: tc_name = getattr(func, "name", None) or ( func.get("name") if isinstance(func, dict) else "unknown" ) tc_args_str = getattr(func, "arguments", None) or ( func.get("arguments", "{}") if isinstance(func, dict) else "{}" ) else: tc_name = "unknown" tc_args_str = "{}" try: tc_args = ( json.loads(tc_args_str) if isinstance(tc_args_str, str) else tc_args_str ) except json.JSONDecodeError: tc_args = {} content_blocks.append( { "type": "tool_use", "id": tc_id or "unknown", "name": tc_name, "input": tc_args, } ) messages.append({"role": "assistant", "content": content_blocks}) continue elif text_content: messages.append({"role": "assistant", "content": text_content}) continue # Fall through to the generic handler if nothing matched content = "" if isinstance(raw_content, str): content = raw_content elif isinstance(raw_content, list): parts = [] for part in raw_content: if hasattr(part, "text"): parts.append(part.text) elif isinstance(part, dict) and "text" in part: parts.append(part["text"]) content = "".join(parts) if content: messages.append({"role": role, "content": content}) # Inject sales pipeline state into system prompt if state exists if system_prompt_override is not None: system = system_prompt_override else: system = SYSTEM_PROMPT if state.todos: todos_json = json.dumps(state.todos, indent=2) system = f"{SYSTEM_PROMPT}\n\nCurrent sales pipeline:\n{todos_json}" thread_id = input_data.thread_id or "default" run_id = input_data.run_id or "run-1" yield encoder.encode( RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=run_id) ) # Agentic loop -- keep calling Claude until no more tool calls while True: response_text = "" tool_calls: list[dict[str, Any]] = [] msg_id = f"msg-{run_id}-{len(messages)}" yield encoder.encode( TextMessageStartEvent( type=EventType.TEXT_MESSAGE_START, message_id=msg_id, role="assistant", ) ) # Build the combined tools list: backend TOOLS + any frontend- # defined tools forwarded by the CopilotKit runtime in # input_data.tools. Frontend tools (registered via useFrontendTool, # useHumanInTheLoop, etc.) are included so the LLM can call them; # the runtime intercepts the resulting events and routes them to # the frontend for resolution. Backend tools are executed locally. backend_tool_names = {t["name"] for t in TOOLS} frontend_tools = _build_frontend_tools(input_data) # Merge: backend tools first, then frontend tools that don't # shadow a backend tool (frontend wins when names collide, because # the frontend registration means the runtime should intercept). frontend_tool_names = {t["name"] for t in frontend_tools} combined_tools: list[dict[str, Any]] = [] for t in TOOLS: if t["name"] not in frontend_tool_names: combined_tools.append(t) combined_tools.extend(frontend_tools) stream_kwargs: dict[str, Any] = { "model": os.getenv("ANTHROPIC_MODEL", "claude-opus-4-5"), "max_tokens": 4096, "system": system, "messages": messages, } if not disable_tools: stream_kwargs["tools"] = combined_tools # type: ignore[assignment] try: async with client.messages.stream(**stream_kwargs) as stream: current_tool_id: str | None = None current_tool_name: str | None = None current_tool_args = "" async for event in stream: etype = type(event).__name__ if etype == "RawContentBlockStartEvent": block = event.content_block # type: ignore[attr-defined] if block.type == "text": pass # streaming text chunks follow elif block.type == "tool_use": current_tool_id = block.id current_tool_name = block.name current_tool_args = "" yield encoder.encode( ToolCallStartEvent( type=EventType.TOOL_CALL_START, tool_call_id=current_tool_id, tool_call_name=current_tool_name, parent_message_id=msg_id, ) ) elif etype == "RawContentBlockDeltaEvent": delta = event.delta # type: ignore[attr-defined] if delta.type == "text_delta": response_text += delta.text yield encoder.encode( TextMessageContentEvent( type=EventType.TEXT_MESSAGE_CONTENT, message_id=msg_id, delta=delta.text, ) ) elif delta.type == "input_json_delta": current_tool_args += delta.partial_json yield encoder.encode( ToolCallArgsEvent( type=EventType.TOOL_CALL_ARGS, tool_call_id=current_tool_id or "", delta=delta.partial_json, ) ) elif etype in ( "RawContentBlockStopEvent", "ParsedContentBlockStopEvent", ): if current_tool_id and current_tool_name: yield encoder.encode( ToolCallEndEvent( type=EventType.TOOL_CALL_END, tool_call_id=current_tool_id, ) ) try: parsed_args = ( json.loads(current_tool_args) if current_tool_args else {} ) except json.JSONDecodeError: parsed_args = {} tool_calls.append( { "id": current_tool_id, "name": current_tool_name, "input": parsed_args, } ) current_tool_id = None current_tool_name = None current_tool_args = "" except Exception: # Surface the error as visible text in the chat so D5 # probes see a non-empty assistant response instead of a # silent broken SSE stream. Full traceback is logged # server-side by FastAPI's exception handler. err_text = f"Agent error: {traceback.format_exc()}" yield encoder.encode( TextMessageContentEvent( type=EventType.TEXT_MESSAGE_CONTENT, message_id=msg_id, delta=err_text, ) ) yield encoder.encode( TextMessageEndEvent( type=EventType.TEXT_MESSAGE_END, message_id=msg_id, ) ) # No tool calls -- we're done if not tool_calls: break # Separate tool calls into backend (locally executed) and frontend # (deferred to the CopilotKit runtime / frontend for resolution). # A tool whose name was registered on the frontend (present in # frontend_tool_names) is a frontend tool even if the backend also # defines it — the frontend registration takes precedence because # hooks like useHumanInTheLoop rely on intercepting the tool call. has_frontend_tool = any(tc["name"] in frontend_tool_names for tc in tool_calls) if has_frontend_tool: # At least one tool call targets a frontend tool. Break the # agentic loop: the CopilotKit runtime will intercept the # pending frontend tool call(s), route them to the frontend # for user interaction, and re-invoke the agent with the # resolved tool result(s) in a subsequent request. # # We do NOT emit ToolCallResultEvent for frontend tools and # we do NOT add them to the message history — the runtime # owns the continuation from here. break # All tool calls are backend-only — execute locally and continue # the agentic loop. # Add assistant turn with tool calls to message history assistant_content: list[dict[str, Any]] = [] if response_text: assistant_content.append({"type": "text", "text": response_text}) for tc in tool_calls: assistant_content.append( { "type": "tool_use", "id": tc["id"], "name": tc["name"], "input": tc["input"], } ) messages.append({"role": "assistant", "content": assistant_content}) # Execute tools and build tool-result turn tool_results: list[dict[str, Any]] = [] for tc in tool_calls: result_text, new_state = _execute_tool( tc["name"], tc["input"], state, conversation_messages=messages ) if new_state is not None: state = new_state yield encoder.encode( StateSnapshotEvent( type=EventType.STATE_SNAPSHOT, snapshot=state.model_dump(), ) ) yield encoder.encode( ToolCallResultEvent( type=EventType.TOOL_CALL_RESULT, tool_call_id=tc["id"], message_id=f"{msg_id}-tool-result-{tc['id']}", content=result_text, ) ) tool_results.append( { "type": "tool_result", "tool_use_id": tc["id"], "content": result_text, } ) messages.append({"role": "user", "content": tool_results}) yield encoder.encode( RunFinishedEvent( type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=run_id ) ) def create_app() -> FastAPI: """Create the FastAPI app with AG-UI endpoint.""" # Local import to avoid a top-level ``agents._header_forwarding`` # dependency in this module (kept agnostic so unit tests that import # individual handlers don't need the starlette middleware shape). from agents._header_forwarding import HeaderForwardingHTTPMiddleware app = FastAPI(title="Claude Agent SDK (Python) Agent Server") app.add_middleware(HealthMiddleware) # Capture inbound CopilotKit ``x-*`` headers (e.g. ``x-aimock-context``) # into a per-request ContextVar so any outbound LLM/provider httpx call # made inside the request scope copies them onto its outbound request. # Paired with ``install_global_httpx_hook`` at the top of agent_server.py. app.add_middleware(HeaderForwardingHTTPMiddleware) app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_methods=["*"], allow_headers=["*"], ) @app.post("/") async def run_agent_endpoint(request: Request) -> StreamingResponse: body = await request.json() input_data = RunAgentInput(**body) async def event_stream() -> AsyncIterator[str]: async for chunk in run_agent(input_data): yield chunk return StreamingResponse( event_stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "X-Accel-Buffering": "no", }, ) return app