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

Split the evaluation benchmark by responsibility

Browse source 
benchmark.py was 1220 lines holding six unrelated jobs: populating a
database, running retrieval, running QA, resolving datasets, moving
databases to and from HuggingFace, and wiring the Typer CLI.

qa.py takes both QA runners with their live summary, refusal metrics and
target resolution. population.py takes populate_db and the batched ingest.
retrieval.py takes run_retrieval_benchmark. artifacts.py takes HF_REPO_ID and
the download/upload bodies. experiment.py takes DEFAULT_JUDGE_MODEL and
build_experiment_metadata, which retrieval and QA both record.

benchmark.py keeps the CLI at 258 lines: the Typer app, config and case-id
loading, dataset resolution, evaluate_dataset, and three commands whose
bodies are now a loop over specs. The module-level side effects stay with it
— load_dotenv before configure_telemetry, so credentials and LOGFIRE_TOKEN
are in the environment before telemetry and model setup read them — so no
importable module carries one.

Test patch targets follow the code. get_model, run_capability_question,
run_capability_conversation, set_eval_attribute and HaikuRAG are patched
inside moved code, so they move with it; run_qa_benchmark,
run_retrieval_benchmark and find_config_file stay patchable on
evaluations.benchmark because evaluate_dataset and _load_config still look
them up there.

One assertion got stronger: a QA test patched benchmark.HaikuRAG to prove the
QA path does not open its own client. qa.py has no HaikuRAG reference at all
now, so the test asserts that instead.
This commit is contained in:
Yiorgis Gozadinos 2026-08-20 14:08:09 +03:00
parent d429ac4996
commit 19d5b2e7f6
No known key found for this signature in database
7 changed files with 1084 additions and 1022 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

106
evaluations/evaluations/artifacts.py Normal file
View file

@ -0,0 +1,106 @@
"""Pre-built evaluation databases on HuggingFace."""
import os
import shutil
import tempfile
from pathlib import Path
from huggingface_hub import HfApi, snapshot_download
from rich.console import Console
from evaluations.config import DatasetSpec
console = Console()
HF_REPO_ID = "ggozad/haiku-rag-eval-dbs"
def download_dataset_db(spec: DatasetSpec, force: bool = False) -> None:
"""Fetch one dataset's database from HuggingFace into its local path."""
db = spec.db_path()
if db.exists() and not force:
console.print(
f"[yellow]Skipping {spec.key}: database already exists at {db}[/yellow]"
)
console.print("Use --force to overwrite.")
return
console.print(f"[blue]Downloading {spec.key}...[/blue]")
try:
downloaded_path = snapshot_download(
repo_id=HF_REPO_ID,
repo_type="dataset",
allow_patterns=f"{spec.db_filename}/*",
)
except Exception as e:
console.print(f"[red]Failed to download {spec.key}: {e}[/red]")
return
source_path = Path(downloaded_path) / spec.db_filename
if not source_path.exists():
console.print(f"[red]Database {spec.key} not found in HuggingFace repo.[/red]")
console.print(
f"[yellow]The database may not have been uploaded yet. "
f"Try running 'evaluations build {spec.key}' to create it locally.[/yellow]"
)
return
if db.exists():
shutil.rmtree(db)
db.parent.mkdir(parents=True, exist_ok=True)
shutil.copytree(source_path, db)
console.print(f"[green]Downloaded {spec.key} to {db}[/green]")
def upload_dataset_db(spec: DatasetSpec) -> None:
"""Push one dataset's database to HuggingFace (maintainer only).
Uses ``upload_large_folder`` for resumable, parallel transfer — important
for the multi-GB ORB databases which would otherwise abort on any transient
network failure under plain ``upload_folder``.
``upload_large_folder`` has no ``path_in_repo`` — it ships the contents of
``folder_path`` to the repo root. Stage the db under a temp parent with
hardlinks so the basename becomes the remote path, leaving everything else
at the root undisturbed.
"""
db = spec.db_path()
if not db.exists():
console.print(f"[red]Database not found at {db}[/red]")
return
api = HfApi()
# Wipe the existing remote path so we don't accumulate orphaned files from
# prior uploads. upload_large_folder doesn't accept delete_patterns, so we
# do this as a separate commit. Safe to run if the path is missing.
try:
api.delete_folder(
path_in_repo=spec.db_filename,
repo_id=HF_REPO_ID,
repo_type="dataset",
)
except Exception:
pass
with tempfile.TemporaryDirectory() as staging:
target = Path(staging) / spec.db_filename
target.mkdir()
for src in db.rglob("*"):
if not src.is_file():
continue
dest = target / src.relative_to(db)
dest.parent.mkdir(parents=True, exist_ok=True)
os.link(src, dest)
console.print(f"[blue]Uploading {spec.key} ({db})...[/blue]")
api.upload_large_folder(
folder_path=staging,
repo_id=HF_REPO_ID,
repo_type="dataset",
)
console.print(f"[green]Uploaded {spec.key} to {HF_REPO_ID}[/green]")

987
evaluations/evaluations/benchmark.py
View file

File diff suppressed because it is too large Load diff

86
evaluations/evaluations/experiment.py Normal file
View file

