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

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

@ -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})
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]
if isinstance(first_content, TextBlock):
return first_content.text
return ""
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 we've exhausted max rounds, return empty string
return ""
except ImportError:
pass

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

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

21
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,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

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

70
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")
with Progress() as progress:
task = progress.add_task("[green]Populating database...", total=len(corpus))
async with HaikuRAG(db_path) as rag:
for i, doc in enumerate(tqdm(corpus)):
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=doc["document_id"], # type: ignore
uri=uri,
)
progress.advance(task)
async def run_match_benchmark():
@ -36,8 +45,13 @@ async def run_match_benchmark():
correct_at_3 = 0
total_queries = 0
with Progress() as progress:
task = progress.add_task(
"[blue]Running retrieval benchmark...", total=len(corpus)
)
async with HaikuRAG(db_path) as rag:
for i, doc in enumerate(tqdm(corpus)):
for doc in corpus:
doc_id = doc["document_id"] # type: ignore
matches = await rag.search(
query=doc["question"], # type: ignore
@ -61,16 +75,18 @@ async def run_match_benchmark():
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,10 +103,13 @@ async def run_qa_benchmark(k: int | None = None):
correct_answers = 0
total_questions = 0
with Progress() as progress:
task = progress.add_task("[yellow]Running QA benchmark...", total=len(corpus))
async with HaikuRAG(db_path) as rag:
qa = get_qa_agent(rag)
for i, doc in enumerate(tqdm(corpus, desc="QA Benchmarking")):
for doc in corpus:
question = doc["question"] # type: ignore
expected_answer = doc["answer"] # type: ignore
@ -98,30 +117,33 @@ async def run_qa_benchmark(k: int | None = None):
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")
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."""