From 2dca72d670d502a7acb4d28505d6b07c9f71c227 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Nov 2025 15:19:28 +0200 Subject: [PATCH] Revert "Handle client-side tool calls in graph ag-ui emitter" This reverts commit b1d19651d59e53e4748aa7046f6de122f5bfd9c8. --- examples/ag-ui-research/backend/main.py | 40 +--------- .../haiku/rag/graph/agui/emitter.py | 42 ---------- haiku_rag_slim/haiku/rag/graph/agui/events.py | 78 +------------------ 3 files changed, 4 insertions(+), 156 deletions(-) diff --git a/examples/ag-ui-research/backend/main.py b/examples/ag-ui-research/backend/main.py index 5ada3b04..4865caba 100644 --- a/examples/ag-ui-research/backend/main.py +++ b/examples/ag-ui-research/backend/main.py @@ -1,4 +1,3 @@ -import json import logging import os from pathlib import Path @@ -6,11 +5,6 @@ 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 @@ -117,40 +111,12 @@ async def stream_research_agent(request: Request) -> StreamingResponse: continue await send_stream.send(format_sse_event(event)) - # Run agent with streaming and capture tool events + # Run agent and event forwarding concurrently async with create_task_group() as tg: tg.start_soon(forward_events) - # 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) - + result = await agent.run(user_message, deps=agent_deps) + emitter.log(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 f1294cb4..09201559 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/emitter.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/emitter.py @@ -18,10 +18,6 @@ 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, ) @@ -158,44 +154,6 @@ 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 a5dcd12b..a7be3d73 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, - "content": content, # Changed from "delta" to "content" for CopilotKit compatibility + "delta": content, } @@ -252,79 +252,3 @@ 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, - }