danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
Improve QA agent with multiple tool calls, provide performance benchmarks
This commit is contained in:
Yiorgis Gozadinos 2025-07-09 10:46:28 +03:00 committed by GitHub
commit 60b072a1ee
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 277 additions and 236 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
BENCHMARKS.md
View file

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

27
docs/benchmarks.md Normal file
View file

@ -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 |

1
mkdocs.yml
View file

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

108
src/haiku/rag/qa/anthropic.py
View file

@ -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:

79
src/haiku/rag/qa/ollama.py
View file

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

119
src/haiku/rag/qa/openai.py
View file

@ -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

23
src/haiku/rag/qa/prompts.py
View file

@ -1,7 +1,20 @@
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. 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.
"""

142
tests/generate_benchmark_db.py
View file

@ -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,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:
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
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()

1
tests/llm_judge.py
View file

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