@ -0,0 +1,86 @@
"""Experiment metadata recorded with every eval run."""
from typing import TYPE_CHECKING, Any
from haiku.rag.config import AppConfig
from haiku.rag.config.models import ModelConfig
if TYPE_CHECKING:
from evaluations.qa import Target
# Pinned judge model. Decoupled from `config.qa.model` so a user changing
# their QA model does not inadvertently change the judge — keeps cross-run
# comparisons stable. Override per-run with `--judge-model provider:name`.
#
# Sampling follows Qwen's recommendation for thinking mode; its model cards
# forbid greedy decoding. Only the keys ollama honours are set: it silently
# ignores `top_k`, `min_p` and `chat_template_kwargs`. The vLLM reference
# configs under `evaluations/configs/` carry those too, plus
# `reasoning_effort`, which qwen3.8 reads from `chat_template_kwargs`.
DEFAULT_JUDGE_MODEL = ModelConfig(
provider="ollama",
name="qwen3.8",
temperature=0.6,
max_tokens=16384,
extra_body={"top_p": 0.95},
)
def build_experiment_metadata(
dataset_key: str,
test_cases: int,
config: AppConfig,
judge_config: ModelConfig | None = None,
target: "Target" = "rag-capability",
capability_config: ModelConfig | None = None,
document_filter: str | None = None,
) -> dict[str, Any]:
"""Build experiment metadata for Logfire tracking."""
metadata: dict[str, Any] = {
"dataset": dataset_key,
"test_cases": test_cases,
"target": target,
"embedder_provider": config.embeddings.model.provider,
"embedder_model": config.embeddings.model.name,
"embedder_dim": config.embeddings.model.vector_dim,
"chunk_size": config.processing.chunk_size,
"search_limit": config.search.limit,
"max_context_chars": config.search.max_context_chars,
"rerank_provider": config.reranking.model.provider
if config.reranking.model
else None,
"rerank_model": config.reranking.model.name if config.reranking.model else None,
"qa_provider": config.qa.model.provider,
"qa_model": config.qa.model.name,
"qa_temperature": config.qa.model.temperature,
"qa_max_tokens": config.qa.model.max_tokens,
"qa_enable_thinking": config.qa.model.enable_thinking,
"qa_extra_body": config.qa.model.extra_body,
"qa_max_searches": config.qa.max_searches,
"document_filter": document_filter,
}
if judge_config is not None:
metadata.update(
{
"judge_provider": judge_config.provider,
"judge_model": judge_config.name,
"judge_temperature": judge_config.temperature,
"judge_max_tokens": judge_config.max_tokens,
"judge_enable_thinking": judge_config.enable_thinking,
# Sampling and thinking reach vLLM through extra_body, so
# without it a trace cannot tell which judge settings ran.
"judge_extra_body": judge_config.extra_body,
}
)
if capability_config is not None:
metadata.update(
{
"capability_provider": capability_config.provider,
"capability_model": capability_config.name,
"capability_temperature": capability_config.temperature,
"capability_max_tokens": capability_config.max_tokens,
"capability_enable_thinking": capability_config.enable_thinking,
"capability_extra_body": capability_config.extra_body,
}
)
return metadata

147
evaluations/evaluations/population.py Normal file
View file

@ -0,0 +1,147 @@
"""Populating an evaluation database from a dataset spec."""
from collections.abc import Callable, Mapping
from pathlib import Path
from typing import Any, cast
from rich.console import Console
from rich.progress import Progress
from evaluations.config import DatasetSpec
from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import DocumentImport
from haiku.rag.config import AppConfig
console = Console()
async def _ingest_batched(
rag: HaikuRAG,
spec: DatasetSpec,
corpus,
batch_size: int,
on_document: Callable[[], None] = lambda: None,
) -> None:
"""Ingest inline-content documents via `import_documents` batches.
Each batch writes the documents/chunks/document_items tables once and
embeds every chunk in one batched pass. A URI is skipped on resume only
when its document has chunks; a chunkless document (crash between the
document and chunk writes) is deleted and re-imported.
"""
uri_rows = await (
rag.store.document_meta_table.query().select(["id", "uri"]).to_list()
)
chunk_rows = await rag.store.chunks_table.query().select(["document_id"]).to_list()
chunked_ids = {row["document_id"] for row in chunk_rows}
complete = {row["uri"] for row in uri_rows if row["id"] in chunked_ids}
chunkless = {
row["uri"]: row["id"] for row in uri_rows if row["id"] not in chunked_ids
}
batch: list[DocumentImport] = []
for doc in corpus:
payload = spec.document_mapper(cast(Mapping[str, Any], doc))
if payload is None or payload.uri in complete:
on_document()
continue
if payload.uri in chunkless:
await rag.delete_document(chunkless[payload.uri])
assert payload.content is not None, "batched ingest requires inline content"
docling_document = await rag.convert(payload.content, format=payload.format)
chunks = await rag.chunk(docling_document)
batch.append(
DocumentImport(
docling_document=docling_document,
chunks=chunks,
uri=payload.uri,
title=payload.title,
metadata=payload.metadata or {},
)
)
if len(batch) >= batch_size:
await rag.import_documents(batch)
batch = []
on_document()
if batch:
await rag.import_documents(batch)
async def populate_db(
spec: DatasetSpec,
config: AppConfig,
db_path: Path | None = None,
vacuum_interval: int = 100,
) -> None:
db = spec.db_path(db_path)
db.parent.mkdir(parents=True, exist_ok=True)
corpus = spec.document_loader()
if spec.document_limit is not None:
corpus = corpus.select(range(min(spec.document_limit, len(corpus))))
# Disable auto_vacuum - we'll vacuum periodically instead to prevent disk exhaustion
config.storage.auto_vacuum = False
with Progress() as progress:
task = progress.add_task("[green]Populating database...", total=len(corpus))
async with HaikuRAG(db, config=config, create=True) as rag:
if spec.ingest_batch_size is not None:
await _ingest_batched(
rag,
spec,
corpus,
batch_size=spec.ingest_batch_size,
on_document=lambda: progress.advance(task),
)
await rag.store.vacuum(retention_seconds=0)
return
docs_since_vacuum = 0
for doc in corpus:
doc_mapping = cast(Mapping[str, Any], doc)
payload = spec.document_mapper(doc_mapping)
if payload is None:
progress.advance(task)
continue
# `payload.uri` is the canonical document identifier and is now
# honored by both `create_document` and (via the `uri=` override)
# `create_document_from_source`, so it's also the right key to
# look up an existing document, regardless of whether the source
# is a file path or inline content.
existing = await rag.get_document_by_uri(payload.uri)
if existing is not None:
assert existing.id
chunks = await rag.chunk_repository.get_by_document_id(existing.id)
if chunks:
progress.advance(task)
continue
await rag.document_repository.delete(existing.id)
if payload.source_path is not None:
await rag.create_document_from_source(
source=payload.source_path,
title=payload.title,
metadata=payload.metadata,
uri=payload.uri,
)
else:
assert payload.content is not None
await rag.create_document(
content=payload.content,
uri=payload.uri,
title=payload.title,
metadata=payload.metadata,
format=payload.format,
)
docs_since_vacuum += 1
progress.advance(task)
# Periodic vacuum to prevent disk exhaustion
if docs_since_vacuum >= vacuum_interval:
await rag.store.vacuum(retention_seconds=0)
docs_since_vacuum = 0
# Final vacuum
await rag.store.vacuum(retention_seconds=0)

