Handle client-side tool calls in graph ag-ui emitter

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:

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.
 

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,
+    }