Merge pull request #545 from ggozad/feat/mtrag-bench

MTRAG multi-turn evaluation with compaction arms

CHANGELOG.md

@@ -4,6 +4,11 @@
 ### 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` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes.
+
+### Changed
+
+- `import_documents` embeds chunks across the whole batch in one pass instead of per document.
 
 ### Removed
 

docs/benchmarks.md

@@ -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`, `mtrag_clapnq_live` and `mtrag_clapnq_live_uncompacted` 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,37 @@ 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.
+
+Four 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, tool history and capability state across turns, with `EvidenceCompactionCapability` registered. `mtrag_clapnq_live_uncompacted` is the same replay without compaction, isolating what compaction contributes. This is the only multi-turn evaluation, so it is the only one where compaction acts at all.
+
+##### 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:Muse-Glimmer-30B-NVFP4` | 224 | 0.76 | 0.35 |
+| Live compacted (`mtrag_clapnq_live`) | `vllm:Muse-Glimmer-30B-NVFP4` | 224/224 scored | 0.83 micro / 0.84 macro | 0.42 |
+| Live uncompacted (`mtrag_clapnq_live_uncompacted`) | `vllm:Muse-Glimmer-30B-NVFP4` | 224/224 scored | 0.78 micro / 0.79 macro | 0.42 |
+
+*Measured on haiku.rag v0.74.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `Qwen3-Reranker-4B`, stock capability instructions with `reasoning_strength: high`, judged by the pinned `vllm:Qwen3.6-35B-A3B-NVFP4` (temperature 0.6, thinking). QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Gold-prefix and live rates answer different judge questions and are not comparable with each other. The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.*
+
+The two live arms replay the same 29 conversations (224 turns) and differ only in registering `EvidenceCompactionCapability`, so they are compared as paired observations:
+
+- Input tokens per model request, computed as total input tokens divided by model requests across the whole arm: 7,461 compacted (5,207,627 tokens over 698 requests) vs 13,539 uncompacted (9,707,530 over 717 requests). The uncompacted arm used 1.81x as many tokens per request, a 44.9% reduction under compaction.
+- Answer pass rate: 185/224 vs 175/224 turns. Of the 18 turns where the arms disagree, 14 pass only compacted and 4 only uncompacted. McNemar exact two-sided p = 0.031. The paired difference is +4.5pp with a Wald 95% CI of +0.8 to +8.1pp, so the honest claim is an improvement of roughly 1 to 8 points, not the point estimate.
+- Citation MAP, macro-averaged over conversations with 208 of 224 turns eligible (turns with gold passages) in each arm: 0.4174 compacted vs 0.4230 uncompacted. The gold-prefix 0.35 is over 208 of 224 eligible cases.
+- Refusal precision and recall against the answerability labels (16 UNANSWERABLE turns per arm): compacted 0.33 precision and 0.44 recall (21 refusals), uncompacted 0.23 and 0.31 (22 refusals). Gold-prefix: 0.24 and 0.44 (29 refusals).

evaluations/README.md

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

evaluations/configs/mtrag_clapnq.yaml

@@ -0,0 +1,56 @@
+# 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, mtrag_clapnq_live and
+# mtrag_clapnq_live_uncompacted.
+# 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: RedHatAI/Muse-Glimmer-30B-NVFP4
+    base_url: http://vllm:11450/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
+    extra_body:
+      chat_template_kwargs:
+        # Part of the measured baseline. vLLM's reasoning parser consumes
+        # enable_thinking before the chat template sees it; reasoning_strength
+        # is the knob Muse Glimmer templates honour, and a template that
+        # defaults it to low silently changes search behavior.
+        reasoning_strength: high
+
+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

evaluations/evaluations/benchmark.py

@@ -1,27 +1,36 @@
 import asyncio
 import shutil
-from collections.abc import Awaitable, Callable, Mapping
+from collections.abc import Callable, Mapping
 from pathlib import Path
-from typing import Any, Literal, cast
+from typing import Any, Literal, NamedTuple, cast
 
 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_ELIGIBLE_LABELS,
