diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index 4865caba..5ada3b04 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -1,3 +1,4 @@ +import json import logging import os from pathlib import Path @@ -5,6 +6,11 @@ from pathlib import Path from agent import AgentDeps, agent from anyio import create_memory_object_stream, create_task_group from anyio.streams.memory import MemoryObjectSendStream +from pydantic_ai import ( + AgentRunResultEvent, + FunctionToolCallEvent, + FunctionToolResultEvent, +) from starlette.applications import Starlette from starlette.middleware import Middleware from starlette.middleware.cors import CORSMiddleware @@ -111,12 +117,40 @@ async def stream_research_agent(request: Request) -> StreamingResponse: continue await send_stream.send(format_sse_event(event)) - # Run agent and event forwarding concurrently + # Run agent with streaming and capture tool events async with create_task_group() as tg: tg.start_soon(forward_events) - result = await agent.run(user_message, deps=agent_deps) - emitter.log(result.output) + # Use run_stream_events to capture all events including tool calls + async for event in agent.run_stream_events( + user_message, deps=agent_deps + ): + # Emit tool call events to AG-UI + if isinstance(event, FunctionToolCallEvent): + # Tool call started + emitter.tool_call_start( + tool_call_id=event.part.tool_call_id, + tool_name=event.part.tool_name, + ) + # Emit args as single delta (they're already complete) + emitter.tool_call_args( + tool_call_id=event.part.tool_call_id, + args_delta=json.dumps(event.part.args), + ) + # End the args stream + emitter.tool_call_end(tool_call_id=event.part.tool_call_id) + + elif isinstance(event, FunctionToolResultEvent): + # Tool call completed with result + emitter.tool_call_result( + tool_call_id=event.tool_call_id, + result=str(event.result.content), + ) + + elif isinstance(event, AgentRunResultEvent): + # Final result from agent + emitter.log(event.result.output) + await emitter.close() except Exception as e: diff --git a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py index 09201559..f1294cb4 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py @@ -18,6 +18,10 @@ from haiku.rag.graph.agui.events import ( emit_step_finished, emit_step_started, emit_text_message, + emit_tool_call_args, + emit_tool_call_end, + emit_tool_call_result, + emit_tool_call_start, ) @@ -154,6 +158,44 @@ class AGUIEmitter[StateT: BaseModel, ResultT]: """ self._emit(emit_run_error(str(error), code)) + def tool_call_start( + self, tool_call_id: str, tool_name: str, parent_message_id: str | None = None + ) -> None: + """Emit ToolCallStart event. + + Args: + tool_call_id: Unique identifier for this tool call + tool_name: Name of the tool being called + parent_message_id: Optional parent message ID + """ + self._emit(emit_tool_call_start(tool_call_id, tool_name, parent_message_id)) + + def tool_call_args(self, tool_call_id: str, args_delta: str) -> None: + """Emit ToolCallArgs event. + + Args: + tool_call_id: Identifier for the tool call + args_delta: Incremental JSON chunk of arguments + """ + self._emit(emit_tool_call_args(tool_call_id, args_delta)) + + def tool_call_end(self, tool_call_id: str) -> None: + """Emit ToolCallEnd event. + + Args: + tool_call_id: Identifier for the tool call + """ + self._emit(emit_tool_call_end(tool_call_id)) + + def tool_call_result(self, tool_call_id: str, result: str) -> None: + """Emit ToolCallResult event. + + Args: + tool_call_id: Identifier for the tool call + result: The result from the tool execution + """ + self._emit(emit_tool_call_result(tool_call_id, result)) + def _emit(self, event: AGUIEvent) -> None: """Put event in queue. diff --git a/haiku_rag_slim/haiku/rag/graph/agui/events.py b/haiku_rag_slim/haiku/rag/graph/agui/events.py index a7be3d73..a5dcd12b 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/events.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/events.py @@ -123,7 +123,7 @@ def emit_text_message(content: str, role: str = "assistant") -> dict[str, Any]: "type": "TEXT_MESSAGE_CHUNK", "messageId": message_id, "role": role, - "delta": content, + "content": content, # Changed from "delta" to "content" for CopilotKit compatibility } @@ -252,3 +252,79 @@ def emit_activity_delta( "activityType": activity_type, "patch": patch, } + + +def emit_tool_call_start( + tool_call_id: str, tool_name: str, parent_message_id: str | None = None +) -> dict[str, Any]: + """Create a ToolCallStart event. + + Args: + tool_call_id: Unique identifier for this tool call + tool_name: Name of the tool being called + parent_message_id: Optional parent message ID + + Returns: + ToolCallStart event dict + """ + event: dict[str, Any] = { + "type": "TOOL_CALL_START", + "toolCallId": tool_call_id, + "toolCallName": tool_name, + } + if parent_message_id: + event["parentMessageId"] = parent_message_id + return event + + +def emit_tool_call_args(tool_call_id: str, args_delta: str) -> dict[str, Any]: + """Create a ToolCallArgs event. + + Args: + tool_call_id: Identifier for the tool call + args_delta: Incremental JSON chunk of arguments + + Returns: + ToolCallArgs event dict + """ + return { + "type": "TOOL_CALL_ARGS", + "toolCallId": tool_call_id, + "delta": args_delta, + } + + +def emit_tool_call_end(tool_call_id: str) -> dict[str, Any]: + """Create a ToolCallEnd event. + + Args: + tool_call_id: Identifier for the tool call + + Returns: + ToolCallEnd event dict + """ + return { + "type": "TOOL_CALL_END", + "toolCallId": tool_call_id, + } + + +def emit_tool_call_result(tool_call_id: str, result: Any) -> dict[str, Any]: + """Create a ToolCallResult event. + + Args: + tool_call_id: Identifier for the tool call + result: The result/output from the tool execution + + Returns: + ToolCallResult event dict + """ + # Convert result to string if needed + if not isinstance(result, str): + result = str(result) + + return { + "type": "TOOL_CALL_RESULT", + "toolCallId": tool_call_id, + "result": result, + }