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

Merge pull request #160 from ggozad/feat/evaluation-enhancements

Browse source 
Evaluation improvements (metrics/utils).
This commit is contained in:
Yiorgis Gozadinos 2025-11-24 15:02:42 +02:00 committed by GitHub
commit 3bdbd10516
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 280 additions and 131 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
CHANGELOG.md
View file

@ -13,9 +13,18 @@
- Tracks retrieval configuration: `retrieval_limit` for number of chunks retrieved
- Tracks reranking configuration: `rerank_provider` and `rerank_model`
- Enables comparison of evaluation runs with different configurations in Logfire
- **Evaluations**: Refactored retrieval evaluation to use pydantic-ai experiment framework
- New `evaluators` module with `MRREvaluator` (Mean Reciprocal Rank) and `MAPEvaluator` (Mean Average Precision)
- Retrieval benchmarks now use `Dataset.evaluate()` with full Logfire experiment tracking
- Dataset specifications now declare their retrieval evaluator (MRR for RepliQA, MAP for Wix)
- Replaced Recall@K and Success@K with industry-standard MRR and MAP metrics
- Unified evaluation framework for both retrieval and QA benchmarks
### Changed
- **Evaluations**: Renamed `--qa-limit` CLI parameter to `--limit`, now applies to both retrieval and QA benchmarks
- **Evaluations**: Retrieval evaluator selection moved from runtime logic to dataset configuration
## [0.18.0] - 2025-11-21
### Added

71
docs/benchmarks.md
View file