559
evaluations/evaluations/qa.py Normal file
View file

@ -0,0 +1,559 @@
"""QA benchmarks: single-question runs and live multi-turn conversations."""
from collections.abc import Mapping
from pathlib import Path
from typing import Any, Literal, NamedTuple, cast
from pydantic_evals import Case, Dataset as EvalDataset, set_eval_attribute
from pydantic_evals.evaluators import Evaluator
from pydantic_evals.reporting import ReportCaseFailure
from rich.console import Console
from evaluations.capability_runner import (
CapabilityFactory,
prefix_to_messages,
run_capability_conversation,
run_capability_question,
)
from evaluations.config import ConversationInput, DatasetSpec
from evaluations.evaluators import (
ANSWER_EQUIVALENCE_RUBRIC,
REFUSAL_ELIGIBLE_LABELS,
REFUSAL_RUBRIC,
ConversationEvaluator,
RefusalJudge,
TranscriptLLMJudge,
)
from evaluations.experiment import DEFAULT_JUDGE_MODEL, build_experiment_metadata
from haiku.rag.config import AppConfig
from haiku.rag.config.models import ModelConfig
from haiku.rag.utils import get_model
console = Console()
Target = Literal["rag-capability", "analysis-capability"]
TARGETS: tuple[Target, ...] = ("rag-capability", "analysis-capability")
def _capability_factory_for_target(target: Target) -> CapabilityFactory:
if target == "rag-capability":
from haiku.rag.capabilities.rag import create_capability
return create_capability
if target == "analysis-capability":
from haiku.rag.capabilities.analysis import create_capability
return create_capability
raise ValueError(f"target {target!r} is not a capability target")
def _attach_relevant_uris(
cases: list[Case[str, str, dict[str, Any]]],
spec: DatasetSpec,
limit: int | None,
) -> None:
"""Augment QA cases with `relevant_uris` joined from retrieval samples.
Mutates each case's metadata in place. Cases with no matching retrieval
sample (by question) are left untouched.
"""
if spec.retrieval_loader is None or spec.retrieval_mapper is None:
return
if not any(isinstance(case.inputs, str) for case in cases):
return
corpus = spec.retrieval_loader()
if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus))))
expected_by_question: dict[str, tuple[str, ...]] = {}
for raw in corpus:
sample = spec.retrieval_mapper(cast(Mapping[str, Any], raw))
if sample is None or sample.skip:
continue
expected_by_question[sample.question] = sample.expected_uris
for case in cases:
if not isinstance(case.inputs, str):
continue
uris = expected_by_question.get(case.inputs)
if uris is None:
continue
metadata = case.metadata if case.metadata is not None else {}
metadata["relevant_uris"] = list(uris)
case.metadata = metadata
def _resolve_capability_config(
target: Target, config: AppConfig, capability_model: ModelConfig | None
) -> ModelConfig:
if target == "analysis-capability":
# Mirror the capability-code resolver: explicit analysis.model wins,
# else fall back to qa.model.
return capability_model or config.analysis.model or config.qa.model
return capability_model or config.qa.model
def _live_summary(report_cases, report_failures) -> dict[str, float | int] | None:
"""Aggregate ConversationEvaluator scores across conversations.
Micro rates weight every turn equally (sums across conversations); macro
rates average per-conversation means, so short conversations don't get
overweighted by micro nor long ones by macro. Failed conversations are
operational exclusions: they count toward the attempted coverage figures
but never toward the rates.
"""
def _score(case, key: str):
result = case.scores.get(key)
return result.value if result is not None else None
scored = [case for case in report_cases if _score(case, "turns_total") is not None]
if not scored:
return None
failed_turns = sum(
len(failure.inputs) if isinstance(failure.inputs, list) else 0
for failure in report_failures
)
turns_total = sum(_score(case, "turns_total") for case in scored)
turns_judged = sum(_score(case, "turns_judged") for case in scored)
turns_passed = sum(_score(case, "turns_passed") for case in scored)
# A conversation with zero judged turns (its judge calls all failed)
# reports turn_pass_rate 0.0; averaging that in would count a judge
# outage as a failed conversation, against the exclusion policy.
judged = [case for case in scored if _score(case, "turns_judged")]
summary: dict[str, float | int] = {
"conversations": len(scored),
"conversations_attempted": len(report_cases) + len(report_failures),
"turns_total": turns_total,
"turns_judged": turns_judged,
"turns_attempted": turns_total + failed_turns,
"micro_pass_rate": turns_passed / turns_judged if turns_judged else 0.0,
"macro_pass_rate": sum(_score(case, "turn_pass_rate") for case in judged)
/ len(judged)
if judged
else 0.0,
}
cited = [case for case in scored if _score(case, "cited_map") is not None]
eligible = sum(_score(case, "cited_eligible") for case in scored)
if cited and eligible:
summary["cited_eligible"] = eligible
summary["cited_map_micro"] = (
sum(
_score(case, "cited_map") * _score(case, "cited_eligible")
for case in cited
)
/ eligible
)
summary["cited_map_macro"] = sum(
_score(case, "cited_map") for case in cited
) / len(cited)
true_refusals = sum(_score(case, "true_refusals") for case in scored)
false_refusals = sum(_score(case, "false_refusals") for case in scored)
unanswerable = sum(_score(case, "unanswerable_turns") for case in scored)
refusals = true_refusals + false_refusals
summary["unanswerable_turns"] = unanswerable
summary["refusals"] = refusals
summary["refusal_precision"] = true_refusals / refusals if refusals else 0.0
summary["refusal_recall"] = true_refusals / unanswerable if unanswerable else 0.0
return summary
def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None:
"""Refusal precision/recall against answerability labels.
Uses cases the refusal judge scored (ANSWERABLE/UNANSWERABLE turns).
Returns (precision, recall, unanswerable_count, refusal_count), or None
when no case was judged.
"""
outcomes: list[tuple[str, bool]] = []
for case in report_cases:
refused = case.assertions.get("refused")
label = (case.metadata or {}).get("answerability")
if refused is None or label not in REFUSAL_ELIGIBLE_LABELS:
continue
outcomes.append((label, bool(refused.value)))
if not outcomes:
return None
refusals = [(label, r) for label, r in outcomes if r]
true_refusals = sum(1 for label, _ in refusals if label == "UNANSWERABLE")
unanswerable = sum(1 for label, _ in outcomes if label == "UNANSWERABLE")
precision = true_refusals / len(refusals) if refusals else 0.0
recall = true_refusals / unanswerable if unanswerable else 0.0
return precision, recall, unanswerable, len(refusals)
def _filter_qa_corpus(corpus, case_ids: set[str] | None):
"""Keep only rows whose ``id`` is in ``case_ids`` (failure-subset reruns).
Returns the corpus unchanged when ``case_ids`` is None.
"""
if case_ids is None:
return corpus
return corpus.filter(lambda row: row.get("id") in case_ids)
class _QARun(NamedTuple):
cases: list[Case[Any, Any, dict[str, Any]]]
db: Path
judge_config: ModelConfig
eval_name: str
experiment_metadata: dict[str, Any]
capability_factory: CapabilityFactory
capability_model: Any
def _prepare_qa_run(
spec: DatasetSpec,
config: AppConfig,
limit: int | None,
name: str | None,
db_path: Path | None,
judge_model: ModelConfig | None,
target: Target,
capability_model: ModelConfig | None,
case_ids: set[str] | None,
document_filter: str | None,
) -> _QARun:
"""Shared setup for the QA runners: cases, models, name and metadata."""
corpus = spec.qa_loader()
corpus = _filter_qa_corpus(corpus, case_ids)
if limit is not None:
corpus = corpus.select(range(min(limit, len(corpus))))
cases = [
spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
for index, doc in enumerate(corpus, start=1)
]
judge_config = judge_model or DEFAULT_JUDGE_MODEL
capability_config = _resolve_capability_config(target, config, capability_model)
eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
experiment_metadata = build_experiment_metadata(
dataset_key=spec.key,
test_cases=len(cases),
config=config,
judge_config=judge_config,
target=target,
capability_config=capability_config,
document_filter=document_filter,
)
experiment_metadata.update(spec.experiment_metadata or {})
return _QARun(
cases=cases,
db=spec.db_path(db_path),
judge_config=judge_config,
eval_name=eval_name,
experiment_metadata=experiment_metadata,
capability_factory=_capability_factory_for_target(target),
capability_model=get_model(capability_config, config),
)
def _print_mean_task_time(report_cases, unit: str = "case") -> None:
if not report_cases:
return
mean = sum(case.task_duration for case in report_cases) / len(report_cases)
console.print(f"Avg task time per {unit}: {mean:.2f}s")
def _print_failures(failures, show_question: bool = False) -> None:
if not failures:
return
console.print("[red]\nSummary of failures:[/red]")
for failure in failures:
console.print(f"Case: {failure.name}")
if show_question:
console.print(f"Question: {failure.inputs}")
console.print(f"Error: {failure.error_message}")
console.print("")
async def run_qa_benchmark(
spec: DatasetSpec,
config: AppConfig,
limit: int | None = None,
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "rag-capability",
capability_model: ModelConfig | None = None,
case_ids: set[str] | None = None,
document_filter: str | None = None,
) -> ReportCaseFailure[str, str, dict[str, str]] | None:
run = _prepare_qa_run(
spec,
config,
limit,
name,
db_path,
judge_model,
target,
capability_model,
case_ids,
document_filter,
)
cases, judge_config = run.cases, run.judge_config
_attach_relevant_uris(cases, spec, limit)
citation_evaluator = spec.citation_evaluator
qa_evaluator = spec.qa_evaluator
evaluators: list[Evaluator]
if qa_evaluator is not None:
evaluators = [qa_evaluator]
else:
evaluators = [
TranscriptLLMJudge(
rubric=ANSWER_EQUIVALENCE_RUBRIC,
include_input=True,
include_expected_output=True,
model=get_model(judge_config, config),
assertion={
"evaluation_name": "answer_equivalent",
"include_reason": True,
},
),
]
if citation_evaluator is not None:
evaluators.append(citation_evaluator)
# RefusalJudge scores only cases whose metadata carries an answerability
# label; on unlabeled datasets it returns no score without a judge call.
evaluators.append(
RefusalJudge(
rubric=REFUSAL_RUBRIC,
model=get_model(judge_config, config),
assertion={"evaluation_name": "refused", "include_reason": False},
)
)
evaluation_dataset = EvalDataset[Any, str, dict[str, Any]](
name=spec.key, cases=cases, evaluators=evaluators
)
async def answer_question(inputs: str | ConversationInput) -> str:
if isinstance(inputs, ConversationInput):
question = inputs.question
message_history = prefix_to_messages(inputs.prefix)
else:
question = inputs
message_history = None
result = await run_capability_question(
capability_factory=run.capability_factory,
db_path=run.db,
config=config,
question=question,
capability_model=run.capability_model,
document_filter=document_filter,
message_history=message_history,
)
set_eval_attribute("cited_uris", result.cited_uris)
set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids)
set_eval_attribute("searched_uris", result.searched_uris)
set_eval_attribute("n_searches", result.n_searches)
set_eval_attribute("n_search_calls", result.n_search_calls)
set_eval_attribute("n_rejected_searches", result.n_rejected_searches)
set_eval_attribute("n_failed_tools", result.n_failed_tools)
set_eval_attribute("n_executions", result.n_executions)
set_eval_attribute("n_requests", result.n_requests)
set_eval_attribute("citation_status", result.citation_status)
return result.answer
report = await evaluation_dataset.evaluate(
answer_question,
name=run.eval_name,
max_concurrency=1,
progress=True,
metadata=run.experiment_metadata,
)
total_processed = len(report.cases)
failures = report.failures
if qa_evaluator is not None:
score_key = qa_evaluator.get_default_evaluation_name()
passing_cases = sum(
1
for case in report.cases
if score_key in case.scores and case.scores[score_key].value >= 1.0
)
scoring = score_key
else:
passing_cases = sum(
1
for case in report.cases
if case.assertions.get("answer_equivalent")
and case.assertions["answer_equivalent"].value
)
scoring = "answer_equivalent"
accuracy = passing_cases / total_processed if total_processed > 0 else 0
console.print("\n=== QA Benchmark Results ===", style="bold cyan")
console.print(f"Scoring: {scoring}")
console.print(f"Total questions: {total_processed}")
console.print(f"Correct answers: {passing_cases}")
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
_print_mean_task_time(report.cases)
if citation_evaluator is not None:
score_key = citation_evaluator.get_default_evaluation_name()
scores = [
case.scores[score_key].value
for case in report.cases
if score_key in case.scores
]
if scores:
cited_count = sum(
1 for case in report.cases if case.attributes.get("cited_uris")
)
mean_citations = sum(
len(case.attributes.get("cited_uris") or []) for case in report.cases
) / len(report.cases)
mean_score = sum(scores) / len(scores)
console.print(
f"\n=== Citation Retrieval ({score_key}) ===", style="bold cyan"
)
console.print(f"Mean {score_key}: {mean_score:.4f}")
console.print(
f"Eligible cases (gold passages known): {len(scores)}/{len(report.cases)}"
)
console.print(
f"Cite rate (≥1 citation): {cited_count / len(report.cases):.2%}"
)
console.print(f"Mean citations per case: {mean_citations:.2f}")
if (metrics := _refusal_metrics(report.cases)) is not None:
precision, recall, unanswerable, refusals = metrics
console.print("\n=== Refusal vs answerability labels ===", style="bold cyan")
console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}")
console.print(
f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} "
"(PARTIAL excluded)"
)
_print_failures(failures, show_question=True)
return failures[0] if failures else None
async def run_live_qa_benchmark(
spec: DatasetSpec,
config: AppConfig,
limit: int | None = None,
name: str | None = None,
db_path: Path | None = None,
judge_model: ModelConfig | None = None,
target: Target = "rag-capability",
capability_model: ModelConfig | None = None,
case_ids: set[str] | None = None,
document_filter: str | None = None,
) -> None:
"""Replay conversations turn by turn through one capability session.
One case per conversation; ``limit`` counts conversations. Answers carry
forward as real message history, so prior-turn compaction is exercised.
"""
run = _prepare_qa_run(
spec,
config,
limit,
name,
db_path,
judge_model,
target,
capability_model,
case_ids,
document_filter,
)
evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]](
name=spec.key,
cases=run.cases,
evaluators=[
ConversationEvaluator(
rubric=ANSWER_EQUIVALENCE_RUBRIC,
model=get_model(run.judge_config, config),
)
],
)
async def answer_conversation(questions: list[str]) -> list[str]:
results = await run_capability_conversation(
capability_factory=run.capability_factory,
db_path=run.db,
config=config,
questions=list(questions),
capability_model=run.capability_model,
document_filter=document_filter,
compaction=spec.compaction,
)
set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results])
set_eval_attribute("turn_n_search_calls", [r.n_search_calls for r in results])
set_eval_attribute(
"turn_n_rejected_searches", [r.n_rejected_searches for r in results]
)
set_eval_attribute("turn_n_failed_tools", [r.n_failed_tools for r in results])
set_eval_attribute("turn_n_requests", [r.n_requests for r in results])
set_eval_attribute("turn_citation_status", [r.citation_status for r in results])
return [r.answer for r in results]
report = await evaluation_dataset.evaluate(
answer_conversation,
name=run.eval_name,
max_concurrency=1,
progress=True,
metadata=run.experiment_metadata,
)
summary = _live_summary(report.cases, report.failures)
console.print("\n=== Live Conversation Results ===", style="bold cyan")
if summary is None:
attempted = len(report.cases) + len(report.failures)
console.print(f"No conversations were scored ({attempted} attempted).")
else:
console.print(
f"Conversations scored: {summary['conversations']}"
f"/{summary['conversations_attempted']} | turns scored: "
f"{summary['turns_total']}/{summary['turns_attempted']}"
)
if summary["turns_judged"] < summary["turns_total"]:
console.print(
f"Turns judged: {summary['turns_judged']}/{summary['turns_total']} "
"(per-turn judge errors excluded from rates)"
)
if report.failures:
console.print(
"Failed conversations are operational exclusions — "
"not counted as wrong answers."
)
console.print(
f"Answer pass rate — micro (per turn): {summary['micro_pass_rate']:.4f} | "
f"macro (per conversation): {summary['macro_pass_rate']:.4f}"
)
if "cited_map_micro" in summary:
console.print(
f"cited_map — micro: {summary['cited_map_micro']:.4f} | "
f"macro: {summary['cited_map_macro']:.4f} "
f"(eligible turns: {summary['cited_eligible']})"
)
console.print(
f"Refusal precision: {summary['refusal_precision']:.2%} | "
f"recall: {summary['refusal_recall']:.2%} "
f"(UNANSWERABLE turns: {summary['unanswerable_turns']}, "
f"refusals: {summary['refusals']})"
)
if report.cases:
mean_task_time = sum(case.task_duration for case in report.cases) / len(
report.cases
)
turns = sum(len(case.output or []) for case in report.cases)
per_turn = (
sum(case.task_duration for case in report.cases) / turns if turns else 0.0
)
console.print(
f"Avg task time: {mean_task_time:.2f}s per conversation | "
f"{per_turn:.2f}s per turn"
)
_print_failures(report.failures)