+    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 +127,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 +198,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 +306,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 +320,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 +353,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 +383,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,
@@ -332,6 +395,8 @@ def _attach_relevant_uris(
     """
     if spec.retrieval_loader is None or spec.retrieval_mapper is None:
         return
+    if not any(isinstance(case.inputs, str) for case in cases):
+        return
     corpus = spec.retrieval_loader()
     if limit is not None:
         corpus = corpus.select(range(min(limit, len(corpus))))
@@ -342,6 +407,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 +417,108 @@ 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") for case in scored)
+    turns_passed = sum(_score(case, "turns_passed") for case in scored)
+    # A conversation with zero judged turns (its judge calls all failed)
+    # reports turn_pass_rate 0.0; averaging that in would count a judge
+    # outage as a failed conversation, against the exclusion policy.
+    judged = [case for case in scored if _score(case, "turns_judged")]
+    summary: dict[str, float | int] = {
+        "conversations": len(scored),
+        "conversations_attempted": len(report_cases) + len(report_failures),
+        "turns_total": turns_total,
+        "turns_judged": turns_judged,
+        "turns_attempted": turns_total + failed_turns,
+        "micro_pass_rate": turns_passed / turns_judged if turns_judged else 0.0,
+        "macro_pass_rate": sum(_score(case, "turn_pass_rate") for case in judged)
+        / len(judged)
+        if judged
+        else 0.0,
+    }
+
+    cited = [case for case in scored if _score(case, "cited_map") is not None]
+    eligible = sum(_score(case, "cited_eligible") for case in scored)
+    if cited and eligible:
+        summary["cited_eligible"] = eligible
+        summary["cited_map_micro"] = (
+            sum(
+                _score(case, "cited_map") * _score(case, "cited_eligible")
+                for case in cited
+            )
+            / eligible
+        )
+        summary["cited_map_macro"] = sum(
+            _score(case, "cited_map") for case in cited
+        ) / len(cited)
+
+    true_refusals = sum(_score(case, "true_refusals") for case in scored)
+    false_refusals = sum(_score(case, "false_refusals") for case in scored)
+    unanswerable = sum(_score(case, "unanswerable_turns") for case in scored)
+    refusals = true_refusals + false_refusals
+    summary["unanswerable_turns"] = unanswerable
+    summary["refusals"] = refusals
+    summary["refusal_precision"] = true_refusals / refusals if refusals else 0.0
+    summary["refusal_recall"] = true_refusals / unanswerable if unanswerable else 0.0
+    return summary
+
+
+def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None:
+    """Refusal precision/recall against answerability labels.
+
+    Uses cases the refusal judge scored (ANSWERABLE/UNANSWERABLE turns).
+    Returns (precision, recall, unanswerable_count, refusal_count), or None
+    when no case was judged.
+    """
+    outcomes: list[tuple[str, bool]] = []
+    for case in report_cases:
+        refused = case.assertions.get("refused")
+        label = (case.metadata or {}).get("answerability")
+        if refused is None or label not in REFUSAL_ELIGIBLE_LABELS:
+            continue
+        outcomes.append((label, bool(refused.value)))
+    if not outcomes:
+        return None
+    refusals = [(label, r) for label, r in outcomes if r]
+    true_refusals = sum(1 for label, _ in refusals if label == "UNANSWERABLE")
+    unanswerable = sum(1 for label, _ in outcomes if label == "UNANSWERABLE")
+    precision = true_refusals / len(refusals) if refusals else 0.0
+    recall = true_refusals / unanswerable if unanswerable else 0.0
+    return precision, recall, unanswerable, len(refusals)
+
+
 def _filter_qa_corpus(corpus, case_ids: set[str] | None):
     """Keep only rows whose ``id`` is in ``case_ids`` (failure-subset reruns).
 
@@ -360,6 +529,84 @@ def _filter_qa_corpus(corpus, case_ids: set[str] | None):
     return corpus.filter(lambda row: row.get("id") in case_ids)
 
 
+class _QARun(NamedTuple):
+    cases: list[Case[Any, Any, dict[str, Any]]]
+    db: Path
+    judge_config: ModelConfig
+    eval_name: str
+    experiment_metadata: dict[str, Any]
+    capability_factory: CapabilityFactory
+    capability_model: Any
+
+
+def _prepare_qa_run(
+    spec: DatasetSpec,
+    config: AppConfig,
+    limit: int | None,
+    name: str | None,
+    db_path: Path | None,
+    judge_model: ModelConfig | None,
+    target: Target,
+    capability_model: ModelConfig | None,
+    case_ids: set[str] | None,
+    document_filter: str | None,
+) -> _QARun:
+    """Shared setup for the QA runners: cases, models, name and metadata."""
+    corpus = spec.qa_loader()
+    corpus = _filter_qa_corpus(corpus, case_ids)
+    if limit is not None:
+        corpus = corpus.select(range(min(limit, len(corpus))))
+
+    cases = [
+        spec.qa_case_builder(index, cast(Mapping[str, Any], doc))
+        for index, doc in enumerate(corpus, start=1)
+    ]
+
+    judge_config = judge_model or DEFAULT_JUDGE_MODEL
+    capability_config = _resolve_capability_config(target, config, capability_model)
+
+    eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
+    experiment_metadata = build_experiment_metadata(
+        dataset_key=spec.key,
+        test_cases=len(cases),
+        config=config,
+        judge_config=judge_config,
+        target=target,
+        capability_config=capability_config,
+        document_filter=document_filter,
+    )
+    experiment_metadata.update(spec.experiment_metadata or {})
+
+    return _QARun(
+        cases=cases,
+        db=spec.db_path(db_path),
+        judge_config=judge_config,
+        eval_name=eval_name,
+        experiment_metadata=experiment_metadata,
+        capability_factory=_capability_factory_for_target(target),
+        capability_model=get_model(capability_config, config),
+    )
+
+
+def _print_mean_task_time(report_cases, unit: str = "case") -> None:
+    if not report_cases:
+        return
+    mean = sum(case.task_duration for case in report_cases) / len(report_cases)
+    console.print(f"Avg task time per {unit}: {mean:.2f}s")
+
+
+def _print_failures(failures, show_question: bool = False) -> None:
+    if not failures:
+        return
+    console.print("[red]\nSummary of failures:[/red]")
+    for failure in failures:
+        console.print(f"Case: {failure.name}")
+        if show_question:
+            console.print(f"Question: {failure.inputs}")
+        console.print(f"Error: {failure.error_message}")
+        console.print("")
+
+
 async def run_qa_benchmark(
     spec: DatasetSpec,
     config: AppConfig,
@@ -372,27 +619,22 @@ async def run_qa_benchmark(
     case_ids: set[str] | None = None,
     document_filter: str | None = None,
 ) -> ReportCaseFailure[str, str, dict[str, str]] | None:
-    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
-    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
-    db = spec.db_path(db_path)
+    run = _prepare_qa_run(
+        spec,
+        config,
+        limit,
+        name,
+        db_path,
+        judge_model,
+        target,
+        capability_model,
+        case_ids,
+        document_filter,
+    )
+    cases, judge_config = run.cases, run.judge_config
 
     _attach_relevant_uris(cases, spec, limit)
-    citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator)
+    citation_evaluator = spec.citation_evaluator
 
     qa_evaluator = spec.qa_evaluator
     evaluators: list[Evaluator]
@@ -400,7 +642,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,42 +655,35 @@ async def run_qa_benchmark(
         ]
     if citation_evaluator is not None:
         evaluators.append(citation_evaluator)
+    # RefusalJudge scores only cases whose metadata carries an answerability
+    # label; on unlabeled datasets it returns no score without a judge call.
+    evaluators.append(
+        RefusalJudge(
+            rubric=REFUSAL_RUBRIC,
+            model=get_model(judge_config, config),
+            assertion={"evaluation_name": "refused", "include_reason": False},
+        )
+    )
 
-    evaluation_dataset = EvalDataset[str, str, dict[str, str]](
+    evaluation_dataset = EvalDataset[Any, str, dict[str, Any]](
         name=spec.key, cases=cases, evaluators=evaluators
     )
 
-    eval_name = name if name is not None else f"{spec.key}_qa_evaluation"
-    experiment_metadata = build_experiment_metadata(
-        dataset_key=spec.key,
-        test_cases=len(cases),
-        config=config,
-        judge_config=judge_config,
-        target=target,
-        capability_config=capability_config,
-        document_filter=document_filter,
-    )
-
-    async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]):
-        return await evaluation_dataset.evaluate(
-            answer_fn,
-            name=eval_name,
-            max_concurrency=1,
-            progress=True,
-            metadata=experiment_metadata,
-        )
-
-    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,
+            capability_factory=run.capability_factory,
+            db_path=run.db,
             config=config,
             question=question,
