From bea96f1b1b9acd1f86d365cdd35a6d7c523d6816 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 4 Jul 2025 13:29:51 +0300 Subject: [PATCH 1/4] Better prompts --- src/haiku/rag/qa/prompts.py | 22 +++++++++++++++++----- tests/llm_judge.py | 1 + 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/haiku/rag/qa/prompts.py b/src/haiku/rag/qa/prompts.py index fc8f2c9b..6a63ccb3 100644 --- a/src/haiku/rag/qa/prompts.py +++ b/src/haiku/rag/qa/prompts.py @@ -1,7 +1,19 @@ SYSTEM_PROMPT = """ -You are a helpful assistant that uses a RAG library to answer the user's prompt. -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. -Never make up information, always use the context to answer the question. -If the context does not contain enough information to answer the question, respond with "I cannot answer that based on the provided context." +You are a knowledgeable assistant that helps users find information from a document knowledge base. + +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 + +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 thorough but concise, and always maintain accuracy over completeness. """ diff --git a/tests/llm_judge.py b/tests/llm_judge.py index 66bfd2cb..5af4cf0e 100644 --- a/tests/llm_judge.py +++ b/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.""" From 2f261a33f5f45ee7e3dfd33dcabdfb6493ddb74a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 8 Jul 2025 12:51:58 +0300 Subject: [PATCH 2/4] 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. """ From 9d72a158a6a5a8fd23bf0afa67842b1346c2b1c8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 8 Jul 2025 19:01:54 +0300 Subject: [PATCH 3/4] Document benchmarks --- docs/benchmarks.md | 28 +++++++++++++++++++++------- mkdocs.yml | 1 + tests/generate_benchmark_db.py | 1 + 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 977e1927..e1fc2bf7 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -1,13 +1,27 @@ -# `haiku.rag` benchmarks +# Benchmarks -We use [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) for the evaluation of `haiku.rag` +We use the [repliqa](https://huggingface.co/datasets/ServiceNow/repliqa) dataset for the evaluation of `haiku.rag`. -* Recall +You can perform your own evaluations using as example the script found at +`tests/generate_benchmark_db.py`. -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. +## Recall -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. +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. -* 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. +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` | | | + +## 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 | diff --git a/mkdocs.yml b/mkdocs.yml index e9767caa..07bcb523 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -63,6 +63,7 @@ nav: - Server: server.md - MCP: mcp.md - Python: python.md + - Benchmarks: benchmarks.md markdown_extensions: - admonition - attr_list diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index b1ade523..20bb01b2 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -106,6 +106,7 @@ async def run_qa_benchmark(k: int | None = None): if is_equivalent: correct_answers += 1 total_questions += 1 + print("Current score:", correct_answers, "/", total_questions) accuracy = correct_answers / total_questions if total_questions > 0 else 0 From 42e62f92a1b44237f6f4d0c5e715bcd92c2fa35d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 9 Jul 2025 10:27:32 +0300 Subject: [PATCH 4/4] Better formatting of benchmark script --- docs/benchmarks.md | 2 +- tests/generate_benchmark_db.py | 143 +++++++++++++++++++-------------- 2 files changed, 83 insertions(+), 62 deletions(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index e1fc2bf7..948cba93 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -15,7 +15,7 @@ The recall obtained is ~0.73 for matching in the top result, raising to ~0.75 fo | Model | Document in top 1 | Document in top 3 | |---------------------------------------|-------------------|-------------------| | Ollama / `mxbai-embed-large` | 0.73 | 0.75 | -| OpenAI / `text-embeddings-3-small` | | | +| OpenAI / `text-embeddings-3-small` | 0.75 | 0.88 | ## Question/Answer evaluation diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index 20bb01b2..70b9c468 100644 --- a/tests/generate_benchmark_db.py +++ b/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: - for i, doc in enumerate(tqdm(corpus)): - await rag.create_document( - content=doc["document_extracted"], # type: ignore - uri=doc["document_id"], # type: ignore - ) + with Progress() as progress: + task = progress.add_task("[green]Populating database...", total=len(corpus)) + + async with HaikuRAG(db_path) as rag: + 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: - for i, doc in enumerate(tqdm(corpus)): - doc_id = doc["document_id"] # type: ignore - matches = await rag.search( - query=doc["question"], # type: ignore - limit=3, - ) + with Progress() as progress: + task = progress.add_task( + "[blue]Running retrieval benchmark...", total=len(corpus) + ) - 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 - for position, (chunk, _) in enumerate(matches): - retrieved = await rag.get_document_by_id(chunk.document_id) - if retrieved and retrieved.uri == doc_id: - if position == 0: # First position - correct_at_1 += 1 - correct_at_2 += 1 - correct_at_3 += 1 - elif position == 1: # Second position - correct_at_2 += 1 - correct_at_3 += 1 - elif position == 2: # Third position - correct_at_3 += 1 - break + total_queries += 1 + + # Check position of correct document in results + for position, (chunk, _) in enumerate(matches): + retrieved = await rag.get_document_by_id(chunk.document_id) + if retrieved and retrieved.uri == doc_id: + if position == 0: # First position + correct_at_1 += 1 + correct_at_2 += 1 + correct_at_3 += 1 + elif position == 1: # Second position + correct_at_2 += 1 + correct_at_3 += 1 + 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 ===") - print(f"Total queries: {total_queries}") - print(f"Recall@1: {recall_at_1:.4f}") - print(f"Recall@2: {recall_at_2:.4f}") - print(f"Recall@3: {recall_at_3:.4f}") + console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan") + console.print(f"Total queries: {total_queries}") + console.print(f"Recall@1: {recall_at_1:.4f}") + console.print(f"Recall@2: {recall_at_2:.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,42 +103,47 @@ async def run_qa_benchmark(k: int | None = None): correct_answers = 0 total_questions = 0 - async with HaikuRAG(db_path) as rag: - qa = get_qa_agent(rag) + with Progress() as progress: + task = progress.add_task("[yellow]Running QA benchmark...", total=len(corpus)) - for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")): - question = doc["question"] # type: ignore - expected_answer = doc["answer"] # type: ignore + async with HaikuRAG(db_path) as rag: + qa = get_qa_agent(rag) - generated_answer = await qa.answer(question) - is_equivalent = await judge.judge_answers( - question, generated_answer, expected_answer - ) - print(f"Question: {question}") - print(f"Expected: {expected_answer}") - print(f"Generated: {generated_answer}") - print(f"Equivalent: {is_equivalent}\n") + for doc in corpus: + question = doc["question"] # type: ignore + expected_answer = doc["answer"] # type: ignore - if is_equivalent: - correct_answers += 1 - total_questions += 1 - print("Current score:", correct_answers, "/", total_questions) + generated_answer = await qa.answer(question) + is_equivalent = await judge.judge_answers( + 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 ===") - print(f"Total questions: {total_questions}") - print(f"Correct answers: {correct_answers}") - print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") + console.print("\n=== QA Benchmark Results ===", style="bold cyan") + console.print(f"Total questions: {total_questions}") + console.print(f"Correct answers: {correct_answers}") + 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()