Merge pull request #143 from ggozad/fix/rerank-config

Increase reranking candidate retrieval multiplier from 3x to 10x for improved result quality

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

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.
 

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/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
         )

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