-            capability_model=resolved_capability_model,
+            capability_model=run.capability_model,
             document_filter=document_filter,
+            message_history=message_history,
         )
         set_eval_attribute("cited_uris", result.cited_uris)
         set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids)
@@ -459,9 +694,16 @@ async def run_qa_benchmark(
         set_eval_attribute("n_failed_tools", result.n_failed_tools)
         set_eval_attribute("n_executions", result.n_executions)
         set_eval_attribute("n_requests", result.n_requests)
+        set_eval_attribute("citation_status", result.citation_status)
         return result.answer
 
-    report = await _evaluate(answer_question)
+    report = await evaluation_dataset.evaluate(
+        answer_question,
+        name=run.eval_name,
+        max_concurrency=1,
+        progress=True,
+        metadata=run.experiment_metadata,
+    )
 
     total_processed = len(report.cases)
     failures = report.failures
@@ -488,6 +730,7 @@ 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}%)")
+    _print_mean_task_time(report.cases)
 
     if citation_evaluator is not None:
         score_key = citation_evaluator.get_default_evaluation_name()
@@ -508,22 +751,150 @@ 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 failures:
-        console.print("[red]\nSummary of failures:[/red]")
-        for failure in failures:
-            console.print(f"Case: {failure.name}")
-            console.print(f"Question: {failure.inputs}")
-            console.print(f"Error: {failure.error_message}")
-            console.print("")
+    if (metrics := _refusal_metrics(report.cases)) is not None:
+        precision, recall, unanswerable, refusals = metrics
+        console.print("\n=== Refusal vs answerability labels ===", style="bold cyan")
+        console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}")
+        console.print(
+            f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} "
+            "(PARTIAL excluded)"
+        )
+
+    _print_failures(failures, show_question=True)
 
     return failures[0] if failures else None
 
 
