diff --git a/src/haiku/rag/qa/anthropic.py b/src/haiku/rag/qa/anthropic.py index 8827c5cb..93b821f9 100644 --- a/src/haiku/rag/qa/anthropic.py +++ b/src/haiku/rag/qa/anthropic.py @@ -1,8 +1,13 @@ from collections.abc import Sequence try: - from anthropic import AsyncAnthropic - from anthropic.types import MessageParam, TextBlock, ToolParam, ToolUseBlock + from anthropic import AsyncAnthropic # type: ignore + 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 = [] - for chunk, score in search_results: - context_chunks.append( - f"Content: {chunk.content}\nScore: {score:.4f}" - ) - - context = "\n\n".join(context_chunks) + context = self._format_search_results(search_results) tool_results.append( { diff --git a/src/haiku/rag/qa/base.py b/src/haiku/rag/qa/base.py index 0ff2a55b..b5293093 100644 --- a/src/haiku/rag/qa/base.py +++ b/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", diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index 7521d8bf..f49315c0 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/src/haiku/rag/qa/ollama.py @@ -41,14 +41,7 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase): search_results = await self._client.search(query, limit=limit) - context_chunks = [] - for chunk, score in search_results: - context_chunks.append( - f"Content: {chunk.content}\nScore: {score:.4f}" - ) - - context = "\n\n".join(context_chunks) - + context = self._format_search_results(search_results) messages.append( { "role": "tool", diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py index 24f58cf9..54b2fc3e 100644 --- a/src/haiku/rag/qa/openai.py +++ b/src/haiku/rag/qa/openai.py @@ -1,15 +1,17 @@ from collections.abc import Sequence try: - from openai import AsyncOpenAI - from openai.types.chat import ( + from openai import AsyncOpenAI # type: ignore + 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 = [] - for chunk, score in search_results: - context_chunks.append( - f"Content: {chunk.content}\nScore: {score:.4f}" - ) - - context = "\n\n".join(context_chunks) + context = self._format_search_results(search_results) messages.append( ChatCompletionToolMessageParam(