From 2f261a33f5f45ee7e3dfd33dcabdfb6493ddb74a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 8 Jul 2025 12:51:58 +0300 Subject: [PATCH] Allow multiple tool calling rounds in QA agent --- BENCHMARKS.md => docs/benchmarks.md | 0 src/haiku/rag/qa/anthropic.py | 108 ++++++++++++------------- src/haiku/rag/qa/ollama.py | 79 +++++++++--------- src/haiku/rag/qa/openai.py | 119 ++++++++++++++-------------- src/haiku/rag/qa/prompts.py | 5 +- 5 files changed, 151 insertions(+), 160 deletions(-) rename BENCHMARKS.md => docs/benchmarks.md (100%) diff --git a/BENCHMARKS.md b/docs/benchmarks.md similarity index 100% rename from BENCHMARKS.md rename to docs/benchmarks.md diff --git a/src/haiku/rag/qa/anthropic.py b/src/haiku/rag/qa/anthropic.py index 5b4479b3..8827c5cb 100644 --- a/src/haiku/rag/qa/anthropic.py +++ b/src/haiku/rag/qa/anthropic.py @@ -37,75 +37,69 @@ try: messages: list[MessageParam] = [{"role": "user", "content": question}] - response = await anthropic_client.messages.create( - model=self._model, - max_tokens=4096, - system=self._system_prompt, - messages=messages, - tools=self.tools, - temperature=0.0, - ) + max_rounds = 5 # Prevent infinite loops - if response.stop_reason == "tool_use": - messages.append({"role": "assistant", "content": response.content}) + for _ in range(max_rounds): + response = await anthropic_client.messages.create( + model=self._model, + max_tokens=4096, + system=self._system_prompt, + messages=messages, + tools=self.tools, + temperature=0.0, + ) - # Process tool calls - tool_results = [] - for content_block in response.content: - if isinstance(content_block, ToolUseBlock): - if content_block.name == "search_documents": - args = content_block.input - query = ( - args.get("query", question) - if isinstance(args, dict) - else question - ) - limit = ( - int(args.get("limit", 3)) - if isinstance(args, dict) - else 3 - ) + if response.stop_reason == "tool_use": + messages.append({"role": "assistant", "content": response.content}) - 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}" + # Process tool calls + tool_results = [] + for content_block in response.content: + if isinstance(content_block, ToolUseBlock): + if content_block.name == "search_documents": + args = content_block.input + query = ( + args.get("query", question) + if isinstance(args, dict) + else question + ) + limit = ( + int(args.get("limit", 3)) + if isinstance(args, dict) + else 3 ) - context = "\n\n".join(context_chunks) + search_results = await self._client.search( + query, limit=limit + ) - tool_results.append( - { - "type": "tool_result", - "tool_use_id": content_block.id, - "content": context, - } - ) + context_chunks = [] + for chunk, score in search_results: + context_chunks.append( + f"Content: {chunk.content}\nScore: {score:.4f}" + ) - if tool_results: - messages.append({"role": "user", "content": tool_results}) + context = "\n\n".join(context_chunks) - final_response = await anthropic_client.messages.create( - model=self._model, - max_tokens=4096, - system=self._system_prompt, - messages=messages, - temperature=0.0, - ) - if final_response.content: - first_content = final_response.content[0] + tool_results.append( + { + "type": "tool_result", + "tool_use_id": content_block.id, + "content": context, + } + ) + + if tool_results: + messages.append({"role": "user", "content": tool_results}) + else: + # No tool use, return the response + if response.content: + first_content = response.content[0] if isinstance(first_content, TextBlock): return first_content.text return "" - if response.content: - first_content = response.content[0] - if isinstance(first_content, TextBlock): - return first_content.text + # If we've exhausted max rounds, return empty string return "" except ImportError: diff --git a/src/haiku/rag/qa/ollama.py b/src/haiku/rag/qa/ollama.py index c8cac4ce..9c4ee01a 100644 --- a/src/haiku/rag/qa/ollama.py +++ b/src/haiku/rag/qa/ollama.py @@ -14,54 +14,51 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase): async def answer(self, question: str) -> str: ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL) - # Define the search tool - messages = [ {"role": "system", "content": self._system_prompt}, {"role": "user", "content": question}, ] - # Initial response with tool calling - response = await ollama_client.chat( - model=self._model, - messages=messages, - tools=self.tools, - options=OLLAMA_OPTIONS, - think=False, - ) + max_rounds = 5 # Prevent infinite loops - if response.get("message", {}).get("tool_calls"): - for tool_call in response["message"]["tool_calls"]: - if tool_call["function"]["name"] == "search_documents": - args = tool_call["function"]["arguments"] - query = args.get("query", question) - limit = int(args.get("limit", 3)) - - 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) - - messages.append(response["message"]) - messages.append( - { - "role": "tool", - "content": context, - "tool_call_id": tool_call.get("id", "search_tool"), - } - ) - - final_response = await ollama_client.chat( + for _ in range(max_rounds): + response = await ollama_client.chat( model=self._model, messages=messages, - think=False, + tools=self.tools, options=OLLAMA_OPTIONS, + think=False, ) - return final_response["message"]["content"] - else: - return response["message"]["content"] + + if response.get("message", {}).get("tool_calls"): + messages.append(response["message"]) + + for tool_call in response["message"]["tool_calls"]: + if tool_call["function"]["name"] == "search_documents": + args = tool_call["function"]["arguments"] + query = args.get("query", question) + limit = int(args.get("limit", 3)) + + 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) + + messages.append( + { + "role": "tool", + "content": context, + "tool_call_id": tool_call.get("id", "search_tool"), + } + ) + else: + # No tool calls, return the response + return response["message"]["content"] + + # If we've exhausted max rounds, return empty string + return "" diff --git a/src/haiku/rag/qa/openai.py b/src/haiku/rag/qa/openai.py index f75a7396..24f58cf9 100644 --- a/src/haiku/rag/qa/openai.py +++ b/src/haiku/rag/qa/openai.py @@ -24,8 +24,6 @@ try: async def answer(self, question: str) -> str: openai_client = AsyncOpenAI() - # Define the search tool - messages: list[ChatCompletionMessageParam] = [ ChatCompletionSystemMessageParam( role="system", content=self._system_prompt @@ -33,69 +31,70 @@ try: ChatCompletionUserMessageParam(role="user", content=question), ] - # Initial response with tool calling - response = await openai_client.chat.completions.create( - model=self._model, - messages=messages, - tools=self.tools, - temperature=0.0, - ) + max_rounds = 5 # Prevent infinite loops - response_message = response.choices[0].message - - if response_message.tool_calls: - messages.append( - ChatCompletionAssistantMessageParam( - role="assistant", - content=response_message.content, - tool_calls=[ - { - "id": tc.id, - "type": "function", - "function": { - "name": tc.function.name, - "arguments": tc.function.arguments, - }, - } - for tc in response_message.tool_calls - ], - ) - ) - - for tool_call in response_message.tool_calls: - if tool_call.function.name == "search_documents": - import json - - args = json.loads(tool_call.function.arguments) - query = args.get("query", question) - limit = int(args.get("limit", 3)) - - 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) - - messages.append( - ChatCompletionToolMessageParam( - role="tool", - content=context, - tool_call_id=tool_call.id, - ) - ) - - final_response = await openai_client.chat.completions.create( + for _ in range(max_rounds): + response = await openai_client.chat.completions.create( model=self._model, messages=messages, + tools=self.tools, temperature=0.0, ) - return final_response.choices[0].message.content or "" - else: - return response_message.content or "" + + response_message = response.choices[0].message + + if response_message.tool_calls: + messages.append( + ChatCompletionAssistantMessageParam( + role="assistant", + content=response_message.content, + tool_calls=[ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in response_message.tool_calls + ], + ) + ) + + for tool_call in response_message.tool_calls: + if tool_call.function.name == "search_documents": + import json + + args = json.loads(tool_call.function.arguments) + query = args.get("query", question) + limit = int(args.get("limit", 3)) + + 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) + + messages.append( + ChatCompletionToolMessageParam( + role="tool", + content=context, + tool_call_id=tool_call.id, + ) + ) + else: + # No tool calls, return the response + return response_message.content or "" + + # If we've exhausted max rounds, return empty string + return "" except ImportError: pass diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index 6a63ccb3..283c40e2 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/src/haiku/rag/qa/prompts.py @@ -5,7 +5,8 @@ Your process: 1. When a user asks a question, use the search_documents tool to find relevant information 2. Search with specific keywords and phrases from the user's question 3. Review the search results and their relevance scores -4. Provide a comprehensive answer based only on the retrieved documents +4. If you need additional context, perform follow-up searches with different keywords +5. Provide a comprehensive answer based only on the retrieved documents Guidelines: - Base your answers strictly on the provided document content @@ -15,5 +16,5 @@ Guidelines: - If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question." - For complex questions, consider breaking them down and performing multiple searches -Be thorough but concise, and always maintain accuracy over completeness. +Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents. """