+async def run_live_qa_benchmark(
+    spec: DatasetSpec,
+    config: AppConfig,
+    limit: int | None = None,
+    name: str | None = None,
+    db_path: Path | None = None,
+    judge_model: ModelConfig | None = None,
+    target: Target = "rag-capability",
+    capability_model: ModelConfig | None = None,
+    case_ids: set[str] | None = None,
+    document_filter: str | None = None,
+) -> None:
+    """Replay conversations turn by turn through one capability session.
+
+    One case per conversation; ``limit`` counts conversations. Answers carry
+    forward as real message history, so prior-turn compaction is exercised.
+    """
+    run = _prepare_qa_run(
+        spec,
+        config,
+        limit,
+        name,
+        db_path,
+        judge_model,
+        target,
+        capability_model,
+        case_ids,
+        document_filter,
+    )
+
+    evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]](
+        name=spec.key,
+        cases=run.cases,
+        evaluators=[
+            ConversationEvaluator(
+                rubric=ANSWER_EQUIVALENCE_RUBRIC,
+                model=get_model(run.judge_config, config),
+            )
+        ],
+    )
+
+    async def answer_conversation(questions: list[str]) -> list[str]:
+        results = await run_capability_conversation(
+            capability_factory=run.capability_factory,
+            db_path=run.db,
+            config=config,
+            questions=list(questions),
+            capability_model=run.capability_model,
+            document_filter=document_filter,
+            compaction=spec.compaction,
+        )
+        set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results])
+        set_eval_attribute("turn_n_search_calls", [r.n_search_calls for r in results])
+        set_eval_attribute(
+            "turn_n_rejected_searches", [r.n_rejected_searches for r in results]
+        )
+        set_eval_attribute("turn_n_failed_tools", [r.n_failed_tools for r in results])
+        set_eval_attribute("turn_n_requests", [r.n_requests for r in results])
+        set_eval_attribute("turn_citation_status", [r.citation_status for r in results])
+        return [r.answer for r in results]
+
+    report = await evaluation_dataset.evaluate(
+        answer_conversation,
+        name=run.eval_name,
+        max_concurrency=1,
+        progress=True,
+        metadata=run.experiment_metadata,
+    )
+
+    summary = _live_summary(report.cases, report.failures)
+    console.print("\n=== Live Conversation Results ===", style="bold cyan")
+    if summary is None:
+        attempted = len(report.cases) + len(report.failures)
+        console.print(f"No conversations were scored ({attempted} attempted).")
+    else:
+        console.print(
+            f"Conversations scored: {summary['conversations']}"
+            f"/{summary['conversations_attempted']} | turns scored: "
+            f"{summary['turns_total']}/{summary['turns_attempted']}"
+        )
+        if summary["turns_judged"] < summary["turns_total"]:
+            console.print(
+                f"Turns judged: {summary['turns_judged']}/{summary['turns_total']} "
+                "(per-turn judge errors excluded from rates)"
+            )
+        if report.failures:
+            console.print(
+                "Failed conversations are operational exclusions — "
+                "not counted as wrong answers."
+            )
+        console.print(
+            f"Answer pass rate — micro (per turn): {summary['micro_pass_rate']:.4f} | "
+            f"macro (per conversation): {summary['macro_pass_rate']:.4f}"
+        )
+        if "cited_map_micro" in summary:
+            console.print(
+                f"cited_map — micro: {summary['cited_map_micro']:.4f} | "
+                f"macro: {summary['cited_map_macro']:.4f} "
+                f"(eligible turns: {summary['cited_eligible']})"
+            )
+        console.print(
+            f"Refusal precision: {summary['refusal_precision']:.2%} | "
+            f"recall: {summary['refusal_recall']:.2%} "
+            f"(UNANSWERABLE turns: {summary['unanswerable_turns']}, "
+            f"refusals: {summary['refusals']})"
+        )
+    if report.cases:
+        mean_task_time = sum(case.task_duration for case in report.cases) / len(
+            report.cases
+        )
+        turns = sum(len(case.output or []) for case in report.cases)
+        per_turn = (
+            sum(case.task_duration for case in report.cases) / turns if turns else 0.0
+        )
+        console.print(
+            f"Avg task time: {mean_task_time:.2f}s per conversation | "
+            f"{per_turn:.2f}s per turn"
+        )
+
+    _print_failures(report.failures)
+
+
 async def evaluate_dataset(
     spec: DatasetSpec,
     config: AppConfig,
@@ -566,7 +937,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 +993,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)]
 
 

evaluations/evaluations/capability_runner.py

@@ -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,14 +6,22 @@ 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 pydantic_ai.capabilities import AbstractCapability
+
+from evaluations.config import Turn
 from haiku.rag.capabilities import RAGCapabilityBase
+from haiku.rag.capabilities.compaction import create_capability as create_compaction
+from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, citation_status
 from haiku.rag.config.models import AppConfig
 from haiku.rag.store.models.chunk import SearchResult
 from haiku.rag.store.models.citation import Citation
@@ -21,10 +29,22 @@ 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]
     citations: list[str]
+    evidence: CapabilityEvidenceRecord
     searches: dict[str, list[SearchResult]]
 
 
@@ -40,6 +60,7 @@ class CapabilityRunResult:
     n_rejected_searches: int = 0
     n_failed_tools: int = 0
     n_requests: int = 0
+    citation_status: str | None = None
 
 
 class ToolTraffic(NamedTuple):
@@ -107,25 +128,15 @@ 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,
+    compaction: bool = False,
+) -> tuple[RAGCapabilityBase[Any], _EvalDeps, Agent[_EvalDeps, str]]:
     capability = capability_factory(
         db_path=db_path,
         config=config,
@@ -138,16 +149,110 @@ async def run_capability_question(
     if document_filter is not None:
         typed.document_filter = document_filter
 
+    capabilities: list[AbstractCapability] = [capability]
+    if compaction:
+        capabilities.append(create_compaction())
     deps = _EvalDeps(state={capability.state_namespace: state.model_dump(mode="json")})
     agent = Agent(
         capability_model,
         deps_type=_EvalDeps,
-        capabilities=[capability],
+        capabilities=capabilities,
     )
-    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,
+    compaction: bool = False,
+) -> 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) and the same state dict, which is what
+    lets ``EvidenceCompactionCapability`` (registered when ``compaction`` is
+    True) replace earlier questions' evidence on the request. 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=document_filter,
+        request_limit=None,
+        compaction=compaction,
+    )
+    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 +273,18 @@ 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,
+    record = typed.evidence
+    status = (
+        citation_status([record], question=record.question)
+        if record.question is not None
+        else None
     )
 
     return CapabilityRunResult(
-        answer=agent_result.output,
+        answer=answer,
         cited_uris=cited_uris,
         cited_chunk_ids=cited_chunk_ids,
         searched_uris=searched_uris,
@@ -192,4 +298,5 @@ async def run_capability_question(
         n_rejected_searches=traffic.n_rejected_searches,
         n_failed_tools=traffic.n_failed_tools,
         n_requests=traffic.n_requests,
+        citation_status=status,
     )

evaluations/evaluations/config.py

@@ -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,7 @@ 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]]]
+CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]]
 
 
 @dataclass
