Add MTRAG ClapNQ multi-turn evaluation

IBM's MTRAG benchmark (ClapNQ domain, pinned repo SHA): retrieval with
Recall@k/nDCG@k against binary qrels, gold-prefix QA replaying reference
conversation prefixes as message history, and live-session replay
carrying the model's own answers and tool history across turns.

Corpus population gains a bounded, resumable batched ingest path.
ConversationInput case type with transcript rendering for the judge,
eligibility-aware citation scoring, refusal precision/recall via a
label-aware RefusalJudge, per-turn verdicts with judged-turn coverage,
and per-turn tool-traffic attributes counted from each turn's new
messages so the arrays survive prior-turn compaction.
This commit is contained in:
Yiorgis Gozadinos 2026-08-06 16:42:54 +03:00
parent cc04f92f28
commit 73d9d93db9
No known key found for this signature in database
25 changed files with 2596 additions and 116 deletions

View file

@ -4,6 +4,7 @@
### Added
- `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata.
- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, and per-turn tool-traffic attributes.
### Changed

View file

@ -1,6 +1,6 @@
# Benchmarks
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, and HotpotQA are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the RAG and analysis capabilities.
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, HotpotQA, and MTRAG are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the RAG and analysis capabilities.
## Running Evaluations
@ -37,6 +37,7 @@ Active datasets:
| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB |
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB |
| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB |
| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite` and `mtrag_clapnq_live` keys | ~2.8 GB |
After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches):
@ -225,3 +226,29 @@ The reranker's contribution is larger here than on the single-doc datasets: hybr
| `vllm:Gemma-4-26B-A4B-NVFP4` | none | 0.83 | 0.75 |
*Measured on haiku.rag v0.66.0 with `qwen3-embedding:4b` (vLLM, dim 2560), judged by `vllm:Qwen3.6-35B-A3B-NVFP4`, 7,405 cases. The reranker lifts QA accuracy +2.7pts and `cited_map` +4.6pts. Without a reranker, `cited_map` (0.75) still exceeds the no-reranker retrieval MAP (0.70): the skill reformulates queries across search calls, partially recovering second-hop documents that a single query misses.*
### MTRAG (ClapNQ)
[MTRAG](https://github.com/IBM/mt-rag-benchmark) is IBM's multi-turn RAG benchmark (TACL 2025, SemEval-2026 Task 8): human-authored conversations with per-turn answerability labels and binary relevance judgments. We evaluate the ClapNQ (Wikipedia) domain: 183,408 passages, 29 conversations, 224 turns, 208 retrieval queries.
Three dataset keys share one database. `mtrag_clapnq` retrieves with the raw last user turn and runs QA by replaying each task's reference conversation prefix as message history. `mtrag_clapnq_rewrite` retrieves with the human standalone rewrites. `mtrag_clapnq_live` replays whole conversations through a single capability session, carrying the model's own answers and tool history across turns.
##### Retrieval (Recall@k / nDCG@k)
Directly comparable with [IBM's published results](https://github.com/IBM/mt-rag-benchmark/tree/main/mtrag-human/retrieval_tasks). Elser is IBM's strongest reported retriever.
| Retriever | Queries | R@5 | R@10 | nDCG@5 | nDCG@10 |
|-----------|---------|----:|-----:|-------:|--------:|
| Elser (IBM) | lastturn | 0.49 | 0.58 | 0.45 | 0.49 |
| `haiku.rag` | lastturn | 0.501 | 0.600 | 0.455 | 0.497 |
| Elser (IBM) | rewrite | 0.52 | 0.64 | 0.48 | 0.54 |
| `haiku.rag` | rewrite | 0.548 | 0.668 | 0.503 | 0.556 |
##### QA accuracy + citation retrieval
| Mode | Capability model | Turns | QA accuracy | Mean `cited_map` |
|------|------------------|------:|-------------|------------------|
| Gold-prefix (`mtrag_clapnq`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 223 | 0.68 | 0.35 |
| Live (`mtrag_clapnq_live`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 195/224 scored | 0.72 micro / 0.73 macro | 0.35 |
*Measured on haiku.rag v0.67.1 with `qwen3-embedding:4b` (vLLM, dim 2560) and `Qwen3-Reranker-4B`, stock capability instructions, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` at temperature 0. The judge sampling has since been re-pinned repo-wide (0.6 with thinking), so future runs re-baseline. QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Live mode additionally reports refusal precision/recall against the per-turn answerability labels and per-turn pass rates; pass rate declines with conversation depth (93% at turn 1 to 38% at turn 9). The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.*

View file

@ -9,6 +9,7 @@ This package is not published to PyPI and is only used for development and testi
Contains evaluation scripts for benchmarking RAG retrieval and QA performance. Available datasets:
- HotpotQA (`hotpotqa`) — multi-hop QA over Wikipedia paragraphs (distractor validation split, 7,405 questions, two gold documents per question)
- MTRAG ClapNQ (`mtrag_clapnq`, `mtrag_clapnq_rewrite`) — IBM's multi-turn RAG benchmark, ClapNQ (Wikipedia) domain: 183,408 passages, 208 retrieval queries with binary qrels, 224 generation tasks. The base key retrieves with the raw last user turn; the `_rewrite` variant uses the human standalone rewrites (both share one database). Retrieval reports Recall@5/@10, nDCG@5/@10, and MAP against IBM's published setup. QA replays each task's reference conversation prefix as message history and answers the final turn; the judge sees the conversation as a transcript, citation MAP is scored only on turns with gold passages, and refusal precision/recall is reported against the answerability labels. Generation scores are internal (our judge and rubric), not comparable with IBM's published generation numbers. The `mtrag_clapnq_live` key replays whole conversations (one case per conversation, `--limit` counts conversations) through a single capability session, carrying the model's own answers and tool history across turns; it reports the same outcomes per turn plus micro (per-turn) and macro (per-conversation) aggregates.
- OpenRAG Bench, two variants:
- `orb_text` — text embedder (`qwen3-embedding:4b`, 2560-dim) with VLM picture descriptions baked into chunk content at ingest. Use for text-only retrieval/QA against figure-rich corpora.
- `orb_multimodal` — multimodal embedder (`qwen3-vl-embedding-8b`, 4096-dim) with picture vectors in the same space as text. Use for cross-modal retrieval (text-as-query → figure hits, image-as-query) and vision QA where the figure itself is the answer.

View file

@ -0,0 +1,48 @@
# Reference config for the `mtrag_clapnq` pre-built evaluation database.
# IBM MTRAG, ClapNQ (Wikipedia) domain: multi-turn retrieval and QA over
# 183,408 passages. Also serves mtrag_clapnq_rewrite and mtrag_clapnq_live.
# Run: evaluations run mtrag_clapnq --config configs/mtrag_clapnq.yaml
# base_url uses the `vllm` host serving each model over an OpenAI-compatible API.
# The corpus is text-only: no multimodal embedder, no vision paths. This eval
# cannot exercise image or vision turn-boundary behavior.
environment: development
storage:
auto_vacuum: false
embeddings:
model:
provider: openai
name: qwen3-embedding-4b
vector_dim: 2560
base_url: http://vllm:11431/v1
reranking:
model:
provider: vllm
name: Qwen/Qwen3-Reranker-4B
base_url: http://vllm:11455
qa:
model:
provider: openai
name: gemma4-26b
base_url: http://vllm:11432/v1
# vLLM enforces input + max_tokens <= max_model_len, so a large output
# budget silently shrinks the input budget. MTRAG answers are sentences.
max_tokens: 8192
evaluations:
judge:
provider: openai
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
base_url: http://vllm:11430/v1
temperature: 0.6
max_tokens: 16384
extra_body:
top_p: 0.95
top_k: 20
min_p: 0
chat_template_kwargs:
enable_thinking: true

View file

@ -8,20 +8,28 @@ import typer
from dotenv import find_dotenv, load_dotenv
from huggingface_hub import HfApi, snapshot_download
from pydantic_evals import Case, Dataset as EvalDataset, set_eval_attribute
from pydantic_evals.evaluators import Evaluator, LLMJudge
from pydantic_evals.evaluators import Evaluator
from pydantic_evals.reporting import ReportCaseFailure
from rich.console import Console
from rich.progress import Progress
from evaluations.config import DatasetSpec
from evaluations.config import ConversationInput, DatasetSpec
from evaluations.datasets import DATASETS
from evaluations.evaluators import (
ANSWER_EQUIVALENCE_RUBRIC,
CitationMAPEvaluator,
MAPEvaluator,
REFUSAL_RUBRIC,
ConversationEvaluator,
RefusalJudge,
TranscriptLLMJudge,
)
from evaluations.capability_runner import (
CapabilityFactory,
prefix_to_messages,
run_capability_conversation,
run_capability_question,
)
from evaluations.capability_runner import CapabilityFactory, run_capability_question
from haiku.rag.client import HaikuRAG
from haiku.rag.client.documents import DocumentImport
from haiku.rag.config import AppConfig, find_config_file, load_yaml_config
from haiku.rag.config.models import ModelConfig
from haiku.rag.logging import configure_cli_logging
@ -118,6 +126,59 @@ def build_experiment_metadata(
return metadata
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,
@ -136,6 +197,17 @@ async def populate_db(
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)
@ -233,16 +305,13 @@ async def run_retrieval_benchmark(
console.print("No retrieval cases to evaluate.")
return None
if spec.retrieval_evaluator is None:
raise ValueError(f"No retrieval evaluator configured for dataset: {spec.key}")
evaluator = spec.retrieval_evaluator
metric_name = evaluator.__class__.__name__.replace("Evaluator", "").upper()
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=[evaluator],
evaluators=list(spec.retrieval_evaluators),
)
db = spec.db_path(db_path)
@ -250,7 +319,10 @@ async def run_retrieval_benchmark(
async def retrieval_target(question: str) -> list[str]:
chunks = await rag.search(
query=question, limit=5, include_images=False, filter=document_filter
query=question,
limit=spec.retrieval_limit,
include_images=False,
filter=document_filter,
)
seen = set()
@ -280,25 +352,22 @@ async def run_retrieval_benchmark(
metadata=experiment_metadata,
)
total_score = 0.0
total_cases = 0
per_metric: dict[str, list[float]] = {}
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
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)}")
console.print(f"{metric_name}: {mean_score:.4f}")
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 {
metric_name.lower(): mean_score,
"queries": len(cases),
}
return results
def _capability_factory_for_target(target: Target) -> CapabilityFactory:
@ -313,13 +382,6 @@ def _capability_factory_for_target(target: Target) -> CapabilityFactory:
raise ValueError(f"target {target!r} is not a capability target")
def _citation_evaluator_for(retrieval_evaluator: Evaluator | None) -> Evaluator | None:
"""Return the citation-scoring twin of the dataset's retrieval evaluator."""
if isinstance(retrieval_evaluator, MAPEvaluator):
return CitationMAPEvaluator()
return None
def _attach_relevant_uris(
cases: list[Case[str, str, dict[str, Any]]],
spec: DatasetSpec,
@ -342,6 +404,8 @@ def _attach_relevant_uris(
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
@ -350,6 +414,102 @@ def _attach_relevant_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") or 0 for case in scored)
turns_passed = sum(_score(case, "turns_passed") for case in scored)
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 scored)
/ len(scored),
}
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") or 0 for case in scored)
false_refusals = sum(_score(case, "false_refusals") or 0 for case in scored)
unanswerable = sum(_score(case, "unanswerable_turns") or 0 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 ("ANSWERABLE", "UNANSWERABLE"):
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).
@ -383,16 +543,11 @@ async def run_qa_benchmark(
]
judge_config = judge_model or DEFAULT_JUDGE_MODEL
if target == "analysis-capability":
# Mirror the capability-code resolver: explicit analysis.model wins,
# else fall back to qa.model.
capability_config = capability_model or config.analysis.model or config.qa.model
else:
capability_config = capability_model or config.qa.model
capability_config = _resolve_capability_config(target, config, capability_model)
db = spec.db_path(db_path)
_attach_relevant_uris(cases, spec, limit)
citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
citation_evaluator = spec.citation_evaluator
qa_evaluator = spec.qa_evaluator
evaluators: list[Evaluator]
@ -400,7 +555,7 @@ async def run_qa_benchmark(
evaluators = [qa_evaluator]
else:
evaluators = [
LLMJudge(
TranscriptLLMJudge(
rubric=ANSWER_EQUIVALENCE_RUBRIC,
include_input=True,
include_expected_output=True,
@ -413,8 +568,16 @@ async def run_qa_benchmark(
]
if citation_evaluator is not None:
evaluators.append(citation_evaluator)
if spec.evaluate_refusal:
evaluators.append(
RefusalJudge(
rubric=REFUSAL_RUBRIC,
model=get_model(judge_config, config),
assertion={"evaluation_name": "refused", "include_reason": False},
)
)
evaluation_dataset = EvalDataset[str, str, dict[str, str]](
evaluation_dataset = EvalDataset[Any, str, dict[str, Any]](
name=spec.key, cases=cases, evaluators=evaluators
)
@ -428,8 +591,9 @@ async def run_qa_benchmark(
capability_config=capability_config,
document_filter=document_filter,
)
experiment_metadata.update(spec.experiment_metadata or {})
async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]):
async def _evaluate(answer_fn: Callable[[Any], Awaitable[str]]):
return await evaluation_dataset.evaluate(
answer_fn,
name=eval_name,
@ -441,7 +605,13 @@ async def run_qa_benchmark(
capability_factory = _capability_factory_for_target(target)
resolved_capability_model = get_model(capability_config, config)
async def answer_question(question: str) -> str:
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=capability_factory,
db_path=db,
@ -449,6 +619,7 @@ async def run_qa_benchmark(
question=question,
capability_model=resolved_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)
@ -488,6 +659,11 @@ async def run_qa_benchmark(
console.print(f"Total questions: {total_processed}")
console.print(f"Correct answers: {passing_cases}")
console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)")
if report.cases:
mean_task_time = sum(case.task_duration for case in report.cases) / len(
report.cases
)
console.print(f"Avg task time per case: {mean_task_time:.2f}s")
if citation_evaluator is not None:
score_key = citation_evaluator.get_default_evaluation_name()
@ -508,11 +684,27 @@ async def run_qa_benchmark(
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 spec.evaluate_refusal:
metrics = _refusal_metrics(report.cases)
if metrics 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)"
)
if failures:
console.print("[red]\nSummary of failures:[/red]")
for failure in failures:
@ -524,6 +716,144 @@ async def run_qa_benchmark(
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,
) -> 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.
"""
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)
db = spec.db_path(db_path)
evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]](
name=spec.key,
cases=cases,
evaluators=[
ConversationEvaluator(
rubric=ANSWER_EQUIVALENCE_RUBRIC,
model=get_model(judge_config, config),
)
],
)
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,
)
experiment_metadata.update(spec.experiment_metadata or {})
capability_factory = _capability_factory_for_target(target)
resolved_capability_model = get_model(capability_config, config)
async def answer_conversation(questions: list[str]) -> list[str]:
results = await run_capability_conversation(
capability_factory=capability_factory,
db_path=db,
config=config,
questions=list(questions),
capability_model=resolved_capability_model,
)
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])
return [r.answer for r in results]
report = await evaluation_dataset.evaluate(
answer_conversation,
name=eval_name,
max_concurrency=1,
progress=True,
metadata=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"
)
if report.failures:
console.print("[red]\nSummary of failures:[/red]")
for failure in report.failures:
console.print(f"Case: {failure.name}")
console.print(f"Error: {failure.error_message}")
console.print("")
async def evaluate_dataset(
spec: DatasetSpec,
config: AppConfig,
@ -566,7 +896,8 @@ async def evaluate_dataset(
console.print(
f"\nRunning QA benchmarks (target={target})...", style="bold yellow"
)
await run_qa_benchmark(
qa_benchmark = run_live_qa_benchmark if spec.live else run_qa_benchmark
await qa_benchmark(
spec,
config,
limit=limit,
@ -621,9 +952,20 @@ def _resolve_dataset(dataset: str) -> DatasetSpec:
def _resolve_datasets(dataset: str) -> list[DatasetSpec]:
"""Resolve 'all' or a single dataset key to a list of DatasetSpecs."""
"""Resolve 'all' or a single dataset key to a list of DatasetSpecs.
'all' yields one spec per database: query variants sharing a db_filename
would otherwise be downloaded/uploaded twice.
"""
if dataset.lower() == "all":
return list(DATASETS.values())
seen: set[str] = set()
specs: list[DatasetSpec] = []
for spec in DATASETS.values():
if spec.db_filename in seen:
continue
seen.add(spec.db_filename)
specs.append(spec)
return specs
return [_resolve_dataset(dataset)]

View file

@ -1,4 +1,4 @@
from collections.abc import Callable
from collections.abc import Callable, Iterable
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, NamedTuple, Protocol, cast
@ -6,13 +6,17 @@ from typing import Any, NamedTuple, Protocol, cast
from pydantic_ai import Agent
from pydantic_ai.messages import (
ModelMessage,
ModelRequest,
ModelResponse,
RetryPromptPart,
TextPart,
ToolCallPart,
ToolReturnPart,
UserPromptPart,
)
from pydantic_ai.models import Model
from evaluations.config import Turn
from haiku.rag.capabilities import RAGCapabilityBase
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
@ -21,6 +25,17 @@ from haiku.rag.store.models.citation import Citation
CapabilityFactory = Callable[..., RAGCapabilityBase[Any]]
def prefix_to_messages(turns: Iterable[Turn]) -> list[ModelMessage]:
"""Render a conversation prefix as pydantic-ai message history."""
messages: list[ModelMessage] = []
for turn in turns:
if turn.speaker == "user":
messages.append(ModelRequest(parts=[UserPromptPart(content=turn.text)]))
else:
messages.append(ModelResponse(parts=[TextPart(content=turn.text)]))
return messages
class _RagLikeState(Protocol):
document_filter: str | None
citation_index: dict[str, Citation]
@ -107,25 +122,14 @@ class _EvalDeps:
state: dict[str, Any] = field(default_factory=dict)
async def run_capability_question(
def _prepare_agent(
capability_factory: CapabilityFactory,
db_path: Path,
config: AppConfig,
question: str,
capability_model: str | Model,
document_filter: str | None = None,
request_limit: int | None = None,
) -> CapabilityRunResult:
"""Run a single question through a capability and return answer + retrieval data.
Builds a native capability via ``capability_factory(db_path=..., config=...)``.
After the run, citations and searched documents
are extracted from the state for downstream eval scoring.
The capability must produce a state with RAG-capability-shaped fields (citation
index, searches, optional document filter) i.e. ``RAGState`` or
``AnalysisState`` from ``haiku.rag.capabilities``.
"""
document_filter: str | None,
request_limit: int | None,
) -> tuple[RAGCapabilityBase[Any], _EvalDeps, Agent[_EvalDeps, str]]:
capability = capability_factory(
db_path=db_path,
config=config,
@ -144,10 +148,98 @@ async def run_capability_question(
deps_type=_EvalDeps,
capabilities=[capability],
)
agent_result = await agent.run(question, deps=deps)
state = capability.state_type.model_validate(deps.state[capability.state_namespace])
typed = cast(_RagLikeState, state)
return capability, deps, agent
def _state_after_run(
capability: RAGCapabilityBase[Any], deps: _EvalDeps
) -> _RagLikeState:
state = capability.state_type.model_validate(deps.state[capability.state_namespace])
return cast(_RagLikeState, state)
async def run_capability_question(
capability_factory: CapabilityFactory,
db_path: Path,
config: AppConfig,
question: str,
capability_model: str | Model,
document_filter: str | None = None,
request_limit: int | None = None,
message_history: list[ModelMessage] | None = None,
) -> CapabilityRunResult:
"""Run a single question through a capability and return answer + retrieval data.
Builds a native capability via ``capability_factory(db_path=..., config=...)``.
After the run, citations and searched documents
are extracted from the state for downstream eval scoring.
The capability must produce a state with RAG-capability-shaped fields (citation
index, searches, optional document filter) i.e. ``RAGState`` or
``AnalysisState`` from ``haiku.rag.capabilities``.
"""
capability, deps, agent = _prepare_agent(
capability_factory,
db_path,
config,
capability_model,
document_filter,
request_limit,
)
agent_result = await agent.run(question, deps=deps, message_history=message_history)
traffic = _count_tool_traffic(
agent_result.new_messages(), capability.state_namespace, capability.tool_names
)
return _result_from_run(
agent_result.output, _state_after_run(capability, deps), traffic
)
async def run_capability_conversation(
capability_factory: CapabilityFactory,
db_path: Path,
config: AppConfig,
questions: list[str],
capability_model: str | Model,
document_filter: str | None = None,
request_limit: int | None = None,
) -> list[CapabilityRunResult]:
"""Run a conversation's user turns sequentially through one capability.
Each turn runs with the previous turn's full ``all_messages()`` as history
(tool calls and returns included), so prior-turn compaction operates on
real evidence. Per-invocation state (citations, searches) is cleared by the
capability on every run, so each returned result reflects only its turn.
"""
capability, deps, agent = _prepare_agent(
capability_factory,
db_path,
config,
capability_model,
document_filter,
request_limit,
)
history: list[ModelMessage] | None = None
results: list[CapabilityRunResult] = []
for question in questions:
agent_result = await agent.run(question, deps=deps, message_history=history)
history = agent_result.all_messages()
traffic = _count_tool_traffic(
agent_result.new_messages(),
capability.state_namespace,
capability.tool_names,
)
results.append(
_result_from_run(
agent_result.output, _state_after_run(capability, deps), traffic
)
)
return results
def _result_from_run(
answer: str, typed: _RagLikeState, traffic: ToolTraffic
) -> CapabilityRunResult:
cited_chunk_ids: list[str] = list(typed.citations)
seen_cited: set[str] = set()
cited_uris: list[str] = []
@ -168,17 +260,11 @@ async def run_capability_question(
seen_searched.add(uri)
searched_uris.append(uri)
executions = getattr(state, "executions", None)
executions = getattr(typed, "executions", None)
n_executions = len(executions) if executions is not None else 0
traffic = _count_tool_traffic(
agent_result.all_messages(),
capability.state_namespace,
capability.tool_names,
)
return CapabilityRunResult(
answer=agent_result.output,
answer=answer,
cited_uris=cited_uris,
cited_chunk_ids=cited_chunk_ids,
searched_uris=searched_uris,

View file

@ -1,13 +1,43 @@
from collections.abc import Callable, Mapping
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from typing import Any, Literal
from datasets import Dataset
from pydantic import BaseModel, model_validator
from pydantic_evals import Case
from pydantic_evals.evaluators import Evaluator
class Turn(BaseModel):
speaker: Literal["user", "agent"]
text: str
class ConversationInput(BaseModel):
"""A conversation prefix plus the final user question (the last turn)."""
turns: list[Turn]
@model_validator(mode="after")
def _ends_with_user_turn(self) -> "ConversationInput":
if not self.turns or self.turns[-1].speaker != "user":
raise ValueError("conversation must end with a user turn")
return self
@property
def question(self) -> str:
return self.turns[-1].text
@property
def prefix(self) -> list[Turn]:
return self.turns[:-1]
@property
def transcript(self) -> str:
return "\n".join(f"{turn.speaker}: {turn.text}" for turn in self.turns)
@dataclass
class DocumentPayload:
uri: str
@ -30,7 +60,8 @@ DocumentLoader = Callable[[], Dataset]
DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None]
RetrievalLoader = Callable[[], Dataset]
RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None]
CaseBuilder = Callable[[int, Mapping[str, Any]], Case[str, str, dict[str, str]]]
QAInput = str | ConversationInput
CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]]
@dataclass
@ -43,9 +74,15 @@ class DatasetSpec:
qa_case_builder: CaseBuilder
retrieval_loader: RetrievalLoader | None = None
retrieval_mapper: RetrievalMapper | None = None
retrieval_evaluator: Evaluator | None = None
retrieval_evaluators: list[Evaluator] | None = None
citation_evaluator: Evaluator | None = None
qa_evaluator: Evaluator | None = None
document_limit: int | None = None
retrieval_limit: int = 5
ingest_batch_size: int | None = None
evaluate_refusal: bool = False
live: bool = False
experiment_metadata: dict[str, Any] | None = None
def db_path(self, override_path: Path | None = None) -> Path:
"""Get the database path.

View file

@ -1,6 +1,11 @@
from evaluations.config import DatasetSpec
from .hotpotqa import HOTPOTQA_SPEC
from .mtrag import (
MTRAG_CLAPNQ_LIVE_SPEC,
MTRAG_CLAPNQ_REWRITE_SPEC,
MTRAG_CLAPNQ_SPEC,
)
from .open_rag_bench import (
ORB_MULTIMODAL_NEMOTRON_SPEC,
ORB_MULTIMODAL_SPEC,
@ -12,6 +17,9 @@ DATASETS: dict[str, DatasetSpec] = {
spec.key: spec
for spec in (
HOTPOTQA_SPEC,
MTRAG_CLAPNQ_SPEC,
MTRAG_CLAPNQ_REWRITE_SPEC,
MTRAG_CLAPNQ_LIVE_SPEC,
ORB_TEXT_SPEC,
ORB_MULTIMODAL_SPEC,
ORB_MULTIMODAL_NEMOTRON_SPEC,

View file

@ -5,7 +5,7 @@ from datasets import Dataset, load_dataset
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MAPEvaluator
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
def load_hotpotqa_validation() -> Dataset:
@ -104,5 +104,6 @@ HOTPOTQA_SPEC = DatasetSpec(
qa_case_builder=build_hotpotqa_case,
retrieval_loader=load_hotpotqa_validation,
retrieval_mapper=map_hotpotqa_retrieval,
retrieval_evaluator=MAPEvaluator(),
retrieval_evaluators=[MAPEvaluator()],
citation_evaluator=CitationMAPEvaluator(),
)

View file

@ -0,0 +1,298 @@
import json
import zipfile
from collections.abc import Iterable, Mapping
from functools import partial
from pathlib import Path
from typing import Any
import httpx
from datasets import Dataset
from pydantic_evals import Case
from evaluations.config import (
ConversationInput,
DatasetSpec,
DocumentPayload,
RetrievalSample,
Turn,
)
from evaluations.evaluators import (
CitationMAPEvaluator,
MAPEvaluator,
NDCGEvaluator,
RecallEvaluator,
)
REPO_SHA = "cc5b1d481b391181b89f7ced860308482e785463"
_BASE_URL = f"https://raw.githubusercontent.com/IBM/mt-rag-benchmark/{REPO_SHA}"
_CORPUS_FILE = "corpora/passage_level/clapnq.jsonl.zip"
_QRELS_FILE = "mtrag-human/retrieval_tasks/clapnq/qrels/dev.tsv"
_QUERY_FILES = {
"lastturn": "mtrag-human/retrieval_tasks/clapnq/clapnq_lastturn.jsonl",
"rewrite": "mtrag-human/retrieval_tasks/clapnq/clapnq_rewrite.jsonl",
}
_GEN_TASKS_FILE = "mtrag-human/generation_tasks/reference.jsonl"
_CLAPNQ_COLLECTION = "mt-rag-clapnq-elser-512-100-20240503"
def get_cache_dir() -> Path:
cache_dir = Path.home() / ".cache" / "haiku.rag" / "evaluations" / "mtrag"
cache_dir.mkdir(parents=True, exist_ok=True)
return cache_dir
def _download(rel_path: str) -> Path:
dest = get_cache_dir() / rel_path.replace("/", "_")
if dest.exists():
return dest
with httpx.stream(
"GET", f"{_BASE_URL}/{rel_path}", timeout=120.0, follow_redirects=True
) as response:
response.raise_for_status()
tmp = dest.with_suffix(dest.suffix + ".part")
with tmp.open("wb") as fh:
for data in response.iter_bytes():
fh.write(data)
tmp.rename(dest)
return dest
def _parse_qrels(lines: Iterable[str]) -> dict[str, list[str]]:
"""Group qrel corpus-ids by query-id, preserving file order."""
qrels: dict[str, list[str]] = {}
rows = iter(lines)
next(rows) # header: query-id / corpus-id / score
for line in rows:
if not line.strip():
continue
query_id, corpus_id, _score = line.rstrip("\n").split("\t")
qrels.setdefault(query_id, []).append(corpus_id)
return qrels
def _validate_qrels_resolve(
corpus_ids: set[str], qrels: Mapping[str, list[str]]
) -> None:
unresolved = sorted(
{cid for ids in qrels.values() for cid in ids if cid not in corpus_ids}
)
if unresolved:
raise ValueError(
f"{len(unresolved)} qrel corpus-ids do not resolve to corpus "
f"passages, e.g. {unresolved[:3]}"
)
def _join_queries_qrels(
queries: Iterable[Mapping[str, Any]], qrels: Mapping[str, list[str]]
) -> list[dict[str, Any]]:
records = []
for query in queries:
query_id = query["_id"]
expected = qrels.get(query_id)
if expected is None:
raise ValueError(f"query {query_id} has no qrels")
records.append(
{
"query_id": query_id,
"question": query["text"],
"expected_uris": expected,
}
)
return records
def _load_qrels() -> dict[str, list[str]]:
path = _download(_QRELS_FILE)
return _parse_qrels(path.read_text().splitlines())
def load_clapnq_corpus() -> Dataset:
path = _download(_CORPUS_FILE)
records: list[dict[str, str]] = []
with zipfile.ZipFile(path) as zf:
with zf.open(zf.namelist()[0]) as fh:
for line in fh:
rec = json.loads(line)
records.append(
{"_id": rec["_id"], "title": rec["title"], "text": rec["text"]}
)
_validate_qrels_resolve({rec["_id"] for rec in records}, _load_qrels())
return Dataset.from_list(records)
def map_mtrag_document(doc: Mapping[str, Any]) -> DocumentPayload:
return DocumentPayload(uri=doc["_id"], content=doc["text"], title=doc["title"])
def load_clapnq_retrieval(variant: str) -> Dataset:
path = _download(_QUERY_FILES[variant])
queries = [json.loads(line) for line in path.read_text().splitlines() if line]
return Dataset.from_list(_join_queries_qrels(queries, _load_qrels()))
def map_mtrag_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
return RetrievalSample(
question=doc["question"],
expected_uris=tuple(doc["expected_uris"]),
)
def _task_to_record(
task: Mapping[str, Any], qrels: Mapping[str, list[str]]
) -> dict[str, Any] | None:
"""Reduce a reference.jsonl generation task to the fields QA cases need.
Task `contexts` are the original system's retrievals, never gold relevance;
gold passages come from the qrels keyed by task_id.
"""
if task["Collection"] != _CLAPNQ_COLLECTION:
return None
return {
"id": task["task_id"],
"turn": task["turn"],
"turns": [
{"speaker": message["speaker"], "text": message["text"]}
for message in task["input"]
],
"answer": task["targets"][0]["text"],
"answerability": task["Answerability"][0],
"multi_turn_type": task["Multi-Turn"][0],
"question_type": list(task["Question Type"]),
"relevant_uris": qrels.get(task["task_id"]),
}
def load_clapnq_qa() -> Dataset:
path = _download(_GEN_TASKS_FILE)
qrels = _load_qrels()
records = []
for line in path.read_text().splitlines():
if not line.strip():
continue
record = _task_to_record(json.loads(line), qrels)
if record is not None:
records.append(record)
return Dataset.from_list(records)
def build_mtrag_case(
index: int, doc: Mapping[str, Any]
) -> Case[ConversationInput, str, dict[str, Any]]:
metadata: dict[str, Any] = {
"task_id": doc["id"],
"turn": doc["turn"],
"answerability": doc["answerability"],
"multi_turn_type": doc["multi_turn_type"],
"question_type": list(doc["question_type"]),
}
if doc["relevant_uris"]:
metadata["relevant_uris"] = list(doc["relevant_uris"])
return Case(
name=f"{index}_{doc['id']}",
inputs=ConversationInput(
turns=[Turn(**turn) for turn in doc["turns"]],
),
expected_output=doc["answer"],
metadata=metadata,
)
def _group_conversations(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
"""Group per-turn generation records into full conversations.
Turns are ordered numerically within each conversation; each turn carries
its user question, reference answer, answerability label, and gold
passages when the turn has qrels.
"""
grouped: dict[str, list[dict[str, Any]]] = {}
for record in records:
conversation_id = record["id"].split("<::>")[0]
grouped.setdefault(conversation_id, []).append(record)
conversations = []
for conversation_id, tasks in grouped.items():
tasks.sort(key=lambda record: int(record["turn"]))
turns = []
for task in tasks:
turn: dict[str, Any] = {
"task_id": task["id"],
"turn": task["turn"],
"question": task["turns"][-1]["text"],
"reference": task["answer"],
"answerability": task["answerability"],
"multi_turn_type": task["multi_turn_type"],
"question_type": list(task["question_type"]),
}
if task["relevant_uris"]:
turn["relevant_uris"] = list(task["relevant_uris"])
turns.append(turn)
conversations.append({"id": conversation_id, "turns": turns})
return conversations
def load_clapnq_conversations() -> Dataset:
corpus = load_clapnq_qa()
return Dataset.from_list(_group_conversations([dict(row) for row in corpus]))
def build_mtrag_live_case(
index: int, doc: Mapping[str, Any]
) -> Case[list[str], list[str], dict[str, Any]]:
questions = [turn["question"] for turn in doc["turns"]]
metadata_turns = [
{
key: value
for key, value in turn.items()
if key != "question" and value is not None
}
for turn in doc["turns"]
]
return Case(
name=f"{index}_{doc['id']}",
inputs=questions,
metadata={"conversation_id": doc["id"], "turns": metadata_turns},
)
def _mtrag_spec(key: str, variant: str) -> DatasetSpec:
return DatasetSpec(
key=key,
db_filename="mtrag_clapnq.lancedb",
document_loader=load_clapnq_corpus,
document_mapper=map_mtrag_document,
qa_loader=load_clapnq_qa,
qa_case_builder=build_mtrag_case,
retrieval_loader=partial(load_clapnq_retrieval, variant),
retrieval_mapper=map_mtrag_retrieval,
retrieval_evaluators=[
RecallEvaluator(k=5),
RecallEvaluator(k=10),
NDCGEvaluator(k=5),
NDCGEvaluator(k=10),
MAPEvaluator(),
],
citation_evaluator=CitationMAPEvaluator(),
retrieval_limit=10,
ingest_batch_size=512,
evaluate_refusal=True,
experiment_metadata={"mtrag_mode": "gold_prefix"},
)
MTRAG_CLAPNQ_SPEC = _mtrag_spec("mtrag_clapnq", "lastturn")
MTRAG_CLAPNQ_REWRITE_SPEC = _mtrag_spec("mtrag_clapnq_rewrite", "rewrite")
MTRAG_CLAPNQ_LIVE_SPEC = DatasetSpec(
key="mtrag_clapnq_live",
db_filename="mtrag_clapnq.lancedb",
document_loader=load_clapnq_corpus,
document_mapper=map_mtrag_document,
qa_loader=load_clapnq_conversations,
qa_case_builder=build_mtrag_live_case,
ingest_batch_size=512,
live=True,
experiment_metadata={"mtrag_mode": "live_session"},
)

View file

@ -10,7 +10,7 @@ from huggingface_hub import hf_hub_download
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MAPEvaluator
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
logger = logging.getLogger(__name__)
@ -226,7 +226,8 @@ def _orb_spec(key: str, db_filename: str) -> DatasetSpec:
qa_case_builder=build_orb_case,
retrieval_loader=load_orb_retrieval,
retrieval_mapper=map_orb_retrieval,
retrieval_evaluator=MAPEvaluator(),
retrieval_evaluators=[MAPEvaluator()],
citation_evaluator=CitationMAPEvaluator(),
)

View file

@ -10,7 +10,11 @@ from huggingface_hub import hf_hub_download
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MAPEvaluator, NumberMatchEvaluator
from evaluations.evaluators import (
CitationMAPEvaluator,
MAPEvaluator,
NumberMatchEvaluator,
)
REPO_ID = "G4KMU/t2-ragbench"
@ -142,7 +146,8 @@ def _t2_spec(subset: str, key: str, db_filename: str) -> DatasetSpec:
qa_case_builder=build_t2_case,
retrieval_loader=partial(load_t2_qa, subset),
retrieval_mapper=map_t2_retrieval,
retrieval_evaluator=MAPEvaluator(),
retrieval_evaluators=[MAPEvaluator()],
citation_evaluator=CitationMAPEvaluator(),
qa_evaluator=NumberMatchEvaluator(),
)

View file

@ -1,4 +1,5 @@
from evaluations.evaluators.citation import CitationMAPEvaluator
from evaluations.evaluators.conversation import ConversationEvaluator
from evaluations.evaluators.judge import (
ANSWER_EQUIVALENCE_RUBRIC,
LLMJudge,
@ -6,12 +7,21 @@ from evaluations.evaluators.judge import (
)
from evaluations.evaluators.map import MAPEvaluator
from evaluations.evaluators.number_match import NumberMatchEvaluator
from evaluations.evaluators.refusal import REFUSAL_RUBRIC, RefusalJudge
from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator
from evaluations.evaluators.transcript import TranscriptLLMJudge
__all__ = [
"ANSWER_EQUIVALENCE_RUBRIC",
"REFUSAL_RUBRIC",
"CitationMAPEvaluator",
"ConversationEvaluator",
"LLMJudge",
"LLMJudgeResponseSchema",
"MAPEvaluator",
"NDCGEvaluator",
"NumberMatchEvaluator",
"RecallEvaluator",
"RefusalJudge",
"TranscriptLLMJudge",
]

View file

@ -1,6 +1,7 @@
from dataclasses import dataclass
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
from pydantic_evals.evaluators.evaluator import EvaluatorOutput
def _cited_uris(ctx: EvaluatorContext) -> list[str]:
@ -13,28 +14,34 @@ def _relevant_uris(ctx: EvaluatorContext) -> set[str]:
return set(ctx.metadata.get("relevant_uris", []))
def average_precision(cited: list[str], relevant: set[str]) -> float:
"""AP of the cited URIs against the relevant set (0.0 when nothing hits)."""
precisions: list[float] = []
found = 0
for rank, uri in enumerate(cited, start=1):
if uri in relevant:
found += 1
precisions.append(found / rank)
if not precisions:
return 0.0
return sum(precisions) / len(relevant)
@dataclass
class CitationMAPEvaluator(Evaluator):
"""Average precision over the URIs the capability cited via the `cite` tool.
Reads ``cited_uris`` from ``ctx.attributes`` (recorded during the task run
via :func:`pydantic_evals.set_eval_attribute`) and ``relevant_uris`` from
``ctx.metadata``.
``ctx.metadata``. Cases without relevant URIs (e.g. unanswerable turns)
are ineligible and produce no score.
"""
def get_default_evaluation_name(self) -> str:
return "cited_map"
def evaluate(self, ctx: EvaluatorContext) -> float:
def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
relevant = _relevant_uris(ctx)
if not relevant:
return 0.0
precisions: list[float] = []
found = 0
for rank, uri in enumerate(_cited_uris(ctx), start=1):
if uri in relevant:
found += 1
precisions.append(found / rank)
if not precisions:
return 0.0
return sum(precisions) / len(relevant)
return {}
return average_precision(_cited_uris(ctx), relevant)

View file

@ -0,0 +1,120 @@
from dataclasses import dataclass
from pydantic_ai import models
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
from pydantic_evals.evaluators.evaluator import EvaluationReason, EvaluatorOutput
from pydantic_evals.evaluators.llm_as_a_judge import (
judge_input_output_expected,
judge_output,
)
from evaluations.evaluators.citation import average_precision
from evaluations.evaluators.refusal import REFUSAL_RUBRIC
_REFUSAL_LABELS = ("ANSWERABLE", "UNANSWERABLE")
@dataclass
class ConversationEvaluator(Evaluator):
"""Score a live-session conversation turn by turn.
Expects the case output to be the list of per-turn answers, case inputs
the list of user questions, ``metadata["turns"]`` the per-turn reference,
answerability label, and optional gold ``relevant_uris``, and the
``turn_cited_uris`` attribute the per-turn cited URIs.
Each turn's answer is judged against the reference with the conversation
so far including the model's own earlier answers — as context. Citation
AP is computed on turns with gold passages; refusal on ANSWERABLE and
UNANSWERABLE turns. Returned counts allow micro aggregation across
conversations; ``turn_pass_rate`` is the per-conversation (macro) rate.
Per-turn verdicts are returned as ``turn_{n}_pass`` (with the judge's
reason), ``turn_{n}_refused``, and ``turn_{n}_cited_ap`` for diagnosis.
"""
rubric: str
model: models.Model | models.KnownModelName | str | None = None
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
questions: list[str] = list(ctx.inputs)
answers: list[str] = list(ctx.output)
turns: list[dict] = (ctx.metadata or {}).get("turns", [])
turn_cited: list[list[str]] = list(
ctx.attributes.get("turn_cited_uris") or [[] for _ in answers]
)
if not (len(questions) == len(answers) == len(turns) == len(turn_cited)):
raise ValueError(
f"conversation arrays disagree: {len(questions)} questions, "
f"{len(answers)} answers, {len(turns)} turn annotations, "
f"{len(turn_cited)} citation lists"
)
passed = 0
judged = 0
citation_scores: list[float] = []
true_refusals = 0
false_refusals = 0
unanswerable = 0
per_turn: dict[str, EvaluationReason | bool | float | str] = {}
transcript_lines: list[str] = []
for index, (question, answer, turn) in enumerate(
zip(questions, answers, turns)
):
number = index + 1
transcript_lines.append(f"user: {question}")
transcript = "\n".join(transcript_lines)
transcript_lines.append(f"agent: {answer}")
try:
grading = await judge_input_output_expected(
transcript, answer, turn["reference"], self.rubric, self.model
)
except Exception as error:
per_turn[f"turn_{number}_judge_error"] = str(error)[:200]
else:
judged += 1
if grading.pass_:
passed += 1
per_turn[f"turn_{number}_pass"] = EvaluationReason(
value=grading.pass_, reason=grading.reason
)
label = turn.get("answerability")
if label in _REFUSAL_LABELS:
try:
refused = (
await judge_output(answer, REFUSAL_RUBRIC, self.model)
).pass_
except Exception as error:
per_turn[f"turn_{number}_judge_error"] = str(error)[:200]
else:
per_turn[f"turn_{number}_refused"] = refused
if label == "UNANSWERABLE":
unanswerable += 1
if refused:
true_refusals += 1
elif refused:
false_refusals += 1
relevant = set(turn.get("relevant_uris") or [])
if relevant:
turn_ap = average_precision(turn_cited[index], relevant)
citation_scores.append(turn_ap)
per_turn[f"turn_{number}_cited_ap"] = turn_ap
total = len(answers)
result: dict[str, EvaluationReason | bool | float | int | str] = {
"turn_pass_rate": passed / judged if judged else 0.0,
"turns_passed": passed,
"turns_judged": judged,
"turns_total": total,
"cited_eligible": len(citation_scores),
"true_refusals": true_refusals,
"false_refusals": false_refusals,
"unanswerable_turns": unanswerable,
}
if citation_scores:
result["cited_map"] = sum(citation_scores) / len(citation_scores)
result.update(per_turn)
return result

View file

@ -0,0 +1,28 @@
from dataclasses import dataclass
from pydantic_evals.evaluators import EvaluatorContext, LLMJudge
from pydantic_evals.evaluators.evaluator import EvaluatorOutput
REFUSAL_RUBRIC = (
"The output declines to answer: it states that the information is not "
"available, cannot be found, or is not known, rather than providing a "
"substantive answer."
)
_ELIGIBLE_LABELS = ("ANSWERABLE", "UNANSWERABLE")
@dataclass
class RefusalJudge(LLMJudge):
"""Judges whether the output declines to answer.
Runs only on turns whose answerability label makes refusal meaningful
(ANSWERABLE/UNANSWERABLE); other turns produce no evaluation and cost no
judge call.
"""
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
label = (ctx.metadata or {}).get("answerability")
if label not in _ELIGIBLE_LABELS:
return {}
return await super().evaluate(ctx)

View file

@ -0,0 +1,54 @@
import math
from dataclasses import dataclass
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
def _relevant_and_retrieved(ctx: EvaluatorContext) -> tuple[set[str], list[str]]:
if ctx.metadata is None:
return set(), []
return set(ctx.metadata.get("relevant_uris", [])), list(ctx.output)
@dataclass
class RecallEvaluator(Evaluator):
"""Recall@k: fraction of relevant documents retrieved in the top k."""
k: int
def get_default_evaluation_name(self) -> str:
return f"recall_{self.k}"
def evaluate(self, ctx: EvaluatorContext) -> float:
relevant, retrieved = _relevant_and_retrieved(ctx)
if not relevant:
return 0.0
found = sum(1 for uri in retrieved[: self.k] if uri in relevant)
return found / len(relevant)
@dataclass
class NDCGEvaluator(Evaluator):
"""Binary nDCG@k: DCG of relevant documents in the top k over the ideal DCG.
Gains are binary (relevant or not), matching qrels without graded scores.
"""
k: int
def get_default_evaluation_name(self) -> str:
return f"ndcg_{self.k}"
def evaluate(self, ctx: EvaluatorContext) -> float:
relevant, retrieved = _relevant_and_retrieved(ctx)
if not relevant:
return 0.0
dcg = sum(
1 / math.log2(rank + 1)
for rank, uri in enumerate(retrieved[: self.k], start=1)
if uri in relevant
)
ideal = sum(
1 / math.log2(rank + 1) for rank in range(1, min(len(relevant), self.k) + 1)
)
return dcg / ideal

View file

@ -0,0 +1,21 @@
from dataclasses import dataclass, replace
from pydantic_evals.evaluators import EvaluatorContext, LLMJudge
from pydantic_evals.evaluators.evaluator import EvaluatorOutput
from evaluations.config import ConversationInput
@dataclass
class TranscriptLLMJudge(LLMJudge):
"""LLMJudge that shows conversation inputs as a readable transcript.
pydantic-evals serializes custom input models as JSON in the judge prompt;
a ConversationInput is rendered as `speaker: text` lines instead. Plain
string inputs pass through unchanged.
"""
async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput:
if isinstance(ctx.inputs, ConversationInput):
ctx = replace(ctx, inputs=ctx.inputs.transcript)
return await super().evaluate(ctx)

View file

@ -1,5 +1,5 @@
from pathlib import Path
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import typer
@ -11,7 +11,7 @@ from evaluations.benchmark import (
evaluate_dataset,
run_qa_benchmark,
)
from evaluations.config import DatasetSpec
from evaluations.config import DatasetSpec, DocumentPayload
from haiku.rag.config.models import AppConfig, ModelConfig
@ -129,6 +129,378 @@ class TestResolveDataset:
_resolve_dataset("nonexistent")
class TestConversationInputDispatch:
@pytest.mark.asyncio
async def test_prefix_rides_as_message_history(self, tmp_path: Path) -> None:
"""A ConversationInput case reaches the capability as final question
plus the prefix converted to message history."""
from dataclasses import dataclass
from pydantic_evals import Case
from pydantic_evals.evaluators import Evaluator, EvaluatorContext
from evaluations.capability_runner import CapabilityRunResult
from evaluations.config import ConversationInput, Turn
@dataclass
class AlwaysOne(Evaluator):
def evaluate(self, ctx: EvaluatorContext) -> float:
return 1.0
def build_case(idx: int, doc) -> Case:
return Case(
name="c1",
inputs=ConversationInput(
turns=[
Turn(speaker="user", text="q1"),
Turn(speaker="agent", text="a1"),
Turn(speaker="user", text="q2"),
]
),
expected_output="ref",
)
spec = DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [{"id": "t1"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=build_case,
qa_evaluator=AlwaysOne(),
)
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.run_capability_question",
new_callable=AsyncMock,
return_value=CapabilityRunResult(answer="answer"),
) as run_question,
):
await run_qa_benchmark(spec, AppConfig(), db_path=tmp_path / "test.lancedb")
assert run_question.await_args is not None
kwargs = run_question.await_args.kwargs
assert kwargs["question"] == "q2"
history = kwargs["message_history"]
assert len(history) == 2
assert history[0].parts[0].content == "q1"
assert history[1].parts[0].content == "a1"
class TestRefusalMetrics:
def _case(self, label: str | None, refused: bool | None) -> MagicMock:
case = MagicMock()
case.metadata = {"answerability": label} if label is not None else {}
case.assertions = (
{"refused": MagicMock(value=refused)} if refused is not None else {}
)
return case
def test_precision_and_recall(self) -> None:
from evaluations.benchmark import _refusal_metrics
cases = [
self._case("UNANSWERABLE", True), # true refusal
self._case("UNANSWERABLE", False), # missed refusal
self._case("ANSWERABLE", True), # false refusal
self._case("ANSWERABLE", False), # answered correctly
self._case("PARTIAL", None), # skipped by the judge, no assertion
self._case(None, None), # no label
]
metrics = _refusal_metrics(cases)
assert metrics is not None
precision, recall, unanswerable, refusals = metrics
assert precision == 0.5 # 1 true refusal of 2 refusals
assert recall == 0.5 # 1 of 2 unanswerable turns refused
assert unanswerable == 2
assert refusals == 2
def test_none_when_no_judged_cases(self) -> None:
from evaluations.benchmark import _refusal_metrics
assert _refusal_metrics([self._case("PARTIAL", None)]) is None
class TestLiveSummary:
def _case(self, scores: dict[str, float | int]) -> MagicMock:
case = MagicMock()
case.scores = {key: MagicMock(value=value) for key, value in scores.items()}
return case
def test_micro_and_macro_aggregation(self) -> None:
from evaluations.benchmark 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).
cases = [
self._case(
{
"turn_pass_rate": 0.25,
"turns_passed": 1,
"turns_judged": 4,
"turns_total": 4,
"cited_map": 0.5,
"cited_eligible": 3,
"true_refusals": 1,
"false_refusals": 1,
"unanswerable_turns": 2,
}
),
self._case(
{
"turn_pass_rate": 1.0,
"turns_passed": 2,
"turns_judged": 2,
"turns_total": 2,
"cited_map": 1.0,
"cited_eligible": 1,
"true_refusals": 0,
"false_refusals": 0,
"unanswerable_turns": 0,
}
),
]
failure = MagicMock()
failure.inputs = ["fq1", "fq2", "fq3"]
summary = _live_summary(cases, [failure])
assert summary is not None
assert summary["conversations"] == 2
assert summary["conversations_attempted"] == 3
assert summary["turns_total"] == 6
assert summary["turns_judged"] == 6
assert summary["turns_attempted"] == 9
assert summary["micro_pass_rate"] == pytest.approx(0.5)
assert summary["macro_pass_rate"] == pytest.approx(0.625)
assert summary["cited_eligible"] == 4
assert summary["cited_map_micro"] == pytest.approx((0.5 * 3 + 1.0 * 1) / 4)
assert summary["cited_map_macro"] == pytest.approx(0.75)
assert summary["refusal_precision"] == pytest.approx(0.5)
assert summary["refusal_recall"] == pytest.approx(0.5)
def test_none_without_scored_cases(self) -> None:
from evaluations.benchmark 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
cases = [
self._case(
{
"turn_pass_rate": 1.0,
"turns_passed": 3,
"turns_judged": 3,
"turns_total": 4, # one turn's judge errored
"cited_eligible": 0,
"true_refusals": 0,
"false_refusals": 0,
"unanswerable_turns": 0,
}
)
]
summary = _live_summary(cases)
assert summary is not None
assert summary["micro_pass_rate"] == 1.0
assert summary["turns_judged"] == 3
assert summary["turns_total"] == 4
def test_failed_conversations_do_not_affect_rates(self) -> None:
from evaluations.benchmark import _live_summary
cases = [
self._case(
{
"turn_pass_rate": 1.0,
"turns_passed": 2,
"turns_judged": 2,
"turns_total": 2,
"cited_eligible": 0,
"true_refusals": 0,
"false_refusals": 0,
"unanswerable_turns": 0,
}
)
]
failure = MagicMock()
failure.inputs = ["fq1", "fq2"]
summary = _live_summary(cases, [failure])
assert summary is not None
assert summary["micro_pass_rate"] == 1.0
assert summary["macro_pass_rate"] == 1.0
assert summary["conversations_attempted"] == 2
assert summary["turns_attempted"] == 4
class TestLiveConversationDispatch:
@pytest.mark.asyncio
async def test_live_spec_replays_conversation(self, tmp_path: Path) -> None:
from pydantic_evals import Case
from evaluations.benchmark import run_live_qa_benchmark
from evaluations.capability_runner import CapabilityRunResult
def build_case(idx: int, doc) -> Case:
return Case(
name="conv1",
inputs=["q1", "q2"],
metadata={
"conversation_id": "conv1",
"turns": [
{"reference": "r1", "answerability": "ANSWERABLE"},
{"reference": "r2", "answerability": "ANSWERABLE"},
],
},
)
spec = DatasetSpec(
key="test_live",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [{"id": "conv1"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=build_case,
live=True,
)
turn_results = [
CapabilityRunResult(answer="a1", cited_uris=["u1"]),
CapabilityRunResult(answer="a2", cited_uris=[]),
]
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.run_capability_conversation",
new_callable=AsyncMock,
return_value=turn_results,
) as run_conversation,
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
return_value=MagicMock(score=None, pass_=True, reason=None),
),
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
return_value=MagicMock(score=None, pass_=False, reason=None),
),
):
await run_live_qa_benchmark(
spec, AppConfig(), db_path=tmp_path / "test.lancedb"
)
assert run_conversation.await_args is not None
assert run_conversation.await_args.kwargs["questions"] == ["q1", "q2"]
@pytest.mark.asyncio
async def test_live_records_per_turn_traffic_arrays(self, tmp_path: Path) -> None:
"""Per-turn tool traffic is recorded as question-length arrays, in the
same list-indexed-by-turn shape as turn_cited_uris."""
from pydantic_evals import Case
from evaluations.benchmark import run_live_qa_benchmark
from evaluations.capability_runner import CapabilityRunResult
def build_case(idx: int, doc) -> Case:
return Case(
name="conv1",
inputs=["q1", "q2"],
metadata={
"conversation_id": "conv1",
"turns": [
{"reference": "r1", "answerability": "ANSWERABLE"},
{"reference": "r2", "answerability": "ANSWERABLE"},
],
},
)
spec = DatasetSpec(
key="test_live",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: None,
qa_loader=lambda: [{"id": "conv1"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=build_case,
live=True,
)
turn_results = [
CapabilityRunResult(
answer="a1",
cited_uris=["u1"],
n_search_calls=2,
n_rejected_searches=1,
n_failed_tools=1,
n_requests=4,
),
CapabilityRunResult(answer="a2"),
]
recorded: dict[str, object] = {}
with (
patch("evaluations.benchmark.get_model", return_value="fake-model"),
patch(
"evaluations.benchmark.set_eval_attribute",
side_effect=lambda key, value: recorded.__setitem__(key, value),
),
patch(
"evaluations.benchmark.run_capability_conversation",
new_callable=AsyncMock,
return_value=turn_results,
),
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
return_value=MagicMock(score=None, pass_=True, reason=None),
),
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
return_value=MagicMock(score=None, pass_=False, reason=None),
),
):
await run_live_qa_benchmark(
spec, AppConfig(), db_path=tmp_path / "test.lancedb"
)
assert recorded["turn_n_search_calls"] == [2, 0]
assert recorded["turn_n_rejected_searches"] == [1, 0]
assert recorded["turn_n_failed_tools"] == [1, 0]
assert recorded["turn_n_requests"] == [4, 0]
questions = 2
for key, value in recorded.items():
if key.startswith("turn_"):
assert isinstance(value, list) and len(value) == questions, key
class TestResolveDatasets:
def test_all_dedupes_shared_databases(self) -> None:
"""Specs sharing a db_filename (mtrag query variants) appear once, so
`download all`/`upload all` do not process the same DB twice."""
from evaluations.benchmark import _resolve_datasets
specs = _resolve_datasets("all")
filenames = [spec.db_filename for spec in specs]
assert len(filenames) == len(set(filenames))
assert "mtrag_clapnq.lancedb" in filenames
def test_single_key_not_deduped(self) -> None:
from evaluations.benchmark import _resolve_datasets
specs = _resolve_datasets("mtrag_clapnq_rewrite")
assert [spec.key for spec in specs] == ["mtrag_clapnq_rewrite"]
class TestLoadConfig:
def test_explicit_path(self, tmp_path: Path) -> None:
config_file = tmp_path / "test.yaml"
@ -383,17 +755,117 @@ class TestRunQaBenchmarkCapabilityTarget:
class TestCitationEvaluatorWiring:
def test_returns_map_twin_for_map_evaluator(self) -> None:
from evaluations.benchmark import _citation_evaluator_for
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
def test_specs_with_retrieval_declare_citation_evaluator(self) -> None:
"""Citation scoring is declared per spec, not inferred: every dataset
that scores retrieval also scores citations."""
from evaluations.datasets import DATASETS
from evaluations.evaluators import CitationMAPEvaluator
result = _citation_evaluator_for(MAPEvaluator())
assert isinstance(result, CitationMAPEvaluator)
for spec in DATASETS.values():
if spec.retrieval_evaluators:
assert isinstance(spec.citation_evaluator, CitationMAPEvaluator), (
spec.key
)
def test_returns_none_for_no_evaluator(self) -> None:
from evaluations.benchmark import _citation_evaluator_for
assert _citation_evaluator_for(None) is None
class TestBatchedIngest:
def _rag(
self,
complete_uris: list[str] | None = None,
chunkless_uris: list[str] | None = None,
) -> MagicMock:
complete_uris = complete_uris or []
chunkless_uris = chunkless_uris or []
def _table(rows: list[dict]) -> MagicMock:
table = MagicMock()
table.query.return_value.select.return_value.to_list = AsyncMock(
return_value=rows
)
return table
rag = MagicMock()
rag.store.document_meta_table = _table(
[{"id": f"id-{uri}", "uri": uri} for uri in complete_uris + chunkless_uris]
)
rag.store.chunks_table = _table(
[{"document_id": f"id-{uri}"} for uri in complete_uris]
)
rag.convert = AsyncMock(side_effect=lambda content, **kw: f"docling:{content}")
rag.chunk = AsyncMock(return_value=[])
rag.import_documents = AsyncMock()
rag.delete_document = AsyncMock()
return rag
def _spec(self) -> DatasetSpec:
return DatasetSpec(
key="test",
db_filename="test.lancedb",
document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
document_mapper=lambda doc: (
None
if doc["uri"] == "bad"
else DocumentPayload(uri=doc["uri"], content=f"text {doc['uri']}")
),
qa_loader=lambda: [], # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
)
@pytest.mark.asyncio
async def test_imports_in_bounded_batches(self) -> None:
from evaluations.benchmark import _ingest_batched
rag = self._rag()
corpus = [{"uri": f"u{i}"} for i in range(5)]
await _ingest_batched(rag, self._spec(), corpus, batch_size=2)
batch_uris = [
[imp.uri for imp in call.args[0]]
for call in rag.import_documents.call_args_list
]
assert batch_uris == [["u0", "u1"], ["u2", "u3"], ["u4"]]
@pytest.mark.asyncio
async def test_resume_skips_complete_uris(self) -> None:
from evaluations.benchmark import _ingest_batched
rag = self._rag(complete_uris=["u0", "u2"])
corpus = [{"uri": f"u{i}"} for i in range(4)]
await _ingest_batched(rag, self._spec(), corpus, batch_size=10)
(batch,), _ = rag.import_documents.call_args
assert [imp.uri for imp in batch] == ["u1", "u3"]
assert rag.convert.await_count == 2
rag.delete_document.assert_not_awaited()
@pytest.mark.asyncio
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
rag = self._rag(complete_uris=["u0"], chunkless_uris=["u1"])
corpus = [{"uri": "u0"}, {"uri": "u1"}]
await _ingest_batched(rag, self._spec(), corpus, batch_size=10)
rag.delete_document.assert_awaited_once_with("id-u1")
(batch,), _ = rag.import_documents.call_args
assert [imp.uri for imp in batch] == ["u1"]
@pytest.mark.asyncio
async def test_unmapped_documents_skipped(self) -> None:
from evaluations.benchmark import _ingest_batched
rag = self._rag()
corpus = [{"uri": "u0"}, {"uri": "bad"}, {"uri": "u1"}]
await _ingest_batched(rag, self._spec(), corpus, batch_size=10)
(batch,), _ = rag.import_documents.call_args
assert [imp.uri for imp in batch] == ["u0", "u1"]
class TestAttachRelevantUris:
@ -433,7 +905,7 @@ class TestAttachRelevantUris:
retrieval_mapper=lambda d: RetrievalSample(
question=d["q"], expected_uris=d["uris"]
),
retrieval_evaluator=MAPEvaluator(),
retrieval_evaluators=[MAPEvaluator()],
)
_attach_relevant_uris(cases, spec, limit=None)
@ -511,7 +983,7 @@ class TestRetrievalTarget:
retrieval_mapper=lambda d: RetrievalSample(
question=d["q"], expected_uris=d["uris"]
),
retrieval_evaluator=MAPEvaluator(),
retrieval_evaluators=[MAPEvaluator()],
)
@pytest.mark.asyncio

View file

@ -144,7 +144,9 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp
with patch(
"evaluations.capability_runner.Agent.run", new_callable=AsyncMock
) as run:
run.return_value = SimpleNamespace(output="done", all_messages=lambda: [])
run.return_value = SimpleNamespace(
output="done", all_messages=lambda: [], new_messages=lambda: []
)
await run_capability_question(
lambda **_kwargs: capability,
@ -157,3 +159,139 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp
assert capability.request_limit == expected
assert "usage_limits" not in run.call_args.kwargs
class TestPrefixToMessages:
def test_maps_turns_to_model_messages(self) -> None:
from pydantic_ai.messages import (
ModelRequest,
ModelResponse,
TextPart,
UserPromptPart,
)
from evaluations.capability_runner import prefix_to_messages
from evaluations.config import Turn
messages = prefix_to_messages(
[
Turn(speaker="user", text="who takes photos of planes?"),
Turn(speaker="agent", text="Ground-to-air photographers."),
]
)
assert len(messages) == 2
assert isinstance(messages[0], ModelRequest)
assert isinstance(messages[0].parts[0], UserPromptPart)
assert messages[0].parts[0].content == "who takes photos of planes?"
assert isinstance(messages[1], ModelResponse)
assert isinstance(messages[1].parts[0], TextPart)
assert messages[1].parts[0].content == "Ground-to-air photographers."
def test_empty_prefix(self) -> None:
from evaluations.capability_runner import prefix_to_messages
assert prefix_to_messages([]) == []
async def test_message_history_passed_to_agent_run(tmp_path):
from evaluations.capability_runner import prefix_to_messages
from evaluations.config import Turn
history = prefix_to_messages([Turn(speaker="user", text="earlier question")])
capability = create_rag(
db_path=tmp_path / "rag.lancedb",
config=AppConfig(),
defer_loading=False,
)
with patch(
"evaluations.capability_runner.Agent.run", new_callable=AsyncMock
) as run:
run.return_value = SimpleNamespace(output="done", new_messages=lambda: [])
await run_capability_question(
lambda **_kwargs: capability,
tmp_path / "rag.lancedb",
AppConfig(),
"follow-up question",
TestModel(call_tools=[]),
message_history=history,
)
assert run.call_args.kwargs["message_history"] is history
async def test_conversation_threads_own_messages_across_turns(tmp_path):
"""Each turn runs with the previous turn's full message history (including
tool traffic), so prior-turn compaction operates on real history."""
from evaluations.capability_runner import run_capability_conversation
capability = create_rag(
db_path=tmp_path / "rag.lancedb",
config=AppConfig(),
defer_loading=False,
)
histories: list[object] = []
async def _run(question, deps=None, message_history=None):
histories.append(message_history)
return SimpleNamespace(
output=f"answer to {question}",
all_messages=lambda: [f"history after {question}"],
new_messages=lambda: [],
)
with patch("evaluations.capability_runner.Agent.run", side_effect=_run):
result = await run_capability_conversation(
lambda **_kwargs: capability,
tmp_path / "rag.lancedb",
AppConfig(),
["q1", "q2", "q3"],
TestModel(call_tools=[]),
)
assert [t.answer for t in result] == [
"answer to q1",
"answer to q2",
"answer to q3",
]
assert histories == [None, ["history after q1"], ["history after q2"]]
async def test_conversation_end_to_end_with_test_model(tmp_path):
from evaluations.capability_runner import run_capability_conversation
result = await run_capability_conversation(
create_rag,
tmp_path / "rag.lancedb",
AppConfig(),
["first question", "follow-up"],
TestModel(call_tools=[]),
)
assert len(result) == 2
assert all(turn.answer == "success (no tool calls)" for turn in result)
assert all(turn.cited_uris == [] for turn in result)
async def test_gold_prefix_run_answers_with_history(tmp_path):
"""End-to-end through a real Agent: the prefix rides along as history."""
from evaluations.capability_runner import prefix_to_messages
from evaluations.config import Turn
history = prefix_to_messages(
[
Turn(speaker="user", text="who takes photos of planes?"),
Turn(speaker="agent", text="Ground-to-air photographers."),
]
)
result = await run_capability_question(
create_rag,
tmp_path / "rag.lancedb",
AppConfig(),
"No, I meant photos in the air.",
TestModel(call_tools=[]),
message_history=history,
)
assert result.answer == "success (no tool calls)"

View file

@ -28,17 +28,25 @@ class TestCitationMAPEvaluator:
def test_no_matches(self) -> None:
assert self.evaluator.evaluate(_ctx(["x", "y"], ["a", "b"])) == 0.0
def test_no_relevant(self) -> None:
assert self.evaluator.evaluate(_ctx(["a"], [])) == 0.0
def test_no_citations(self) -> None:
assert self.evaluator.evaluate(_ctx([], ["a"])) == 0.0
def test_metadata_none(self) -> None:
def test_ineligible_when_no_relevant_uris(self) -> None:
"""Turns without gold passages (unanswerable) produce no score at all,
not a penalizing zero."""
assert self.evaluator.evaluate(_ctx(["a"], [])) == {}
def test_ineligible_when_relevant_uris_missing(self) -> None:
ctx = MagicMock()
ctx.metadata = {"answerability": "UNANSWERABLE"}
ctx.attributes = {"cited_uris": ["a"]}
assert self.evaluator.evaluate(ctx) == {}
def test_ineligible_when_metadata_none(self) -> None:
ctx = MagicMock()
ctx.metadata = None
ctx.attributes = {"cited_uris": ["a"]}
assert self.evaluator.evaluate(ctx) == 0.0
assert self.evaluator.evaluate(ctx) == {}
def test_evaluation_name(self) -> None:
assert self.evaluator.get_default_evaluation_name() == "cited_map"

View file

@ -1,7 +1,16 @@
from pathlib import Path
from unittest.mock import patch
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
import pytest
from pydantic import ValidationError
from evaluations.config import (
ConversationInput,
DatasetSpec,
DocumentPayload,
RetrievalSample,
Turn,
)
def _make_spec(**kwargs: object) -> DatasetSpec:
@ -49,8 +58,53 @@ class TestDatasetSpecDefaults:
spec = _make_spec()
assert spec.retrieval_loader is None
assert spec.retrieval_mapper is None
assert spec.retrieval_evaluator is None
assert spec.retrieval_evaluators is None
assert spec.citation_evaluator is None
assert spec.document_limit is None
assert spec.retrieval_limit == 5
class TestConversationInput:
def _conversation(self) -> ConversationInput:
return ConversationInput(
turns=[
Turn(speaker="user", text="who takes photos of planes?"),
Turn(speaker="agent", text="Ground-to-air photographers."),
Turn(speaker="user", text="No, I meant photos in the air."),
]
)
def test_question_is_last_turn(self) -> None:
assert self._conversation().question == "No, I meant photos in the air."
def test_prefix_excludes_last_turn(self) -> None:
prefix = self._conversation().prefix
assert [t.speaker for t in prefix] == ["user", "agent"]
def test_transcript_renders_speaker_lines(self) -> None:
assert self._conversation().transcript == (
"user: who takes photos of planes?\n"
"agent: Ground-to-air photographers.\n"
"user: No, I meant photos in the air."
)
def test_single_turn_has_empty_prefix(self) -> None:
conversation = ConversationInput(turns=[Turn(speaker="user", text="hi")])
assert conversation.prefix == []
assert conversation.question == "hi"
def test_must_end_with_user_turn(self) -> None:
with pytest.raises(ValidationError, match="user turn"):
ConversationInput(
turns=[
Turn(speaker="user", text="q"),
Turn(speaker="agent", text="a"),
]
)
def test_must_have_turns(self) -> None:
with pytest.raises(ValidationError, match="user turn"):
ConversationInput(turns=[])
class TestDocumentPayload:

View file

@ -0,0 +1,257 @@
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic_evals.evaluators import EvaluatorContext
from pydantic_evals.evaluators.evaluator import EvaluationReason
from evaluations.evaluators.conversation import ConversationEvaluator
def _ctx(
questions: list[str],
answers: list[str],
turns: list[dict],
turn_cited_uris: list[list[str]] | None = None,
) -> EvaluatorContext:
return EvaluatorContext(
name="conv",
inputs=questions,
metadata={"conversation_id": "conv1", "turns": turns},
expected_output=None,
output=answers,
duration=0.0,
_span_tree=MagicMock(),
attributes={"turn_cited_uris": turn_cited_uris or [[] for _ in answers]},
metrics={},
)
def _grading(pass_: bool) -> MagicMock:
return MagicMock(score=None, pass_=pass_, reason=None)
class TestConversationEvaluator:
@pytest.mark.asyncio
async def test_per_turn_scores_and_aggregates(self) -> None:
evaluator = ConversationEvaluator(rubric="equivalence rubric", model="test")
ctx = _ctx(
questions=["q1", "q2", "q3"],
answers=["a1", "a2", "a3"],
turns=[
{
"reference": "r1",
"answerability": "ANSWERABLE",
"relevant_uris": ["p1", "p2"],
},
{"reference": "r2", "answerability": "UNANSWERABLE"},
{
"reference": "r3",
"answerability": "PARTIAL",
"relevant_uris": ["p3"],
},
],
turn_cited_uris=[["p1"], [], ["p3"]],
)
with (
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
side_effect=[_grading(True), _grading(False), _grading(True)],
) as judge_answer,
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
side_effect=[_grading(False), _grading(True)],
) as judge_refusal,
):
result = await evaluator.evaluate(ctx)
assert isinstance(result, dict)
assert result == {
"turn_pass_rate": pytest.approx(2 / 3),
"turns_passed": 2,
"turns_judged": 3,
"turns_total": 3,
"cited_map": pytest.approx((0.5 + 1.0) / 2),
"cited_eligible": 2,
"true_refusals": 1,
"false_refusals": 0,
"unanswerable_turns": 1,
"turn_1_pass": EvaluationReason(value=True, reason=None),
"turn_2_pass": EvaluationReason(value=False, reason=None),
"turn_3_pass": EvaluationReason(value=True, reason=None),
"turn_1_refused": False,
"turn_2_refused": True,
"turn_1_cited_ap": 0.5,
"turn_3_cited_ap": 1.0,
}
# Refusal judged only on ANSWERABLE/UNANSWERABLE turns.
assert judge_refusal.await_count == 2
assert judge_answer.await_count == 3
@pytest.mark.asyncio
async def test_judge_sees_live_transcript(self) -> None:
"""Turn 2 is judged against the conversation so far with OUR answer to
turn 1, not the reference."""
evaluator = ConversationEvaluator(rubric="rubric", model="test")
ctx = _ctx(
questions=["q1", "q2"],
answers=["my a1", "my a2"],
turns=[
{"reference": "r1", "answerability": "ANSWERABLE"},
{"reference": "r2", "answerability": "ANSWERABLE"},
],
)
with (
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
return_value=_grading(True),
) as judge_answer,
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
return_value=_grading(False),
),
):
await evaluator.evaluate(ctx)
second_call = judge_answer.await_args_list[1]
transcript, answer, reference = second_call.args[:3]
assert transcript == "user: q1\nagent: my a1\nuser: q2"
assert answer == "my a2"
assert reference == "r2"
@pytest.mark.asyncio
async def test_no_citation_scores_without_eligible_turns(self) -> None:
evaluator = ConversationEvaluator(rubric="rubric", model="test")
ctx = _ctx(
questions=["q1"],
answers=["a1"],
turns=[{"reference": "r1", "answerability": "UNANSWERABLE"}],
)
with (
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
return_value=_grading(False),
),
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
return_value=_grading(True),
),
):
result = await evaluator.evaluate(ctx)
assert result == {
"turn_pass_rate": 0.0,
"turns_passed": 0,
"turns_judged": 1,
"turns_total": 1,
"cited_eligible": 0,
"true_refusals": 1,
"false_refusals": 0,
"unanswerable_turns": 1,
"turn_1_pass": EvaluationReason(value=False, reason=None),
"turn_1_refused": True,
}
@pytest.mark.asyncio
async def test_mismatched_arrays_raise(self) -> None:
evaluator = ConversationEvaluator(rubric="rubric", model="test")
ctx = _ctx(
questions=["q1", "q2"],
answers=["a1"],
turns=[{"reference": "r1", "answerability": "ANSWERABLE"}],
)
with pytest.raises(ValueError, match="conversation arrays disagree"):
await evaluator.evaluate(ctx)
@pytest.mark.asyncio
async def test_judge_error_voids_one_turn_not_the_conversation(self) -> None:
evaluator = ConversationEvaluator(rubric="rubric", model="test")
ctx = _ctx(
questions=["q1", "q2", "q3"],
answers=["a1", "a2", "a3"],
turns=[
{"reference": "r1", "answerability": "ANSWERABLE"},
{"reference": "r2", "answerability": "ANSWERABLE"},
{"reference": "r3", "answerability": "ANSWERABLE"},
],
)
with (
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
side_effect=[
_grading(True),
RuntimeError("token limit exceeded"),
_grading(True),
],
),
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
return_value=_grading(False),
),
):
result = await evaluator.evaluate(ctx)
assert result == {
"turn_pass_rate": 1.0,
"turns_passed": 2,
"turns_judged": 2,
"turns_total": 3,
"cited_eligible": 0,
"true_refusals": 0,
"false_refusals": 0,
"unanswerable_turns": 0,
"turn_1_pass": EvaluationReason(value=True, reason=None),
"turn_3_pass": EvaluationReason(value=True, reason=None),
"turn_2_judge_error": "token limit exceeded",
"turn_1_refused": False,
"turn_2_refused": False,
"turn_3_refused": False,
}
@pytest.mark.asyncio
async def test_refusal_judge_error_skips_refusal_verdict_only(self) -> None:
evaluator = ConversationEvaluator(rubric="rubric", model="test")
ctx = _ctx(
questions=["q1"],
answers=["a1"],
turns=[{"reference": "r1", "answerability": "UNANSWERABLE"}],
)
with (
patch(
"evaluations.evaluators.conversation.judge_input_output_expected",
new_callable=AsyncMock,
return_value=_grading(True),
),
patch(
"evaluations.evaluators.conversation.judge_output",
new_callable=AsyncMock,
side_effect=RuntimeError("boom"),
),
):
result = await evaluator.evaluate(ctx)
assert result == {
"turn_pass_rate": 1.0,
"turns_passed": 1,
"turns_judged": 1,
"turns_total": 1,
"cited_eligible": 0,
"true_refusals": 0,
"false_refusals": 0,
"unanswerable_turns": 0,
"turn_1_pass": EvaluationReason(value=True, reason=None),
"turn_1_judge_error": "boom",
}

View file

@ -1,9 +1,14 @@
from unittest.mock import MagicMock
import math
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from pydantic_evals.evaluators import EvaluatorContext
from evaluations.evaluators import REFUSAL_RUBRIC, RefusalJudge
from evaluations.evaluators.map import MAPEvaluator
from evaluations.evaluators.number_match import NumberMatchEvaluator
from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator
class TestMAPEvaluator:
@ -57,6 +62,189 @@ class TestMAPEvaluator:
assert self.evaluator.evaluate(ctx) == 0.0
def _retrieval_ctx(relevant_uris: list[str], retrieved_uris: list[str]) -> MagicMock:
ctx = MagicMock()
ctx.metadata = {"relevant_uris": relevant_uris}
ctx.output = retrieved_uris
return ctx
class TestRecallEvaluator:
def test_evaluation_name_includes_k(self) -> None:
assert RecallEvaluator(k=5).get_default_evaluation_name() == "recall_5"
assert RecallEvaluator(k=10).get_default_evaluation_name() == "recall_10"
def test_all_relevant_within_k(self) -> None:
ctx = _retrieval_ctx(["a", "b"], ["a", "b", "c"])
assert RecallEvaluator(k=5).evaluate(ctx) == 1.0
def test_partial_recall(self) -> None:
ctx = _retrieval_ctx(["a", "b"], ["a", "c", "d"])
assert RecallEvaluator(k=3).evaluate(ctx) == 0.5
def test_relevant_beyond_k_not_counted(self) -> None:
ctx = _retrieval_ctx(["a"], ["b", "c", "d", "e", "f", "a"])
assert RecallEvaluator(k=5).evaluate(ctx) == 0.0
assert RecallEvaluator(k=10).evaluate(ctx) == 1.0
def test_empty_relevant(self) -> None:
ctx = _retrieval_ctx([], ["a"])
assert RecallEvaluator(k=5).evaluate(ctx) == 0.0
def test_none_metadata(self) -> None:
ctx = MagicMock()
ctx.metadata = None
ctx.output = ["a"]
assert RecallEvaluator(k=5).evaluate(ctx) == 0.0
class TestNDCGEvaluator:
def test_evaluation_name_includes_k(self) -> None:
assert NDCGEvaluator(k=5).get_default_evaluation_name() == "ndcg_5"
def test_perfect_ranking(self) -> None:
ctx = _retrieval_ctx(["a", "b"], ["a", "b", "c"])
assert NDCGEvaluator(k=5).evaluate(ctx) == pytest.approx(1.0)
def test_single_relevant_at_rank_two(self) -> None:
# DCG = 1/log2(3); IDCG = 1/log2(2) = 1
ctx = _retrieval_ctx(["a"], ["b", "a"])
expected = 1 / math.log2(3)
assert NDCGEvaluator(k=5).evaluate(ctx) == pytest.approx(expected)
def test_two_relevant_with_gap(self) -> None:
# Relevant at ranks 1 and 3: DCG = 1 + 1/log2(4) = 1.5
# IDCG = 1 + 1/log2(3)
ctx = _retrieval_ctx(["a", "b"], ["a", "c", "b"])
expected = 1.5 / (1 + 1 / math.log2(3))
assert NDCGEvaluator(k=3).evaluate(ctx) == pytest.approx(expected)
def test_relevant_beyond_k_not_counted(self) -> None:
ctx = _retrieval_ctx(["a"], ["b", "c", "d", "e", "f", "a"])
assert NDCGEvaluator(k=5).evaluate(ctx) == 0.0
def test_ideal_dcg_capped_at_k(self) -> None:
# 3 relevant but k=2: IDCG uses only the top-2 ideal ranks, so a
# retrieval with both top-2 slots relevant scores 1.0.
ctx = _retrieval_ctx(["a", "b", "c"], ["a", "b"])
assert NDCGEvaluator(k=2).evaluate(ctx) == pytest.approx(1.0)
def test_empty_relevant(self) -> None:
ctx = _retrieval_ctx([], ["a"])
assert NDCGEvaluator(k=5).evaluate(ctx) == 0.0
def test_none_metadata(self) -> None:
ctx = MagicMock()
ctx.metadata = None
ctx.output = ["a"]
assert NDCGEvaluator(k=5).evaluate(ctx) == 0.0
def _evaluator_ctx(inputs: object, metadata: dict | None = None) -> EvaluatorContext:
return EvaluatorContext(
name="case",
inputs=inputs,
metadata=metadata,
expected_output="expected",
output="answer",
duration=0.0,
_span_tree=MagicMock(),
attributes={},
metrics={},
)
class TestTranscriptLLMJudge:
@pytest.mark.asyncio
async def test_conversation_inputs_judged_as_transcript(self) -> None:
from evaluations.config import ConversationInput, Turn
from evaluations.evaluators import TranscriptLLMJudge
judge = TranscriptLLMJudge(
rubric="rubric",
include_input=True,
include_expected_output=True,
model="test",
)
conversation = ConversationInput(
turns=[
Turn(speaker="user", text="q1"),
Turn(speaker="agent", text="a1"),
Turn(speaker="user", text="q2"),
]
)
grading = MagicMock(score=None, pass_=True, reason="ok")
with patch(
"pydantic_evals.evaluators.llm_as_a_judge.judge_input_output_expected",
new_callable=AsyncMock,
return_value=grading,
) as judge_call:
await judge.evaluate(_evaluator_ctx(conversation))
assert judge_call.await_args is not None
assert judge_call.await_args.args[0] == "user: q1\nagent: a1\nuser: q2"
@pytest.mark.asyncio
async def test_string_inputs_pass_through(self) -> None:
from evaluations.evaluators import TranscriptLLMJudge
judge = TranscriptLLMJudge(
rubric="rubric",
include_input=True,
include_expected_output=True,
model="test",
)
grading = MagicMock(score=None, pass_=True, reason="ok")
with patch(
"pydantic_evals.evaluators.llm_as_a_judge.judge_input_output_expected",
new_callable=AsyncMock,
return_value=grading,
) as judge_call:
await judge.evaluate(_evaluator_ctx("plain question"))
assert judge_call.await_args is not None
assert judge_call.await_args.args[0] == "plain question"
class TestRefusalJudge:
def _judge(self) -> RefusalJudge:
return RefusalJudge(
rubric=REFUSAL_RUBRIC,
model="test",
assertion={"evaluation_name": "refused", "include_reason": False},
)
@pytest.mark.asyncio
@pytest.mark.parametrize("label", ["ANSWERABLE", "UNANSWERABLE"])
async def test_judges_eligible_labels(self, label: str) -> None:
grading = MagicMock(score=None, pass_=True, reason=None)
with patch(
"pydantic_evals.evaluators.llm_as_a_judge.judge_output",
new_callable=AsyncMock,
return_value=grading,
) as judge_call:
result = await self._judge().evaluate(
_evaluator_ctx("q", metadata={"answerability": label})
)
judge_call.assert_awaited_once()
assert result == {"refused": True}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"metadata", [{"answerability": "PARTIAL"}, {"answerability": None}, {}, None]
)
async def test_ineligible_turns_skip_the_judge(self, metadata) -> None:
with patch(
"pydantic_evals.evaluators.llm_as_a_judge.judge_output",
new_callable=AsyncMock,
) as judge_call:
result = await self._judge().evaluate(_evaluator_ctx("q", metadata))
judge_call.assert_not_awaited()
assert result == {}
class TestNumberMatchEvaluator:
def setup_method(self) -> None:
self.evaluator = NumberMatchEvaluator()

View file

@ -0,0 +1,268 @@
import pytest
from evaluations.config import ConversationInput
from evaluations.datasets import DATASETS
from evaluations.datasets.mtrag import (
MTRAG_CLAPNQ_LIVE_SPEC,
MTRAG_CLAPNQ_REWRITE_SPEC,
MTRAG_CLAPNQ_SPEC,
_group_conversations,
_join_queries_qrels,
_parse_qrels,
_task_to_record,
_validate_qrels_resolve,
build_mtrag_case,
build_mtrag_live_case,
map_mtrag_document,
map_mtrag_retrieval,
)
from evaluations.evaluators import (
CitationMAPEvaluator,
MAPEvaluator,
NDCGEvaluator,
RecallEvaluator,
)
GENERATION_TASK = {
"task_id": "conv1<::>2",
"conversation_id": "conv1",
"turn": "2",
"Collection": "mt-rag-clapnq-elser-512-100-20240503",
"Answerability": ["ANSWERABLE"],
"Multi-Turn": ["Follow-up"],
"Question Type": ["Factoid"],
"input": [
{"speaker": "user", "text": "q1", "metadata": {}},
{"speaker": "agent", "text": "a1", "metadata": {}},
{"speaker": "user", "text": "q2", "metadata": {}},
],
"targets": [{"text": "reference answer"}],
"contexts": [{"document_id": "retrieved-not-gold"}],
}
class TestDocumentMapper:
def test_maps_passage_to_payload(self) -> None:
payload = map_mtrag_document(
{"_id": "837799097_6931-7548-0-617", "title": "T", "text": "body"}
)
assert payload.uri == "837799097_6931-7548-0-617"
assert payload.title == "T"
assert payload.content == "body"
class TestQrels:
QRELS_TSV = (
"query-id\tcorpus-id\tscore\n"
"conv1<::>2\tdoc1_0-10-0-10\t1\n"
"conv1<::>2\tdoc2_5-20-0-15\t1\n"
"conv2<::>1\tdoc3_0-9-0-9\t1\n"
)
def test_parse_groups_by_query_preserving_order(self) -> None:
qrels = _parse_qrels(self.QRELS_TSV.splitlines())
assert qrels == {
"conv1<::>2": ["doc1_0-10-0-10", "doc2_5-20-0-15"],
"conv2<::>1": ["doc3_0-9-0-9"],
}
def test_join_builds_records(self) -> None:
qrels = _parse_qrels(self.QRELS_TSV.splitlines())
queries = [
{"_id": "conv1<::>2", "text": "q one"},
{"_id": "conv2<::>1", "text": "q two"},
]
records = _join_queries_qrels(queries, qrels)
assert records == [
{
"query_id": "conv1<::>2",
"question": "q one",
"expected_uris": ["doc1_0-10-0-10", "doc2_5-20-0-15"],
},
{
"query_id": "conv2<::>1",
"question": "q two",
"expected_uris": ["doc3_0-9-0-9"],
},
]
def test_join_raises_on_query_without_qrels(self) -> None:
with pytest.raises(ValueError, match="no qrels"):
_join_queries_qrels([{"_id": "missing<::>1", "text": "q"}], {})
def test_validation_passes_when_all_resolve(self) -> None:
qrels = {"q1": ["a", "b"]}
_validate_qrels_resolve({"a", "b", "c"}, qrels)
def test_validation_raises_on_unresolved_id(self) -> None:
qrels = {"q1": ["a", "ghost"]}
with pytest.raises(ValueError, match="ghost"):
_validate_qrels_resolve({"a"}, qrels)
class TestRetrievalMapper:
def test_maps_joined_record(self) -> None:
sample = map_mtrag_retrieval(
{
"query_id": "conv1<::>2",
"question": "who?",
"expected_uris": ["u1", "u2"],
}
)
assert sample is not None
assert sample.question == "who?"
assert sample.expected_uris == ("u1", "u2")
class TestSpecs:
def test_registered(self) -> None:
assert DATASETS["mtrag_clapnq"] is MTRAG_CLAPNQ_SPEC
assert DATASETS["mtrag_clapnq_rewrite"] is MTRAG_CLAPNQ_REWRITE_SPEC
def test_variants_share_db(self) -> None:
assert MTRAG_CLAPNQ_SPEC.db_filename == MTRAG_CLAPNQ_REWRITE_SPEC.db_filename
def test_retrieval_configuration(self) -> None:
for spec in (MTRAG_CLAPNQ_SPEC, MTRAG_CLAPNQ_REWRITE_SPEC):
assert spec.retrieval_limit == 10
assert spec.ingest_batch_size == 512
assert spec.retrieval_evaluators is not None
kinds = {
(type(e), getattr(e, "k", None)) for e in spec.retrieval_evaluators
}
assert kinds == {
(RecallEvaluator, 5),
(RecallEvaluator, 10),
(NDCGEvaluator, 5),
(NDCGEvaluator, 10),
(MAPEvaluator, None),
}
assert isinstance(spec.citation_evaluator, CitationMAPEvaluator)
def test_refusal_evaluation_enabled(self) -> None:
assert MTRAG_CLAPNQ_SPEC.evaluate_refusal is True
class TestGenerationTasks:
def test_task_to_record(self) -> None:
record = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p1", "p2"]})
assert record == {
"id": "conv1<::>2",
"turn": "2",
"turns": [
{"speaker": "user", "text": "q1"},
{"speaker": "agent", "text": "a1"},
{"speaker": "user", "text": "q2"},
],
"answer": "reference answer",
"answerability": "ANSWERABLE",
"multi_turn_type": "Follow-up",
"question_type": ["Factoid"],
"relevant_uris": ["p1", "p2"],
}
def test_task_without_qrels_has_no_relevant_uris(self) -> None:
record = _task_to_record(GENERATION_TASK, {})
assert record is not None
assert record["relevant_uris"] is None
def test_other_collections_excluded(self) -> None:
task = {**GENERATION_TASK, "Collection": "mt-rag-govt-elser-512-100-20240611"}
assert _task_to_record(task, {}) is None
def test_build_case_conversation_and_metadata(self) -> None:
record = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p1"]})
assert record is not None
case = build_mtrag_case(3, record)
assert isinstance(case.inputs, ConversationInput)
assert case.inputs.question == "q2"
assert [t.speaker for t in case.inputs.turns] == ["user", "agent", "user"]
assert case.expected_output == "reference answer"
assert case.metadata == {
"task_id": "conv1<::>2",
"turn": "2",
"answerability": "ANSWERABLE",
"multi_turn_type": "Follow-up",
"question_type": ["Factoid"],
"relevant_uris": ["p1"],
}
def test_build_case_omits_relevant_uris_when_absent(self) -> None:
record = _task_to_record(
{**GENERATION_TASK, "Answerability": ["UNANSWERABLE"]}, {}
)
assert record is not None
case = build_mtrag_case(1, record)
assert case.metadata is not None
assert "relevant_uris" not in case.metadata
assert case.metadata["answerability"] == "UNANSWERABLE"
class TestLiveConversations:
def _records(self) -> list[dict]:
turn1 = _task_to_record(
{
**GENERATION_TASK,
"task_id": "conv1<::>1",
"turn": "1",
"input": [{"speaker": "user", "text": "q1", "metadata": {}}],
"targets": [{"text": "r1"}],
},
{"conv1<::>1": ["p1"]},
)
turn2 = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p2", "p3"]})
other = _task_to_record(
{
**GENERATION_TASK,
"task_id": "conv2<::>1",
"turn": "1",
"input": [{"speaker": "user", "text": "other q", "metadata": {}}],
"targets": [{"text": "other r"}],
"Answerability": ["UNANSWERABLE"],
},
{},
)
assert turn1 and turn2 and other
# turn 2 first: grouping must sort turns numerically within a conversation
return [turn2, turn1, other]
def test_grouping_sorts_turns_within_conversations(self) -> None:
conversations = _group_conversations(self._records())
assert [c["id"] for c in conversations] == ["conv1", "conv2"]
conv1 = conversations[0]
assert [t["question"] for t in conv1["turns"]] == ["q1", "q2"]
assert [t["reference"] for t in conv1["turns"]] == ["r1", "reference answer"]
assert conv1["turns"][1]["relevant_uris"] == ["p2", "p3"]
def test_build_live_case(self) -> None:
conversations = _group_conversations(self._records())
case = build_mtrag_live_case(1, conversations[0])
assert case.inputs == ["q1", "q2"]
assert case.metadata is not None
assert case.metadata["conversation_id"] == "conv1"
turns = case.metadata["turns"]
assert turns[0] == {
"task_id": "conv1<::>1",
"turn": "1",
"reference": "r1",
"answerability": "ANSWERABLE",
"multi_turn_type": "Follow-up",
"question_type": ["Factoid"],
"relevant_uris": ["p1"],
}
other_case = build_mtrag_live_case(2, conversations[1])
assert other_case.metadata is not None
assert "relevant_uris" not in other_case.metadata["turns"][0]
def test_live_spec(self) -> None:
assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC
assert MTRAG_CLAPNQ_LIVE_SPEC.db_filename == MTRAG_CLAPNQ_SPEC.db_filename
assert MTRAG_CLAPNQ_LIVE_SPEC.live is True
assert MTRAG_CLAPNQ_LIVE_SPEC.retrieval_loader is None
assert MTRAG_CLAPNQ_LIVE_SPEC.experiment_metadata == {
"mtrag_mode": "live_session"
}
assert MTRAG_CLAPNQ_SPEC.experiment_metadata == {"mtrag_mode": "gold_prefix"}