128
evaluations/evaluations/retrieval.py Normal file
View file

@ -0,0 +1,128 @@
"""Retrieval benchmark: search the corpus and score the ranking."""
from collections.abc import Mapping
from pathlib import Path
from typing import Any, cast
from pydantic_evals import Case, Dataset as EvalDataset
from rich.console import Console
from rich.progress import Progress
from evaluations.config import DatasetSpec
from evaluations.experiment import build_experiment_metadata
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
console = Console()
async def run_retrieval_benchmark(
spec: DatasetSpec,
config: AppConfig,
limit: int | None = None,
name: str | None = None,
db_path: Path | None = None,
multimodal_only: bool = False,
document_filter: 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))))
cases = []
with Progress() as progress:
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
# Filter for multimodal queries if requested
if multimodal_only:
if sample.source_type is None or "image" not in sample.source_type:
progress.advance(task)
continue
case = Case(
inputs=sample.question,
metadata={
"relevant_uris": sample.expected_uris,
"source_type": sample.source_type,
},
)
cases.append(case)
progress.advance(task)
if not cases:
console.print("No retrieval cases to evaluate.")
return None
if not spec.retrieval_evaluators:
raise ValueError(f"No retrieval evaluators configured for dataset: {spec.key}")
dataset = EvalDataset(
name=f"{spec.key}-retrieval",
cases=cases,
evaluators=list(spec.retrieval_evaluators),
)
db = spec.db_path(db_path)
async with HaikuRAG(db, config=config, read_only=True) as rag:
async def retrieval_target(question: str) -> list[str]:
chunks = await rag.search(
query=question,
limit=spec.retrieval_limit,
include_images=False,
filter=document_filter,
)
seen = set()
identifiers = []
for result in chunks:
uri = result.document_uri
if uri and uri not in seen:
identifiers.append(uri)
seen.add(uri)
return identifiers
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,
document_filter=document_filter,
)
report = await dataset.evaluate(
retrieval_target,
name=eval_name,
max_concurrency=1,
progress=True,
metadata=experiment_metadata,
)
per_metric: dict[str, list[float]] = {}
for case in report.cases:
for key, score_result in case.scores.items():
per_metric.setdefault(key, []).append(score_result.value)
console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
console.print(f"Dataset: {spec.key}")
console.print(f"Total queries: {len(cases)}")
results: dict[str, float] = {"queries": len(cases)}
for key, values in per_metric.items():
mean_score = sum(values) / len(values)
metric_name = key.replace("Evaluator", "").upper()
console.print(f"{metric_name}: {mean_score:.4f}")
results[metric_name.lower()] = mean_score
return results

