Allow multiple tool calling rounds in QA agent

src/haiku/rag/qa/anthropic.py

@@ -37,6 +37,9 @@ try:
 
             messages: list[MessageParam] = [{"role": "user", "content": question}]
 
+            max_rounds = 5  # Prevent infinite loops
+
+            for _ in range(max_rounds):
                 response = await anthropic_client.messages.create(
                     model=self._model,
                     max_tokens=4096,
@@ -88,25 +91,16 @@ try:
 
                     if tool_results:
                         messages.append({"role": "user", "content": tool_results})
-
+                else:
-                    final_response = await anthropic_client.messages.create(
+                    # No tool use, return the response
-                        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]
-                        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
                     return ""
 
+            # If we've exhausted max rounds, return empty string
+            return ""
+
 except ImportError:
     pass

src/haiku/rag/qa/ollama.py

@@ -14,14 +14,14 @@ 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
+        max_rounds = 5  # Prevent infinite loops
+
+        for _ in range(max_rounds):
             response = await ollama_client.chat(
                 model=self._model,
                 messages=messages,
@@ -31,6 +31,8 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
             )
 
             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"]
@@ -47,7 +49,6 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
 
                         context = "\n\n".join(context_chunks)
 
-                    messages.append(response["message"])
                         messages.append(
                             {
                                 "role": "tool",
@@ -55,13 +56,9 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
                                 "tool_call_id": tool_call.get("id", "search_tool"),
                             }
                         )
-
-            final_response = await ollama_client.chat(
-                model=self._model,
-                messages=messages,
-                think=False,
-                options=OLLAMA_OPTIONS,
-            )
-            return final_response["message"]["content"]
             else:
+                # No tool calls, return the response
                 return response["message"]["content"]
+
+        # If we've exhausted max rounds, return empty string
+        return ""

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,7 +31,9 @@ try:
                 ChatCompletionUserMessageParam(role="user", content=question),
             ]
 
-            # Initial response with tool calling
+            max_rounds = 5  # Prevent infinite loops
+
+            for _ in range(max_rounds):
                 response = await openai_client.chat.completions.create(
                     model=self._model,
                     messages=messages,
@@ -70,7 +70,9 @@ try:
                             query = args.get("query", question)
                             limit = int(args.get("limit", 3))
 
-                        search_results = await self._client.search(query, limit=limit)
+                            search_results = await self._client.search(
+                                query, limit=limit
+                            )
 
                             context_chunks = []
                             for chunk, score in search_results:
@@ -87,15 +89,12 @@ try:
                                     tool_call_id=tool_call.id,
                                 )
                             )
-
-                final_response = await openai_client.chat.completions.create(
-                    model=self._model,
-                    messages=messages,
-                    temperature=0.0,
-                )
-                return final_response.choices[0].message.content or ""
                 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

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.
 """