Give the results of search to the qa agent as json including the document uri

src/haiku/rag/qa/anthropic.py

@@ -1,8 +1,13 @@
 from collections.abc import Sequence
 
 try:
-    from anthropic import AsyncAnthropic
+    from anthropic import AsyncAnthropic  # type: ignore
-    from anthropic.types import MessageParam, TextBlock, ToolParam, ToolUseBlock
+    from anthropic.types import (  # type: ignore
+        MessageParam,
+        TextBlock,
+        ToolParam,
+        ToolUseBlock,
+    )
 
     from haiku.rag.client import HaikuRAG
     from haiku.rag.qa.base import QuestionAnswerAgentBase
@@ -73,13 +78,7 @@ try:
                                     query, limit=limit
                                 )
 
-                                context_chunks = []
+                                context = self._format_search_results(search_results)
-                                for chunk, score in search_results:
-                                    context_chunks.append(
-                                        f"Content: {chunk.content}\nScore: {score:.4f}"
-                                    )
-
-                                context = "\n\n".join(context_chunks)
 
                                 tool_results.append(
                                     {

src/haiku/rag/qa/base.py

@@ -1,3 +1,5 @@
+import json
+
 from haiku.rag.client import HaikuRAG
 from haiku.rag.qa.prompts import SYSTEM_PROMPT
 
@@ -15,6 +17,19 @@ class QuestionAnswerAgentBase:
             "QABase is an abstract class. Please implement the answer method in a subclass."
         )
 
+    def _format_search_results(self, search_results) -> str:
+        """Format search results as JSON list of {content, score, document_uri}"""
+        formatted_results = []
+        for chunk, score in search_results:
+            formatted_results.append(
+                {
+                    "content": chunk.content,
+                    "score": score,
+                    "document_uri": chunk.document_uri,
+                }
+            )
+        return json.dumps(formatted_results, indent=2)
+
     tools = [
         {
             "type": "function",

src/haiku/rag/qa/ollama.py

@@ -41,14 +41,7 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
 
                         search_results = await self._client.search(query, limit=limit)
 
-                        context_chunks = []
+                        context = self._format_search_results(search_results)
-                        for chunk, score in search_results:
-                            context_chunks.append(
-                                f"Content: {chunk.content}\nScore: {score:.4f}"
-                            )
-
-                        context = "\n\n".join(context_chunks)
-
                         messages.append(
                             {
                                 "role": "tool",

src/haiku/rag/qa/openai.py

@@ -1,15 +1,17 @@
 from collections.abc import Sequence
 
 try:
-    from openai import AsyncOpenAI
+    from openai import AsyncOpenAI  # type: ignore
-    from openai.types.chat import (
+    from openai.types.chat import (  # type: ignore
         ChatCompletionAssistantMessageParam,
         ChatCompletionMessageParam,
         ChatCompletionSystemMessageParam,
         ChatCompletionToolMessageParam,
         ChatCompletionUserMessageParam,
     )
-    from openai.types.chat.chat_completion_tool_param import ChatCompletionToolParam
+    from openai.types.chat.chat_completion_tool_param import (  # type: ignore
+        ChatCompletionToolParam,
+    )
 
     from haiku.rag.client import HaikuRAG
     from haiku.rag.qa.base import QuestionAnswerAgentBase
@@ -74,13 +76,7 @@ try:
                                 query, limit=limit
                             )
 
-                            context_chunks = []
+                            context = self._format_search_results(search_results)
-                            for chunk, score in search_results:
-                                context_chunks.append(
-                                    f"Content: {chunk.content}\nScore: {score:.4f}"
-                                )
-
-                            context = "\n\n".join(context_chunks)
 
                             messages.append(
                                 ChatCompletionToolMessageParam(