93
evaluations/tests/test_benchmark.py
View file

@ -7,10 +7,10 @@ import typer
from evaluations.benchmark import (
_load_config,
_resolve_dataset,
build_experiment_metadata,
evaluate_dataset,
run_qa_benchmark,
)
from evaluations.experiment import build_experiment_metadata
from evaluations.qa import run_qa_benchmark
from evaluations.config import DatasetSpec, DocumentPayload
from haiku.rag.config.models import AppConfig, ModelConfig
@ -171,9 +171,9 @@ class TestConversationInputDispatch:
)
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch("evaluations.qa.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.run_capability_question",
"evaluations.qa.run_capability_question",
new_callable=AsyncMock,
return_value=CapabilityRunResult(answer="answer"),
) as run_question,
@ -217,13 +217,13 @@ class TestConversationInputDispatch:
recorded: dict[str, object] = {}
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch("evaluations.qa.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.set_eval_attribute",
"evaluations.qa.set_eval_attribute",
side_effect=lambda key, value: recorded.__setitem__(key, value),
),
patch(
"evaluations.benchmark.run_capability_question",
"evaluations.qa.run_capability_question",
new_callable=AsyncMock,
return_value=CapabilityRunResult(
answer="answer", citation_status="ungrounded"
@ -245,7 +245,7 @@ class TestRefusalMetrics:
return case
def test_precision_and_recall(self) -> None:
from evaluations.benchmark import _refusal_metrics
from evaluations.qa import _refusal_metrics
cases = [
self._case("UNANSWERABLE", True), # true refusal
@ -266,7 +266,7 @@ class TestRefusalMetrics:
assert refusals == 2
def test_none_when_no_judged_cases(self) -> None:
from evaluations.benchmark import _refusal_metrics
from evaluations.qa import _refusal_metrics
assert _refusal_metrics([self._case("PARTIAL", None)]) is None
@ -278,7 +278,7 @@ class TestLiveSummary:
return case
def test_micro_and_macro_aggregation(self) -> None:
from evaluations.benchmark import _live_summary
from evaluations.qa import _live_summary
# Conversation A: 1/4 turns pass; B: 2/2 pass. Micro weights turns
# (3/6); macro averages conversations ((0.25 + 1.0) / 2).
@ -330,12 +330,12 @@ class TestLiveSummary:
assert summary["refusal_recall"] == pytest.approx(0.5)
def test_none_without_scored_cases(self) -> None:
from evaluations.benchmark import _live_summary
from evaluations.qa import _live_summary
assert _live_summary([self._case({})], []) is None
def test_micro_rate_uses_judged_turns(self) -> None:
from evaluations.benchmark import _live_summary
from evaluations.qa import _live_summary
cases = [
self._case(
@ -363,7 +363,7 @@ class TestLiveSummary:
"""A conversation whose every turn lost its judge reports
turn_pass_rate 0.0; treating that as a failed conversation would
contradict the exclusion policy. It must not enter the macro average."""
from evaluations.benchmark import _live_summary
from evaluations.qa import _live_summary
cases = [
self._case(
@ -401,7 +401,7 @@ class TestLiveSummary:
assert summary["turns_total"] == 10
def test_failed_conversations_do_not_affect_rates(self) -> None:
from evaluations.benchmark import _live_summary
from evaluations.qa import _live_summary
cases = [
self._case(
@ -466,9 +466,9 @@ class TestLiveConversationDispatch:
CapabilityRunResult(answer="a2", cited_uris=[]),
]
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch("evaluations.qa.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.run_capability_conversation",
"evaluations.qa.run_capability_conversation",
new_callable=AsyncMock,
return_value=turn_results,
) as run_conversation,
@ -545,13 +545,13 @@ class TestLiveConversationDispatch:
recorded: dict[str, object] = {}
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch("evaluations.qa.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.set_eval_attribute",
"evaluations.qa.set_eval_attribute",
side_effect=lambda key, value: recorded.__setitem__(key, value),
),
patch(
"evaluations.benchmark.run_capability_conversation",
"evaluations.qa.run_capability_conversation",
new_callable=AsyncMock,
return_value=turn_results,
),
@ -639,10 +639,8 @@ class TestRunQaBenchmarkJudgeModel:
custom_judge = ModelConfig(provider="openai", name="gpt-4o")
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch(
"evaluations.benchmark.run_capability_question", new_callable=AsyncMock
),
patch("evaluations.qa.get_model") as mock_get_model,
patch("evaluations.qa.run_capability_question", new_callable=AsyncMock),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -656,13 +654,11 @@ class TestRunQaBenchmarkJudgeModel:
@pytest.mark.asyncio
async def test_defaults_to_pinned_judge_model(self, tmp_path: Path) -> None:
from evaluations.benchmark import DEFAULT_JUDGE_MODEL
from evaluations.experiment import DEFAULT_JUDGE_MODEL
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch(
"evaluations.benchmark.run_capability_question", new_callable=AsyncMock
),
patch("evaluations.qa.get_model") as mock_get_model,
patch("evaluations.qa.run_capability_question", new_callable=AsyncMock),
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -674,7 +670,7 @@ class TestRunQaBenchmarkJudgeModel:
mock_get_model.assert_any_call(DEFAULT_JUDGE_MODEL, AppConfig())
def test_pinned_judge_avoids_greedy_decoding(self) -> None:
from evaluations.benchmark import DEFAULT_JUDGE_MODEL
from evaluations.experiment import DEFAULT_JUDGE_MODEL
assert DEFAULT_JUDGE_MODEL.temperature == 0.6
assert DEFAULT_JUDGE_MODEL.name == "qwen3.8"
@ -815,11 +811,10 @@ class TestRunQaBenchmarkCapabilityTarget:
return_value=CapabilityRunResult(answer="from capability")
)
with (
patch("evaluations.benchmark.get_model") as mock_get_model,
patch("evaluations.qa.get_model") as mock_get_model,
patch(
"evaluations.benchmark.run_capability_question", new=capability_run
"evaluations.qa.run_capability_question", new=capability_run
) as mock_run_capability,
patch("evaluations.benchmark.HaikuRAG") as mock_haiku,
):
mock_get_model.return_value = "fake-model"
await run_qa_benchmark(
@ -829,9 +824,11 @@ class TestRunQaBenchmarkCapabilityTarget:
target="rag-capability",
)
# When target is rag-capability, HaikuRAG context manager is NOT entered
# (the capability manages its own client via lifespan).
mock_haiku.assert_not_called()
# The capability manages its own client, so the QA runner never opens
# one — it has no HaikuRAG reference to open.
import evaluations.qa as qa_module
assert not hasattr(qa_module, "HaikuRAG")
# capability model defaults to qa.model when not provided
assert any(
call[0][0] == AppConfig().qa.model for call in mock_get_model.call_args_list
@ -842,7 +839,7 @@ class TestRunQaBenchmarkCapabilityTarget:
async def test_analysis_capability_target_resolves_factory(
self, tmp_path: Path
) -> None:
from evaluations.benchmark import _capability_factory_for_target
from evaluations.qa import _capability_factory_for_target
from haiku.rag.capabilities.analysis import (
create_capability as analysis_factory,
)
@ -913,7 +910,7 @@ class TestBatchedIngest:
@pytest.mark.asyncio
async def test_imports_in_bounded_batches(self) -> None:
from evaluations.benchmark import _ingest_batched
from evaluations.population import _ingest_batched
rag = self._rag()
corpus = [{"uri": f"u{i}"} for i in range(5)]
@ -928,7 +925,7 @@ class TestBatchedIngest:
@pytest.mark.asyncio
async def test_resume_skips_complete_uris(self) -> None:
from evaluations.benchmark import _ingest_batched
from evaluations.population import _ingest_batched
rag = self._rag(complete_uris=["u0", "u2"])
corpus = [{"uri": f"u{i}"} for i in range(4)]
@ -944,7 +941,7 @@ class TestBatchedIngest:
async def test_resume_reimports_chunkless_documents(self) -> None:
"""A crash between the document and chunk writes leaves a document
without chunks; resume must delete and re-import it, not skip it."""
from evaluations.benchmark import _ingest_batched
from evaluations.population import _ingest_batched
rag = self._rag(complete_uris=["u0"], chunkless_uris=["u1"])
corpus = [{"uri": "u0"}, {"uri": "u1"}]
@ -957,7 +954,7 @@ class TestBatchedIngest:
@pytest.mark.asyncio
async def test_unmapped_documents_skipped(self) -> None:
from evaluations.benchmark import _ingest_batched
from evaluations.population import _ingest_batched
rag = self._rag()
corpus = [{"uri": "u0"}, {"uri": "bad"}, {"uri": "u1"}]
@ -972,7 +969,7 @@ class TestAttachRelevantUris:
def test_joins_by_question(self) -> None:
from pydantic_evals import Case
from evaluations.benchmark import _attach_relevant_uris
from evaluations.qa import _attach_relevant_uris
from evaluations.config import RetrievalSample
from evaluations.evaluators import MAPEvaluator
@ -1021,7 +1018,7 @@ class TestAttachRelevantUris:
def test_no_op_without_retrieval_loader(self) -> None:
from pydantic_evals import Case
from evaluations.benchmark import _attach_relevant_uris
from evaluations.qa import _attach_relevant_uris
cases: list[Case[str, str, dict]] = [
Case(name="c1", inputs="q", expected_output="a"),
@ -1042,7 +1039,7 @@ class TestFilterQaCorpus:
def test_keeps_only_matching_ids(self) -> None:
from datasets import Dataset
from evaluations.benchmark import _filter_qa_corpus
from evaluations.qa import _filter_qa_corpus
corpus = Dataset.from_list(
[{"id": "a", "q": 1}, {"id": "b", "q": 2}, {"id": "c", "q": 3}]
@ -1053,7 +1050,7 @@ class TestFilterQaCorpus:
def test_none_returns_corpus_unchanged(self) -> None:
from datasets import Dataset
from evaluations.benchmark import _filter_qa_corpus
from evaluations.qa import _filter_qa_corpus
corpus = Dataset.from_list([{"id": "a"}])
assert _filter_qa_corpus(corpus, None) is corpus
@ -1114,7 +1111,7 @@ class TestRetrievalTarget:
)
fake = FakeRag()
with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
with patch("evaluations.retrieval.HaikuRAG") as mock_haiku:
mock_haiku.return_value.__aenter__.return_value = fake
result = await run_retrieval_benchmark(
self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb"
@ -1142,7 +1139,7 @@ class TestRetrievalTarget:
_result("uri-x", 0.7),
]
with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
with patch("evaluations.retrieval.HaikuRAG") as mock_haiku:
mock_haiku.return_value.__aenter__.return_value = FakeRag()
result = await run_retrieval_benchmark(
self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb"
@ -1195,7 +1192,7 @@ class TestDocumentFilterThreading:
retrieval_evaluators=[MAPEvaluator()],
)
with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
with patch("evaluations.retrieval.HaikuRAG") as mock_haiku:
mock_haiku.return_value.__aenter__.return_value = FakeRag()
await run_retrieval_benchmark(
spec,
@ -1225,7 +1222,7 @@ class TestDocumentFilterThreading:
)
with patch(
"evaluations.benchmark.run_capability_question",
"evaluations.qa.run_capability_question",
new_callable=AsyncMock,
return_value=CapabilityRunResult(answer="ANSWER: 42"),
) as mock_run: