Merge pull request #14 from ggozad/feat/better-qa

Improve QA agent with multiple tool calls, provide performance benchmarks

BENCHMARKS.md

@@ -1,13 +0,0 @@
-# `haiku.rag` benchmarks
-
-We use [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) for the evaluation of `haiku.rag`
-
-* Recall
-
-We load the `News Stories` from `repliqa_3` which is 1035 documents, using `tests/generate_benchmark_db.py`, using the `mxbai-embed-large` Ollama embeddings.
-
-Subsequently, we run a search over the `question` for each row of the dataset and check whether we match the document that answers the question. The recall obtained is ~0.75 for matching in the top result, raising to ~0.75 for the top 3 results.
-
-* Question/Answer evaluation
-
-We use the  `News Stories` from `repliqa_3` using the `mxbai-embed-large` Ollama embeddings, with a QA agent also using Ollama with the `qwen3` model (8b). For each story we ask the `question` and use an LLM judge (also `qwen3`) to evaluate whether the answer is correct or not. Thus we obtain accuracy of ~0.54.

docs/benchmarks.md

@@ -0,0 +1,27 @@
+# Benchmarks
+
+We use the [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) dataset for the evaluation of `haiku.rag`.
+
+You can perform your own evaluations using as example the script found at
+`tests/generate_benchmark_db.py`.
+
+## Recall
+
+In order to calculate recall, we load the `News Stories` from `repliqa_3` which is 1035 documents and index them in a sqlite db. Subsequently, we run a search over the `question` field for each row of the dataset and check whether we match the document that answers the question.
+
+
+The recall obtained is ~0.73 for matching in the top result, raising to ~0.75 for the top 3 results.
+
+| Model                                 | Document in top 1 | Document in top 3 |
+|---------------------------------------|-------------------|-------------------|
+| Ollama / `mxbai-embed-large`          | 0.73              | 0.75              |
+| OpenAI / `text-embeddings-3-small`    | 0.75              | 0.88              |
+
+## Question/Answer evaluation
+
+Again using the same dataset, we use a QA agent to answer the question. In addition we use an LLM judge (using the Ollama `qwen3`) to evaluate whether the answer is correct or not. The obtained accuracy is as follows:
+
+| Embedding Model              | QA Model                          | Accuracy  |
+|------------------------------|-----------------------------------|-----------|
+| Ollama / `mxbai-embed-large` | Ollama / `qwen3`                  | 0.64      |
+| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7`   | 0.79      |

mkdocs.yml

@@ -63,6 +63,7 @@ nav:
       - Server: server.md
       - MCP: mcp.md
       - Python: python.md
+      - Benchmarks: benchmarks.md
 markdown_extensions:
   - admonition
   - attr_list

src/haiku/rag/qa/anthropic.py

@@ -37,75 +37,69 @@ try:
 
             messages: list[MessageParam] = [{"role": "user", "content": question}]
 
-            response = await anthropic_client.messages.create(
+            max_rounds = 5  # Prevent infinite loops
-                model=self._model,
-                max_tokens=4096,
-                system=self._system_prompt,
-                messages=messages,
-                tools=self.tools,
-                temperature=0.0,
-            )
 
-            if response.stop_reason == "tool_use":
+            for _ in range(max_rounds):
-                messages.append({"role": "assistant", "content": response.content})
+                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
+                if response.stop_reason == "tool_use":
-                tool_results = []
+                    messages.append({"role": "assistant", "content": response.content})
-                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
-                            )
 
-                            search_results = await self._client.search(
+                    # Process tool calls
-                                query, limit=limit
+                    tool_results = []
-                            )
+                    for content_block in response.content:
-
+                        if isinstance(content_block, ToolUseBlock):
-                            context_chunks = []
+                            if content_block.name == "search_documents":
-                            for chunk, score in search_results:
+                                args = content_block.input
-                                context_chunks.append(
+                                query = (
-                                    f"Content: {chunk.content}\nScore: {score:.4f}"
+                                    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(
+                                context_chunks = []
-                                {
+                                for chunk, score in search_results:
-                                    "type": "tool_result",
+                                    context_chunks.append(
-                                    "tool_use_id": content_block.id,
+                                        f"Content: {chunk.content}\nScore: {score:.4f}"
-                                    "content": context,
+                                    )
-                                }
-                            )
 
-                if tool_results:
+                                context = "\n\n".join(context_chunks)
-                    messages.append({"role": "user", "content": tool_results})
 
-                    final_response = await anthropic_client.messages.create(
+                                tool_results.append(
-                        model=self._model,
+                                    {
-                        max_tokens=4096,
+                                        "type": "tool_result",
-                        system=self._system_prompt,
+                                        "tool_use_id": content_block.id,
-                        messages=messages,
+                                        "content": context,
-                        temperature=0.0,
+                                    }
-                    )
+                                )
-                    if final_response.content:
+
-                        first_content = final_response.content[0]
+                    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:
+            # If we've exhausted max rounds, return empty string
-                first_content = response.content[0]
-                if isinstance(first_content, TextBlock):
-                    return first_content.text
             return ""
 
 except ImportError:

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
+        max_rounds = 5  # Prevent infinite loops
-        response = await ollama_client.chat(
-            model=self._model,
-            messages=messages,
-            tools=self.tools,
-            options=OLLAMA_OPTIONS,
-            think=False,
-        )
 
-        if response.get("message", {}).get("tool_calls"):
+        for _ in range(max_rounds):
-            for tool_call in response["message"]["tool_calls"]:
+            response = await ollama_client.chat(
-                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(
                 model=self._model,
                 messages=messages,
-                think=False,
+                tools=self.tools,
                 options=OLLAMA_OPTIONS,
+                think=False,
             )
-            return final_response["message"]["content"]
+
-        else:
+            if response.get("message", {}).get("tool_calls"):
-            return response["message"]["content"]
+                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 ""

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
+            max_rounds = 5  # Prevent infinite loops
-            response = await openai_client.chat.completions.create(
-                model=self._model,
-                messages=messages,
-                tools=self.tools,
-                temperature=0.0,
-            )
 
-            response_message = response.choices[0].message
+            for _ in range(max_rounds):
-
+                response = await openai_client.chat.completions.create(
-            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(
                     model=self._model,
                     messages=messages,
+                    tools=self.tools,
                     temperature=0.0,
                 )
-                return final_response.choices[0].message.content or ""
+
-            else:
+                response_message = response.choices[0].message
-                return response_message.content or ""
+
+                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

src/haiku/rag/qa/prompts.py

@@ -1,7 +1,20 @@
 SYSTEM_PROMPT = """
-You are a helpful assistant that uses a RAG library to answer the user's prompt.
+You are a knowledgeable assistant that helps users find information from a document knowledge base.
-Your task is to provide a concise and accurate answer based on the provided context.
+
-You should ask the provided tools to find relevant documents and then use the content of those documents to answer the question.
+Your process:
-Never make up information, always use the context to answer the question.
+1. When a user asks a question, use the search_documents tool to find relevant information
-If the context does not contain enough information to answer the question, respond with "I cannot answer that based on the provided context."
+2. Search with specific keywords and phrases from the user's question
+3. Review the search results and their relevance scores
+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
+- Quote or reference specific information when possible
+- If multiple documents contain relevant information, synthesize them coherently
+- Indicate when information is incomplete or when you need to search for additional context
+- 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 concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
 """

tests/generate_benchmark_db.py

@@ -3,28 +3,37 @@ from pathlib import Path
 
 from datasets import Dataset, load_dataset
 from llm_judge import LLMJudge
-from tqdm import tqdm
+from rich.console import Console
+from rich.progress import Progress
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.qa import get_qa_agent
 
+console = Console()
+
 db_path = Path(__file__).parent / "data" / "benchmark.sqlite"
 
 
 async def populate_db():
-    if (db_path).exists():
-        print("Benchmark database already exists. Skipping creation.")
-        return
-
     ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]  # type: ignore
     corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
 
-    async with HaikuRAG(db_path) as rag:
+    with Progress() as progress:
-        for i, doc in enumerate(tqdm(corpus)):
+        task = progress.add_task("[green]Populating database...", total=len(corpus))
-            await rag.create_document(
+
-                content=doc["document_extracted"],  # type: ignore
+        async with HaikuRAG(db_path) as rag:
-                uri=doc["document_id"],  # type: ignore
+            for doc in corpus:
-            )
+                uri = doc["document_id"]  # type: ignore
+                existing_doc = await rag.get_document_by_uri(uri)
+                if existing_doc is not None:
+                    progress.advance(task)
+                    continue
+
+                await rag.create_document(
+                    content=doc["document_extracted"],  # type: ignore
+                    uri=uri,
+                )
+                progress.advance(task)
 
 
 async def run_match_benchmark():
@@ -36,41 +45,48 @@ async def run_match_benchmark():
     correct_at_3 = 0
     total_queries = 0
 
-    async with HaikuRAG(db_path) as rag:
+    with Progress() as progress:
-        for i, doc in enumerate(tqdm(corpus)):
+        task = progress.add_task(
-            doc_id = doc["document_id"]  # type: ignore
+            "[blue]Running retrieval benchmark...", total=len(corpus)
-            matches = await rag.search(
+        )
-                query=doc["question"],  # type: ignore
-                limit=3,
-            )
 
-            total_queries += 1
+        async with HaikuRAG(db_path) as rag:
+            for doc in corpus:
+                doc_id = doc["document_id"]  # type: ignore
+                matches = await rag.search(
+                    query=doc["question"],  # type: ignore
+                    limit=3,
+                )
 
-            # Check position of correct document in results
+                total_queries += 1
-            for position, (chunk, _) in enumerate(matches):
+
-                retrieved = await rag.get_document_by_id(chunk.document_id)
+                # Check position of correct document in results
-                if retrieved and retrieved.uri == doc_id:
+                for position, (chunk, _) in enumerate(matches):
-                    if position == 0:  # First position
+                    retrieved = await rag.get_document_by_id(chunk.document_id)
-                        correct_at_1 += 1
+                    if retrieved and retrieved.uri == doc_id:
-                        correct_at_2 += 1
+                        if position == 0:  # First position
-                        correct_at_3 += 1
+                            correct_at_1 += 1
-                    elif position == 1:  # Second position
+                            correct_at_2 += 1
-                        correct_at_2 += 1
+                            correct_at_3 += 1
-                        correct_at_3 += 1
+                        elif position == 1:  # Second position
-                    elif position == 2:  # Third position
+                            correct_at_2 += 1
-                        correct_at_3 += 1
+                            correct_at_3 += 1
-                    break
+                        elif position == 2:  # Third position
+                            correct_at_3 += 1
+                        break
+
+                progress.advance(task)
 
     # Calculate recall metrics
     recall_at_1 = correct_at_1 / total_queries
     recall_at_2 = correct_at_2 / total_queries
     recall_at_3 = correct_at_3 / total_queries
 
-    print("\n=== Retrieval Benchmark Results ===")
+    console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
-    print(f"Total queries: {total_queries}")
+    console.print(f"Total queries: {total_queries}")
-    print(f"Recall@1: {recall_at_1:.4f}")
+    console.print(f"Recall@1: {recall_at_1:.4f}")
-    print(f"Recall@2: {recall_at_2:.4f}")
+    console.print(f"Recall@2: {recall_at_2:.4f}")
-    print(f"Recall@3: {recall_at_3:.4f}")
+    console.print(f"Recall@3: {recall_at_3:.4f}")
 
     return {"recall@1": recall_at_1, "recall@2": recall_at_2, "recall@3": recall_at_3}
 
@@ -87,41 +103,47 @@ async def run_qa_benchmark(k: int | None = None):
     correct_answers = 0
     total_questions = 0
 
-    async with HaikuRAG(db_path) as rag:
+    with Progress() as progress:
-        qa = get_qa_agent(rag)
+        task = progress.add_task("[yellow]Running QA benchmark...", total=len(corpus))
 
-        for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")):
+        async with HaikuRAG(db_path) as rag:
-            question = doc["question"]  # type: ignore
+            qa = get_qa_agent(rag)
-            expected_answer = doc["answer"]  # type: ignore
 
-            generated_answer = await qa.answer(question)
+            for doc in corpus:
-            is_equivalent = await judge.judge_answers(
+                question = doc["question"]  # type: ignore
-                question, generated_answer, expected_answer
+                expected_answer = doc["answer"]  # type: ignore
-            )
-            print(f"Question: {question}")
-            print(f"Expected: {expected_answer}")
-            print(f"Generated: {generated_answer}")
-            print(f"Equivalent: {is_equivalent}\n")
 
-            if is_equivalent:
+                generated_answer = await qa.answer(question)
-                correct_answers += 1
+                is_equivalent = await judge.judge_answers(
-            total_questions += 1
+                    question, generated_answer, expected_answer
+                )
+                console.print(f"Question: {question}")
+                console.print(f"Expected: {expected_answer}")
+                console.print(f"Generated: {generated_answer}")
+                console.print(f"Equivalent: {is_equivalent}\n")
+
+                if is_equivalent:
+                    correct_answers += 1
+                total_questions += 1
+                console.print("Current score:", correct_answers, "/", total_questions)
+
+                progress.advance(task)
 
     accuracy = correct_answers / total_questions if total_questions > 0 else 0
 
-    print("\n=== QA Benchmark Results ===")
+    console.print("\n=== QA Benchmark Results ===", style="bold cyan")
-    print(f"Total questions: {total_questions}")
+    console.print(f"Total questions: {total_questions}")
-    print(f"Correct answers: {correct_answers}")
+    console.print(f"Correct answers: {correct_answers}")
-    print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
+    console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
 
 
 async def main():
     await populate_db()
 
-    print("Running retrieval benchmarks...")
+    console.print("Running retrieval benchmarks...", style="bold blue")
     await run_match_benchmark()
 
-    print("\nRunning QA benchmarks...")
+    console.print("\nRunning QA benchmarks...", style="bold yellow")
     await run_qa_benchmark()
 
 

tests/llm_judge.py

@@ -49,6 +49,7 @@ class LLMJudge:
         1. Do both answers provide the same answer?
         2. Do both answers directly address the question asked?
         3. Minor differences in wording or style are acceptable if the meaning of the answer is the same.
+        4. If one answer is more detailed but the other is correct, they can still be considered equivalent.
 
         Be strict but fair in your evaluation. Focus on factual correctness and whether both answers would satisfy someone asking the question."""