From b1d19651d59e53e4748aa7046f6de122f5bfd9c8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Nov 2025 10:50:16 +0200 Subject: [PATCH 1/2] Handle client-side tool calls in graph ag-ui emitter --- 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, 156 insertions(+), 4 deletions(-) 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, + } From cc09ce87c2c8ec6ac7e6429f7e0197eb12fc39c0 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 18 Nov 2025 14:35:57 +0200 Subject: [PATCH 2/2] Increase reranking candidate retrieval multiplier from 3x to 10x for improved result quality --- CHANGELOG.md | 2 ++ docs/configuration.md | 2 +- haiku_rag_slim/haiku/rag/client.py | 4 ++-- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 30c422de..2281297f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog ## [Unreleased] +- Increase reranking candidate retrieval multiplier from 3x to 10x for improved result quality + ## [0.17.0] - 2025-11-17 ### Added diff --git a/docs/configuration.md b/docs/configuration.md index 26961fb7..70291d6f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -476,7 +476,7 @@ See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the com ## Reranking -Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. +Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (10x the requested limit) and then reranks them to return the most relevant results. Reranking is **disabled by default** (`provider: ""`) for faster searches. You can enable it by configuring one of the providers below. diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index f26f611b..1cac8faa 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -466,8 +466,8 @@ class HaikuRAG: # No reranking - return direct search results return await self.chunk_repository.search(query, limit, search_type, filter) - # Get more initial results (3X) for reranking - search_limit = limit * 3 + # Get more initial results (10X) for reranking + search_limit = limit * 10 search_results = await self.chunk_repository.search( query, search_limit, search_type, filter )