@@ -43,9 +73,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
+    live: bool = False
+    compaction: bool = False
+    experiment_metadata: dict[str, Any] | None = None
 
     def db_path(self, override_path: Path | None = None) -> Path:
         """Get the database path.

evaluations/evaluations/datasets/__init__.py

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

evaluations/evaluations/datasets/hotpotqa.py

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

evaluations/evaluations/datasets/mtrag.py

@@ -0,0 +1,304 @@
+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 _qa_records() -> list[dict[str, Any]]:
+    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 records
+
+
+def load_clapnq_qa() -> Dataset:
+    return Dataset.from_list(_qa_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"]),
+                "relevant_uris": list(task["relevant_uris"] or []),
+            }
+            turns.append(turn)
+        conversations.append({"id": conversation_id, "turns": turns})
+    return conversations
+
+
+def load_clapnq_conversations() -> Dataset:
+    return Dataset.from_list(_group_conversations(_qa_records()))
+
+
+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"}
+        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,
+        experiment_metadata={"mtrag_mode": "gold_prefix"},
+    )
+
+
+MTRAG_CLAPNQ_SPEC = _mtrag_spec("mtrag_clapnq", "lastturn")
+MTRAG_CLAPNQ_REWRITE_SPEC = _mtrag_spec("mtrag_clapnq_rewrite", "rewrite")
+
+
+def _mtrag_live_spec(key: str, compaction: bool) -> DatasetSpec:
+    return DatasetSpec(
+        key=key,
+        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,
+        compaction=compaction,
+        experiment_metadata={"mtrag_mode": "live_session", "compaction": compaction},
+    )
+
+
+MTRAG_CLAPNQ_LIVE_SPEC = _mtrag_live_spec("mtrag_clapnq_live", compaction=True)
+MTRAG_CLAPNQ_LIVE_UNCOMPACTED_SPEC = _mtrag_live_spec(
+    "mtrag_clapnq_live_uncompacted", compaction=False
+)

evaluations/evaluations/datasets/open_rag_bench.py

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

evaluations/evaluations/datasets/t2_ragbench.py

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

evaluations/evaluations/evaluators/__init__.py

@@ -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,26 @@ from evaluations.evaluators.judge import (
 )
 from evaluations.evaluators.map import MAPEvaluator
 from evaluations.evaluators.number_match import NumberMatchEvaluator
+from evaluations.evaluators.refusal import (
+    REFUSAL_ELIGIBLE_LABELS,
+    REFUSAL_RUBRIC,
+    RefusalJudge,
+)
+from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator
+from evaluations.evaluators.transcript import TranscriptLLMJudge
 
 __all__ = [
     "ANSWER_EQUIVALENCE_RUBRIC",
+    "REFUSAL_ELIGIBLE_LABELS",
+    "REFUSAL_RUBRIC",
     "CitationMAPEvaluator",
+    "ConversationEvaluator",
     "LLMJudge",
     "LLMJudgeResponseSchema",
     "MAPEvaluator",
+    "NDCGEvaluator",
     "NumberMatchEvaluator",
+    "RecallEvaluator",
+    "RefusalJudge",
+    "TranscriptLLMJudge",
 ]

evaluations/evaluations/evaluators/citation.py

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

evaluations/evaluations/evaluators/conversation.py

@@ -0,0 +1,118 @@
+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_ELIGIBLE_LABELS, REFUSAL_RUBRIC
+
+
+@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_ELIGIBLE_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

evaluations/evaluations/evaluators/map.py

@@ -2,6 +2,8 @@ from dataclasses import dataclass
 
 from pydantic_evals.evaluators import Evaluator, EvaluatorContext
 
+from evaluations.evaluators.citation import average_precision
+
 
 @dataclass
 class MAPEvaluator(Evaluator):
@@ -28,22 +30,6 @@ class MAPEvaluator(Evaluator):
         if ctx.metadata is None:
             return 0.0
         relevant_uris = set(ctx.metadata.get("relevant_uris", []))
-        retrieved_uris = ctx.output
-
         if not relevant_uris:
             return 0.0
-
-        num_relevant = len(relevant_uris)
-        precisions = []
-        num_relevant_found = 0
-
-        for rank, uri in enumerate(retrieved_uris, start=1):
-            if uri in relevant_uris:
-                num_relevant_found += 1
-                precision_at_k = num_relevant_found / rank
-                precisions.append(precision_at_k)
-
-        if not precisions:
-            return 0.0
-
-        return sum(precisions) / num_relevant
+        return average_precision(list(ctx.output), relevant_uris)

evaluations/evaluations/evaluators/refusal.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+from typing import Final
+
+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."
+)
+
+REFUSAL_ELIGIBLE_LABELS: Final = ("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 REFUSAL_ELIGIBLE_LABELS:
+            return {}
+        return await super().evaluate(ctx)

evaluations/evaluations/evaluators/retrieval.py

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

evaluations/evaluations/evaluators/transcript.py

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

evaluations/tests/test_benchmark.py

@@ -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,476 @@ 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"
+
+    @pytest.mark.asyncio
+    async def test_records_citation_status_attribute(self, tmp_path: Path) -> None:
+        from dataclasses import dataclass
+
+        from pydantic_evals import Case
+        from pydantic_evals.evaluators import Evaluator, EvaluatorContext
+
+        from evaluations.capability_runner import CapabilityRunResult
+
+        @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="q1", 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(),
+        )
+
+        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_question",
+                new_callable=AsyncMock,
+                return_value=CapabilityRunResult(
+                    answer="answer", citation_status="ungrounded"
+                ),
+            ),
+        ):
+            await run_qa_benchmark(spec, AppConfig(), db_path=tmp_path / "test.lancedb")
+
+        assert recorded["citation_status"] == "ungrounded"
+
+
+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_macro_rate_excludes_fully_unjudged_conversations(self) -> None:
+        """A conversation whose every turn lost its judge reports
+        turn_pass_rate 0.0; treating that as a failed conversation would
+        contradict the exclusion policy. It must not enter the macro average."""
+        from evaluations.benchmark import _live_summary
+
+        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,
+                }
+            ),
+            self._case(
+                {
+                    "turn_pass_rate": 0.0,
+                    "turns_passed": 0,
+                    "turns_judged": 0,  # total judge outage for this conversation
+                    "turns_total": 8,
+                    "cited_eligible": 0,
+                    "true_refusals": 0,
+                    "false_refusals": 0,
+                    "unanswerable_turns": 0,
+                }
+            ),
+        ]
+
+        summary = _live_summary(cases, [])
+
+        assert summary is not None
+        assert summary["macro_pass_rate"] == pytest.approx(1.0)
+        assert summary["micro_pass_rate"] == pytest.approx(1.0)
+        assert summary["turns_judged"] == 2
+        assert summary["turns_total"] == 10
+
+    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,
+            compaction=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",
+                document_filter="uri = 'manual.pdf'",
+            )
+
+        assert run_conversation.await_args is not None
+        assert run_conversation.await_args.kwargs["questions"] == ["q1", "q2"]
+        assert run_conversation.await_args.kwargs["compaction"] is True
+        assert (
+            run_conversation.await_args.kwargs["document_filter"]
+            == "uri = 'manual.pdf'"
+        )
+
+    @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,
+                citation_status="grounded",
+            ),
+            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]
+        assert recorded["turn_citation_status"] == ["grounded", None]
+        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"
@@ -362,8 +832,9 @@ class TestRunQaBenchmarkCapabilityTarget:
         # (the capability manages its own client via lifespan).
         mock_haiku.assert_not_called()
         # capability model defaults to qa.model when not provided
-        capability_call = mock_get_model.call_args_list[-1]
-        assert capability_call[0][0] == AppConfig().qa.model
+        assert any(
+            call[0][0] == AppConfig().qa.model for call in mock_get_model.call_args_list
+        )
         assert mock_run_capability is capability_run
 
     @pytest.mark.asyncio
@@ -383,17 +854,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 +1004,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 +1082,7 @@ class TestRetrievalTarget:
             retrieval_mapper=lambda d: RetrievalSample(
                 question=d["q"], expected_uris=d["uris"]
             ),
-            retrieval_evaluator=MAPEvaluator(),
+            retrieval_evaluators=[MAPEvaluator()],
         )
 
     @pytest.mark.asyncio
@@ -620,7 +1191,7 @@ class TestDocumentFilterThreading:
             retrieval_mapper=lambda d: RetrievalSample(
                 question=d["q"], expected_uris=d["uris"]
             ),
-            retrieval_evaluator=MAPEvaluator(),
+            retrieval_evaluators=[MAPEvaluator()],
         )
 
         with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:

evaluations/tests/test_capability_runner.py

@@ -13,7 +13,11 @@ from pydantic_ai.messages import (
 )
 from pydantic_ai.models.test import TestModel
 
-from evaluations.capability_runner import _count_tool_traffic, run_capability_question
+from evaluations.capability_runner import (
+    CapabilityRunResult,
+    _count_tool_traffic,
+    run_capability_question,
+)
 from haiku.rag.capabilities.analysis import create_capability as create_analysis
 from haiku.rag.capabilities.rag import create_capability as create_rag
 from haiku.rag.config.models import AppConfig
@@ -118,6 +122,7 @@ async def test_runs_rag_capability_without_legacy_capability_layer(tmp_path):
     assert result.answer == "success (no tool calls)"
     assert result.cited_uris == []
     assert result.n_searches == 0
+    assert result.citation_status == "missing"
 
 
 async def test_runs_analysis_capability_without_legacy_capability_layer(tmp_path):
@@ -144,7 +149,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 +164,291 @@ 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 TestCitationStatusDerivation:
+    """`citation_status` distinguishes an answer that declared nothing
+    (`missing`) from one that declared ungrounded (`ungrounded`) — refusals
+    now cite an empty list."""
+
+    def _result(self, record) -> CapabilityRunResult:
+        from evaluations.capability_runner import ToolTraffic, _result_from_run
+        from haiku.rag.capabilities.rag import RAGState
+
+        state = RAGState(evidence=record)
+        return _result_from_run("answer", state, ToolTraffic(0, 0, 0, 1))
+
+    def test_grounded(self) -> None:
+        from haiku.rag.capabilities.ledger import (
+            CapabilityEvidenceRecord,
+            CitationDeclaration,
+            EvidenceRef,
+        )
+
+        record = CapabilityEvidenceRecord(
+            question=2,
+            latest_evidence_epoch=3,
+            declaration=CitationDeclaration(
+                question=2,
+                epoch=5,
+                refs=[EvidenceRef(capability="rag", chunk_id="c1")],
+            ),
+        )
+        assert self._result(record).citation_status == "grounded"
+
+    def test_ungrounded(self) -> None:
+        from haiku.rag.capabilities.ledger import (
+            CapabilityEvidenceRecord,
+            CitationDeclaration,
+        )
+
+        record = CapabilityEvidenceRecord(
+            question=2,
+            latest_evidence_epoch=3,
+            declaration=CitationDeclaration(question=2, epoch=5, refs=[]),
+        )
+        assert self._result(record).citation_status == "ungrounded"
+
+    def test_missing(self) -> None:
+        from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
+
+        record = CapabilityEvidenceRecord(question=2, latest_evidence_epoch=3)
+        assert self._result(record).citation_status == "missing"
+
+    def test_none_without_a_question(self) -> None:
+        """A record no run ever stamped (mocked runs) derives no status."""
+        from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
+
+        assert self._result(CapabilityEvidenceRecord()).citation_status is None
+
+
+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_applies_document_filter(tmp_path):
+    """The filter must reach the capability state so every search in the
+    conversation is restricted, same as the single-question runner."""
+    from evaluations.capability_runner import run_capability_conversation
+
+    deps_seen = []
+
+    async def _run(question, deps=None, message_history=None):
+        deps_seen.append(deps)
+        return SimpleNamespace(
+            output="a", all_messages=lambda: [], new_messages=lambda: []
+        )
+
+    with patch("evaluations.capability_runner.Agent.run", side_effect=_run):
+        await run_capability_conversation(
+            create_rag,
+            tmp_path / "rag.lancedb",
+            AppConfig(),
+            ["q1"],
+            TestModel(call_tools=[]),
+            document_filter="uri = 'manual.pdf'",
+        )
+
+    assert deps_seen[0].state["rag"]["document_filter"] == "uri = 'manual.pdf'"
+
+
+async def test_conversation_carries_one_state_dict_across_turns(tmp_path):
+    """Capabilities read and write state through the deps dict; carrying the
+    same dict across turns is what lets compaction see earlier questions'
+    records instead of refusing."""
+    from evaluations.capability_runner import run_capability_conversation
+
+    deps_seen: list[object] = []
+
+    async def _run(question, deps=None, message_history=None):
+        deps_seen.append(deps)
+        return SimpleNamespace(
+            output="a", all_messages=lambda: [], new_messages=lambda: []
+        )
+
+    with patch("evaluations.capability_runner.Agent.run", side_effect=_run):
+        await run_capability_conversation(
+            create_rag,
+            tmp_path / "rag.lancedb",
+            AppConfig(),
+            ["q1", "q2", "q3"],
+            TestModel(call_tools=[]),
+        )
+
+    assert deps_seen[0] is deps_seen[1] is deps_seen[2]
+
+
+@pytest.mark.parametrize(("compaction", "expected"), [(False, 0), (True, 1)])
+async def test_conversation_compaction_registration(tmp_path, compaction, expected):
+    from haiku.rag.capabilities.compaction import EvidenceCompactionCapability
+
+    from evaluations.capability_runner import run_capability_conversation
+
+    with patch("evaluations.capability_runner.Agent") as agent_cls:
+        agent_cls.return_value.run = AsyncMock(
+            return_value=SimpleNamespace(
+                output="a", all_messages=lambda: [], new_messages=lambda: []
+            )
+        )
+        await run_capability_conversation(
+            create_rag,
+            tmp_path / "rag.lancedb",
+            AppConfig(),
+            ["q1"],
+            TestModel(call_tools=[]),
+            compaction=compaction,
+        )
+        capabilities = agent_cls.call_args.kwargs["capabilities"]
+
+    compactors = [
+        c for c in capabilities if isinstance(c, EvidenceCompactionCapability)
+    ]
+    assert len(compactors) == expected
+    assert len(capabilities) == 1 + expected
+
+
+async def test_conversation_end_to_end_with_compaction(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=[]),
+        compaction=True,
+    )
+
+    assert [turn.answer for turn in result] == ["success (no tool calls)"] * 2
+
+
+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)"

evaluations/tests/test_citation_evaluators.py

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

evaluations/tests/test_config.py

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

evaluations/tests/test_conversation_evaluator.py

@@ -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",
+        }

evaluations/tests/test_evaluators.py

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

evaluations/tests/test_mtrag.py

@@ -0,0 +1,278 @@
+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)
+
+
+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 other_case.metadata["turns"][0]["relevant_uris"] == []
+
+    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",
+            "compaction": True,
+        }
+        assert MTRAG_CLAPNQ_SPEC.experiment_metadata == {"mtrag_mode": "gold_prefix"}
+
+    def test_live_compaction_arms(self) -> None:
+        assert MTRAG_CLAPNQ_LIVE_SPEC.compaction is True
+        uncompacted = DATASETS["mtrag_clapnq_live_uncompacted"]
+        assert uncompacted.compaction is False
+        assert uncompacted.live is True
+        assert uncompacted.db_filename == MTRAG_CLAPNQ_LIVE_SPEC.db_filename
+        assert uncompacted.qa_case_builder is MTRAG_CLAPNQ_LIVE_SPEC.qa_case_builder
+        assert uncompacted.experiment_metadata == {
+            "mtrag_mode": "live_session",
+            "compaction": False,
+        }

haiku_rag_slim/haiku/rag/client/documents.py

@@ -299,10 +299,16 @@ async def _store_documents_with_chunks(
     Embeds any chunks that lack embeddings, then writes the documents, chunks,
     and document_items tables once apiece. Restores all tables on any failure.
     """