@ -25,24 +25,33 @@ You can also use command-line options:
- `--skip-db` - Skip updating the evaluation database
- `--skip-retrieval` - Skip retrieval benchmark
- `--skip-qa` - Skip QA benchmark
- `--qa-limit N` - Limit number of QA cases to evaluate
- `--name NAME` - Override the evaluation name (defaults to `{dataset}_qa_evaluation`)
- `--limit N` - Limit number of test cases for both retrieval and QA
- `--name NAME` - Override the evaluation name (defaults to `{dataset}_retrieval_evaluation` or `{dataset}_qa_evaluation`)
## Recall
## RepliQA Retrieval
In order to calculate recall, we load the `News Stories` from `repliqa_3` (1035 documents) and index them. 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. Questions for which the answer cannot be found in the documents are ignored.
We use the [RepliQA](https://huggingface.co/datasets/ServiceNow/repliqa) dataset to evaluate retrieval performance. We load the `News Stories` from `repliqa_3` (1035 documents) and index them. 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. Questions for which the answer cannot be found in the documents are ignored.
For RepliQA, we use **Mean Reciprocal Rank (MRR)** as the primary metric since each query has exactly one relevant document.
The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 for the top 3 results with the "bare" default settings (Ollama `qwen3`, `mxbai-embed-large` embeddings, no reranking).
**How MRR is calculated:**
- For each query, we retrieve the top-K documents and find the rank (position) of the first relevant document
- The reciprocal rank for that query is `1/rank` (e.g., if the relevant document is at position 3, the score is 1/3 ≈ 0.333)
- If no relevant document is found in the top-K results, the score is 0
- MRR is the mean of these reciprocal ranks across all queries
- Scores range from 0 (never found) to 1 (always found at rank 1)
| Embedding Model | Document in top 1 | Document in top 3 | Reranker |
|---------------------------------------|-------------------|-------------------|------------------------|
| Ollama / `qwen3-embedding:8b` | 0.81 | 0.95 | None |
| Ollama / `qwen3-embedding:0.6b` | 0.77 | 0.97 | None |
| Ollama / `qwen3-embedding:8b` | 0.91 | 0.98 | `mxbai-rerank-base-v2` |
| Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None |
| Ollama / `mxbai-embed-large` | 0.90 | 0.95 | `mxbai-rerank-base-v2` |
| Ollama / `nomic-embed-text-v1.5` | 0.74 | 0.90 | None |
**Example:** If we run 3 queries and the relevant documents are found at positions 1, 2, and not found:
- Query 1: 1/1 = 1.0
- Query 2: 1/2 = 0.5
- Query 3: 0 (not found)
- MRR = (1.0 + 0.5 + 0) / 3 = 0.5
### MRR Results
| Embedding Model | MRR | Reranker |
|---------------------------------------|-------|------------------------|
| Ollama / `qwen3-embedding:8b` | 0.91 | - |
## Question/Answer evaluation
@ -60,7 +69,7 @@ determine whether the answer is correct. The obtained accuracy is as follows:
Note the significant degradation when very small models are used such as `qwen3:0.6b`.
## Wix dataset
## Wix Retrieval
We also track retrieval performance on [WixQA](https://huggingface.co/datasets/Wix/WixQA),
a dataset of real customer support questions paired with curated answers from
@ -70,26 +79,30 @@ system handles conversational, product-specific support queries.
For retrieval evaluation, we index the reference answer passages shipped with the dataset and
run retrieval against each user question. Each sample supplies one or more
relevant passage URIs. We track two complementary metrics:
relevant passage URIs.
- **Recall@K**: Fraction of relevant documents retrieved in top K results. Measures coverage.
- **Success@K**: Fraction of queries with at least one relevant document in top K. Most relevant for RAG, where finding one good document is often sufficient.
For Wix, we use **Mean Average Precision (MAP)** as the primary metric since each query has multiple relevant documents. MAP accounts for both the presence and ranking of all relevant documents.
### Recall@K Results
**How MAP is calculated:**
- For each query, we retrieve the top-K documents and identify which ones are relevant
- For each relevant document found at position k, we calculate precision@k = (number of relevant docs in top k) / k
- Average Precision (AP) for that query is the mean of these precision values, divided by the total number of relevant documents
- MAP is the mean of AP scores across all queries
- Scores range from 0 (no relevant documents found) to 1 (all relevant documents ranked at the top)
| Embedding Model | Recall@1 | Recall@3 | Recall@5 | Reranker |
|----------------------------|----------|----------|----------|------------------------|
| `qwen3-embedding:8b` | 0.31 | 0.48 | 0.54 | None |
| `qwen3-embedding:8b` | 0.36 | 0.57 | 0.68 | `mxbai-rerank-base-v2` |
| `qwen3-embedding:8b` | 0.36 | 0.58 | 0.67 | `zeroentropy` |
**Example:** If a query has 2 relevant documents (A and B), and we retrieve 5 documents [A, X, B, Y, Z]:
- A is at position 1: precision@1 = 1/1 = 1.0 (1 relevant out of top 1)
- B is at position 3: precision@3 = 2/3 ≈ 0.667 (2 relevant out of top 3)
- AP = (1.0 + 0.667) / 2 = 0.833
- If we had another query with AP = 0.5, then MAP = (0.833 + 0.5) / 2 = 0.667
### Success@K Results
MAP rewards systems that rank relevant documents higher, not just finding them.
| Embedding Model | Success@1 | Success@3 | Success@5 | Reranker |
|----------------------------|-----------|-----------|-----------|------------------------|
| `qwen3-embedding:8b` | 0.36 | 0.54 | 0.62 | None |
| `qwen3-embedding:8b` | 0.42 | 0.66 | 0.76 | `mxbai-rerank-base-v2` |
| `qwen3-embedding:8b` | 0.41 | 0.66 | 0.76 | `zeroentropy` |
### MAP Results
| Embedding Model | MAP | Reranker |
|----------------------------|-------|------------------------|
| - | - | - |
## QA Accuracy

219
evaluations/evaluations/benchmark.py
View file

@ -8,7 +8,7 @@ import typer
from dotenv import load_dotenv
from pydantic_ai.models.openai import OpenAIChatModel
from pydantic_ai.providers.ollama import OllamaProvider
from pydantic_evals import Dataset as EvalDataset
from pydantic_evals import Case, Dataset as EvalDataset
from pydantic_evals.evaluators import LLMJudge
from pydantic_evals.reporting import ReportCaseFailure
from rich.console import Console
@ -16,7 +16,7 @@ from rich.progress import Progress
from evaluations.config import DatasetSpec
from evaluations.datasets import DATASETS
from evaluations.llm_judge import ANSWER_EQUIVALENCE_RUBRIC
from evaluations.evaluators import ANSWER_EQUIVALENCE_RUBRIC
from evaluations.prompts import WIX_SUPPORT_PROMPT
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
@ -33,6 +33,30 @@ configure_cli_logging()
console = Console()
def build_experiment_metadata(
dataset_key: str,
test_cases: int,
config: AppConfig,
judge_model: str,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
return {
"dataset": dataset_key,
"test_cases": test_cases,
"embedder_provider": config.embeddings.provider,
"embedder_model": config.embeddings.model,
"embedder_dim": config.embeddings.vector_dim,
"chunk_size": config.processing.chunk_size,
"context_chunk_radius": config.processing.context_chunk_radius,
"rerank_provider": config.reranking.provider,
"rerank_model": config.reranking.model,
"qa_provider": config.qa.provider,
"qa_model": config.qa.model,
"judge_provider": "ollama",
"judge_model": judge_model,
}
async def populate_db(spec: DatasetSpec, config: AppConfig) -> None:
spec.db_path.parent.mkdir(parents=True, exist_ok=True)
corpus = spec.document_loader()
@ -68,108 +92,115 @@ async def populate_db(spec: DatasetSpec, config: AppConfig) -> None:
async def run_retrieval_benchmark(
spec: DatasetSpec, config: AppConfig
spec: DatasetSpec,
config: AppConfig,
limit: int | None = None,
name: str | None = None,
) -> dict[str, float] | None:
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
console.print("Skipping retrieval benchmark; no retrieval config.")
return None
corpus = spec.retrieval_loader()
if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus))))
recall_totals = {
1: 0.0,
3: 0.0,
5: 0.0,
}
success_totals = {
1: 0.0,
3: 0.0,
5: 0.0,
}
total_queries = 0
cases = []
with Progress() as progress:
task = progress.add_task(
"[blue]Running retrieval benchmark...", total=len(corpus)
)
async with HaikuRAG(spec.db_path, config=config) as rag:
for doc in corpus:
doc_mapping = cast(Mapping[str, Any], doc)
sample = spec.retrieval_mapper(doc_mapping)
if sample is None or sample.skip:
progress.advance(task)
continue
matches = await rag.search(query=sample.question, limit=5)
if not matches:
progress.advance(task)
continue
total_queries += 1
retrieved_uris: list[str] = []
for chunk, _ in matches:
if chunk.document_id is None:
continue
retrieved_doc = await rag.get_document_by_id(chunk.document_id)
if retrieved_doc and retrieved_doc.uri:
retrieved_uris.append(retrieved_doc.uri)
# Compute metrics for each cutoff
for cutoff in (1, 3, 5):
top_k = set(retrieved_uris[:cutoff])
relevant = set(sample.expected_uris)
if relevant:
matched = len(top_k & relevant)
# Recall: fraction of relevant docs retrieved
recall_totals[cutoff] += matched / len(relevant)
# Success: binary - did we get at least one relevant doc?
success_totals[cutoff] += 1.0 if matched > 0 else 0.0
task = progress.add_task("[blue]Building retrieval cases...", total=len(corpus))
for doc in corpus:
doc_mapping = cast(Mapping[str, Any], doc)
sample = spec.retrieval_mapper(doc_mapping)
if sample is None or sample.skip:
progress.advance(task)
continue
if total_queries == 0:
case = Case(
inputs=sample.question,
metadata={"relevant_uris": sample.expected_uris},
)
cases.append(case)
progress.advance(task)
if not cases:
console.print("No retrieval cases to evaluate.")
return None
recall_at_1 = recall_totals[1] / total_queries
recall_at_3 = recall_totals[3] / total_queries
recall_at_5 = recall_totals[5] / total_queries
if spec.retrieval_evaluator is None:
raise ValueError(f"No retrieval evaluator configured for dataset: {spec.key}")
success_at_1 = success_totals[1] / total_queries
success_at_3 = success_totals[3] / total_queries
success_at_5 = success_totals[5] / total_queries
evaluator = spec.retrieval_evaluator
metric_name = evaluator.__class__.__name__.replace("Evaluator", "").upper()
dataset = EvalDataset(
cases=cases,
evaluators=[evaluator],
)
async with HaikuRAG(spec.db_path, config=config) as rag:
async def retrieval_target(question: str) -> list[str]:
chunks = await rag.search(query=question, limit=5)
seen = set()
uris = []
for chunk, _ in chunks:
if chunk.document_id is None:
continue
doc = await rag.get_document_by_id(chunk.document_id)
if doc and doc.uri and doc.uri not in seen:
uris.append(doc.uri)
seen.add(doc.uri)
return uris
eval_name = name if name is not None else f"{spec.key}_retrieval_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_model=QA_JUDGE_MODEL,
)
report = await dataset.evaluate(
retrieval_target,
name=eval_name,
max_concurrency=1,
progress=True,
metadata=experiment_metadata,
)
total_score = 0.0
total_cases = 0
for case in report.cases:
if case.scores:
for score_result in case.scores.values():
total_score += score_result.value
total_cases += 1
mean_score = total_score / total_cases if total_cases > 0 else 0.0
console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
console.print(f"Total queries: {total_queries}")
console.print("\nRecall@K (fraction of relevant docs retrieved):")
console.print(f" Recall@1: {recall_at_1:.4f}")
console.print(f" Recall@3: {recall_at_3:.4f}")
console.print(f" Recall@5: {recall_at_5:.4f}")
console.print("\nSuccess@K (queries with at least one relevant doc):")
console.print(f" Success@1: {success_at_1:.4f} ({success_at_1 * 100:.1f}%)")
console.print(f" Success@3: {success_at_3:.4f} ({success_at_3 * 100:.1f}%)")
console.print(f" Success@5: {success_at_5:.4f} ({success_at_5 * 100:.1f}%)")
console.print(f"Dataset: {spec.key}")
console.print(f"Total queries: {len(cases)}")
console.print(f"{metric_name}: {mean_score:.4f}")
return {
"recall@1": recall_at_1,
"recall@3": recall_at_3,
"recall@5": recall_at_5,
"success@1": success_at_1,
"success@3": success_at_3,
"success@5": success_at_5,
metric_name.lower(): mean_score,
"queries": len(cases),
}
async def run_qa_benchmark(
spec: DatasetSpec,
config: AppConfig,
qa_limit: int | None = None,
limit: int | None = None,
name: str | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
corpus = spec.qa_loader()
if qa_limit is not None:
corpus = corpus.select(range(min(qa_limit, len(corpus))))
if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus))))
cases = [
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
@ -207,22 +238,12 @@ async def run_qa_benchmark(
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = {
"dataset": spec.key,
"test_cases": len(cases),
"embedder_provider": config.embeddings.provider,
"embedder_model": config.embeddings.model,
"embedder_dim": config.embeddings.vector_dim,
"qa_provider": config.qa.provider,
"qa_model": config.qa.model,
"judge_provider": "ollama",
"judge_model": QA_JUDGE_MODEL,
"chunk_size": config.processing.chunk_size,
"context_chunk_radius": config.processing.context_chunk_radius,
"retrieval_limit": 3,
"rerank_provider": config.reranking.provider,
"rerank_model": config.reranking.model,
}
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_model=QA_JUDGE_MODEL,
)
report = await evaluation_dataset.evaluate(
answer_question,
@ -266,7 +287,7 @@ async def evaluate_dataset(
skip_db: bool,
skip_retrieval: bool,
skip_qa: bool,
qa_limit: int | None,
limit: int | None,
name: str | None,
) -> None:
if not skip_db:
@ -275,11 +296,11 @@ async def evaluate_dataset(
if not skip_retrieval:
console.print("Running retrieval benchmarks...", style="bold blue")
await run_retrieval_benchmark(spec, config)
await run_retrieval_benchmark(spec, config, limit=limit, name=name)
if not skip_qa:
console.print("\nRunning QA benchmarks...", style="bold yellow")
await run_qa_benchmark(spec, config, qa_limit=qa_limit, name=name)
await run_qa_benchmark(spec, config, limit=limit, name=name)
app = typer.Typer(help="Run retrieval and QA benchmarks for configured datasets.")
@ -298,8 +319,8 @@ def run(
False, "--skip-retrieval", help="Skip retrieval benchmark."
),
skip_qa: bool = typer.Option(False, "--skip-qa", help="Skip QA benchmark."),
qa_limit: int | None = typer.Option(
None, "--qa-limit", help="Limit number of QA cases."
limit: int | None = typer.Option(
None, "--limit", help="Limit number of test cases for both retrieval and QA."
),
name: str | None = typer.Option(None, "--name", help="Override evaluation name."),
) -> None:
@ -335,7 +356,7 @@ def run(
skip_db=skip_db,
skip_retrieval=skip_retrieval,
skip_qa=skip_qa,
qa_limit=qa_limit,
limit=limit,
name=name,
)
)

2
evaluations/evaluations/config.py
View file

@ -5,6 +5,7 @@ from typing import Any
from datasets import Dataset
from pydantic_evals import Case
from pydantic_evals.evaluators import Evaluator
@dataclass
@ -39,6 +40,7 @@ class DatasetSpec:
qa_case_builder: CaseBuilder
retrieval_loader: RetrievalLoader | None = None
retrieval_mapper: RetrievalMapper | None = None
retrieval_evaluator: Evaluator | None = None
document_limit: int | None = None
@property

2
evaluations/evaluations/datasets/repliqa.py
View file

@ -5,6 +5,7 @@ from datasets import Dataset, DatasetDict, load_dataset
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MRREvaluator
def load_repliqa_corpus() -> Dataset:
@ -55,4 +56,5 @@ REPLIQ_SPEC = DatasetSpec(
qa_case_builder=build_repliqa_case,
retrieval_loader=load_repliqa_corpus,
retrieval_mapper=map_repliqa_retrieval,
retrieval_evaluator=MRREvaluator(),
)

2
evaluations/evaluations/datasets/wix.py
View file

@ -6,6 +6,7 @@ from datasets import Dataset, DatasetDict, load_dataset
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MAPEvaluator
def load_wix_corpus() -> Dataset:
@ -78,4 +79,5 @@ WIX_SPEC = DatasetSpec(
qa_case_builder=build_wix_case,
retrieval_loader=load_wix_qa,
retrieval_mapper=map_wix_retrieval,
retrieval_evaluator=MAPEvaluator(),
)

15
evaluations/evaluations/evaluators/__init__.py Normal file
View file

@ -0,0 +1,15 @@
from evaluations.evaluators.judge import (
ANSWER_EQUIVALENCE_RUBRIC,
LLMJudge,
LLMJudgeResponseSchema,
)
from evaluations.evaluators.map import MAPEvaluator
from evaluations.evaluators.mrr import MRREvaluator
__all__ = [
"ANSWER_EQUIVALENCE_RUBRIC",
"LLMJudge",
"LLMJudgeResponseSchema",
"MAPEvaluator",
"MRREvaluator",
]

1
evaluations/evaluations/llm_judge.py → evaluations/evaluations/evaluators/judge.py
View file

@ -5,7 +5,6 @@ from pydantic_ai.providers.ollama import OllamaProvider
from haiku.rag.config import Config
# Shared rubric/prompt for answer equivalence evaluation
ANSWER_EQUIVALENCE_RUBRIC = """You are evaluating whether two answers to the same question are semantically equivalent.
EVALUATION CRITERIA:

49
evaluations/evaluations/evaluators/map.py Normal file
View file

@ -0,0 +1,49 @@
from dataclasses import dataclass
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
@dataclass
class MAPEvaluator(Evaluator):
"""
Mean Average Precision evaluator for multi-document retrieval.
AP = (sum of P@k for each relevant doc) / total relevant docs
where P@k is precision at position k.
Appropriate for retrieval tasks where queries have multiple relevant documents.
"""
def evaluate(self, ctx: EvaluatorContext) -> float:
"""
Calculate average precision for a single query.
Expected context:
- ctx.metadata['relevant_uris']: set/list of relevant document URIs
- ctx.output: list of retrieved document URIs (ordered by rank)
Returns:
float: Average precision (0.0-1.0)
"""
if ctx.metadata is None:
return 0.0
relevant_uris = set(ctx.metadata.get("relevant_uris", []))
retrieved_uris = ctx.output
if not relevant_uris:
return 0.0
num_relevant = len(relevant_uris)
precisions = []
num_relevant_found = 0
for rank, uri in enumerate(retrieved_uris, start=1):
if uri in relevant_uris:
num_relevant_found += 1
precision_at_k = num_relevant_found / rank
precisions.append(precision_at_k)
if not precisions:
return 0.0
return sum(precisions) / num_relevant

37
evaluations/evaluations/evaluators/mrr.py Normal file
View file

@ -0,0 +1,37 @@
from dataclasses import dataclass
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
@dataclass
class MRREvaluator(Evaluator):
"""
Mean Reciprocal Rank evaluator for single-document retrieval.
MRR = 1/rank where rank is the position of the first relevant document.
Returns 0 if no relevant document is found.
Appropriate for retrieval tasks where each query has exactly one relevant document.
"""
def evaluate(self, ctx: EvaluatorContext) -> float:
"""
Calculate reciprocal rank for a single query.
Expected context:
- ctx.metadata['relevant_uris']: set/list of relevant document URIs
- ctx.output: list of retrieved document URIs (ordered by rank)
Returns:
float: 1/rank of first relevant doc, or 0.0 if not found
"""
if ctx.metadata is None:
return 0.0
relevant_uris = set(ctx.metadata.get("relevant_uris", []))
retrieved_uris = ctx.output
for rank, uri in enumerate(retrieved_uris, start=1):
if uri in relevant_uris:
return 1.0 / rank
return 0.0

2
haiku_rag_slim/haiku/rag/qa/agent.py
View file

@ -51,7 +51,7 @@ class QuestionAnswerAgent:
async def search_documents(
ctx: RunContext[Dependencies],
query: str,
limit: int = 3,
limit: int = 5,
) -> list[SearchResult]:
"""Search the knowledge base for relevant documents."""
search_results = await ctx.deps.client.search(query, limit=limit)

2
tests/test_qa.py
View file

@ -2,7 +2,7 @@ import os
import pytest
from datasets import Dataset
from evaluations.llm_judge import LLMJudge
from evaluations.evaluators import LLMJudge
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config