-    embedded: list[list[Chunk]] = [
-        await ensure_chunks_embedded(client._config, chunks, client.embedder)
-        for _, chunks, _ in prepared
-    ]
+    flat = await ensure_chunks_embedded(
+        client._config,
+        [chunk for _, chunks, _ in prepared for chunk in chunks],
+        client.embedder,
+    )
+    embedded: list[list[Chunk]] = []
+    position = 0
+    for _, chunks, _ in prepared:
+        embedded.append(flat[position : position + len(chunks)])
+        position += len(chunks)
 
     def _extract_all_items():
         return [extract_items("", d) for _, _, d in prepared]

haiku_rag_slim/haiku/rag/client/processing.py

@@ -352,16 +352,10 @@ async def ensure_chunks_embedded(
 
     embedded = await embed_chunks(chunks_to_embed, embedder, config)
 
-    # Build result maintaining original order
-    embedded_map = {(c.content, c.order): c for c in embedded}
-    result = []
-    for ch in chunks:
-        if ch.embedding is not None:
-            result.append(ch)
-        else:
-            result.append(embedded_map[(ch.content, ch.order)])
-
-    return result
+    # embed_chunks preserves input order; fill positionally, since duplicate
+    # chunk texts across documents make a content-keyed lookup ambiguous.
+    filled = iter(embedded)
+    return [ch if ch.embedding is not None else next(filled) for ch in chunks]
 
 
 def get_extension_from_content_type_or_url(url: str, content_type: str) -> str:

tests/test_client.py

@@ -18,6 +18,7 @@ from haiku.rag.client.documents import (
     check_source_accessible,
 )
 from haiku.rag.config import Config
+from haiku.rag.embeddings import EmbedderWrapper
 from haiku.rag.store.compression import decompress_json
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
@@ -901,6 +902,94 @@ async def test_client_import_documents_empty(temp_db_path):
         assert after == before
 
 
+class _CountingEmbedder(EmbedderWrapper):
+    def __init__(self, vector_dim: int):
+        super().__init__(None, vector_dim)
+        self.batches: list[int] = []
+
+    async def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        self.batches.append(len(texts))
+        return [[0.1] * self.vector_dim for _ in texts]
+
+
+async def test_client_import_documents_batches_embeddings(temp_db_path):
+    """Chunks missing embeddings are embedded in one pass across the whole
+    batch, not one embedder call per document. Duplicate chunk texts across
+    documents keep their per-document embeddings."""
+    dim = Config.embeddings.model.vector_dim
+    embedder = _CountingEmbedder(dim)
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        client.store.embedder = embedder
+        imports = [
+            DocumentImport(
+                docling_document=_docling_doc(name, text),
+                chunks=[Chunk(content=text, order=0)],
+                uri=f"mem://{name}",
+                title=name,
+            )
+            for name, text in (
+                ("a", "Alpha document body"),
+                ("b", "Beta document body"),
+                ("c", "Alpha document body"),
+            )
+        ]
+
+        docs = await client.import_documents(imports)
+
+        assert embedder.batches == [3]
+        rows = await (
+            client.store.chunks_table.query()
+            .select(["document_id", "vector"])
+            .to_list()
+        )
+        assert {row["document_id"] for row in rows} == {doc.id for doc in docs}
+        assert all(len(row["vector"]) == dim for row in rows)
+
+
+async def test_client_import_documents_mixed_embeddings(temp_db_path):
+    """Pre-embedded chunks keep their vectors; only the unembedded ones go
+    through the embedder, in one batch."""
+    dim = Config.embeddings.model.vector_dim
+    embedder = _CountingEmbedder(dim)
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        client.store.embedder = embedder
+        pre_embedded = DocumentImport(
+            docling_document=_docling_doc("b", "Beta document body"),
+            chunks=[
+                Chunk(content="Beta document body", embedding=[0.5] * dim, order=0)
+            ],
+            uri="mem://b",
+            title="b",
+        )
+        unembedded = [
+            DocumentImport(
+                docling_document=_docling_doc(name, text),
+                chunks=[Chunk(content=text, order=0)],
+                uri=f"mem://{name}",
+                title=name,
+            )
+            for name, text in (("a", "Alpha document body"), ("c", "Gamma body"))
+        ]
+
+        docs = await client.import_documents(
+            [unembedded[0], pre_embedded, unembedded[1]]
+        )
+
+        assert embedder.batches == [2]
+        by_uri = {doc.uri: doc.id for doc in docs}
+        rows = await (
+            client.store.chunks_table.query()
+            .select(["document_id", "vector"])
+            .to_list()
+        )
+        vectors = {row["document_id"]: list(row["vector"]) for row in rows}
+        assert vectors[by_uri["mem://b"]] == pytest.approx([0.5] * dim)
+        assert vectors[by_uri["mem://a"]] == pytest.approx([0.1] * dim)
+        assert vectors[by_uri["mem://c"]] == pytest.approx([0.1] * dim)
+
+
 async def test_client_update_document_replaces_rows_with_bounded_versions(
     temp_db_path,
 ):