From cc04f92f28b601b2b5017265621b4f59a7ef199b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 6 Aug 2026 16:42:17 +0300 Subject: [PATCH 1/9] Batch embeddings across import_documents batches _store_documents_with_chunks embedded each document's chunks in its own embed_chunks call; chunks missing embeddings are now flattened across the whole batch, embedded in one pass honoring embeddings.batch_size, and assigned back positionally. --- CHANGELOG.md | 4 + haiku_rag_slim/haiku/rag/client/documents.py | 15 +++- tests/test_client.py | 89 ++++++++++++++++++++ 3 files changed, 106 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index bda57732..5a549e05 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,10 @@ - `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. +### Changed + +- `import_documents` embeds chunks across the whole batch in one pass instead of per document. + ### Removed - `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`. diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 8facaba2..1f5c4aae 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -299,10 +299,21 @@ 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) + missing = [ + chunk for _, chunks, _ in prepared + for chunk in chunks + if chunk.embedding is None ] + if missing: + from haiku.rag.embeddings import embed_chunks + + embedded_flat = await embed_chunks(missing, client.embedder, client._config) + # Assign positionally: duplicate chunk texts across documents make a + # content-keyed lookup ambiguous. + for chunk, with_embedding in zip(missing, embedded_flat): + chunk.embedding = with_embedding.embedding + embedded: list[list[Chunk]] = [chunks for _, chunks, _ in prepared] def _extract_all_items(): return [extract_items("", d) for _, _, d in prepared] diff --git a/tests/test_client.py b/tests/test_client.py index db2da16f..939d1f07 100644 --- a/tests/test_client.py +++ b/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, ): From 73d9d93db941f4d7072d390c6309e82b5ed9bad4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 6 Aug 2026 16:42:54 +0300 Subject: [PATCH 2/9] Add MTRAG ClapNQ multi-turn evaluation IBM's MTRAG benchmark (ClapNQ domain, pinned repo SHA): retrieval with Recall@k/nDCG@k against binary qrels, gold-prefix QA replaying reference conversation prefixes as message history, and live-session replay carrying the model's own answers and tool history across turns. Corpus population gains a bounded, resumable batched ingest path. ConversationInput case type with transcript rendering for the judge, eligibility-aware citation scoring, refusal precision/recall via a label-aware RefusalJudge, per-turn verdicts with judged-turn coverage, and per-turn tool-traffic attributes counted from each turn's new messages so the arrays survive prior-turn compaction. --- CHANGELOG.md | 1 + docs/benchmarks.md | 29 +- evaluations/README.md | 1 + evaluations/configs/mtrag_clapnq.yaml | 48 ++ evaluations/evaluations/benchmark.py | 434 +++++++++++++-- evaluations/evaluations/capability_runner.py | 140 ++++- evaluations/evaluations/config.py | 43 +- evaluations/evaluations/datasets/__init__.py | 8 + evaluations/evaluations/datasets/hotpotqa.py | 5 +- evaluations/evaluations/datasets/mtrag.py | 298 +++++++++++ .../evaluations/datasets/open_rag_bench.py | 5 +- .../evaluations/datasets/t2_ragbench.py | 9 +- .../evaluations/evaluators/__init__.py | 10 + .../evaluations/evaluators/citation.py | 31 +- .../evaluations/evaluators/conversation.py | 120 +++++ evaluations/evaluations/evaluators/refusal.py | 28 + .../evaluations/evaluators/retrieval.py | 54 ++ .../evaluations/evaluators/transcript.py | 21 + evaluations/tests/test_benchmark.py | 496 +++++++++++++++++- evaluations/tests/test_capability_runner.py | 140 ++++- evaluations/tests/test_citation_evaluators.py | 18 +- evaluations/tests/test_config.py | 58 +- .../tests/test_conversation_evaluator.py | 257 +++++++++ evaluations/tests/test_evaluators.py | 190 ++++++- evaluations/tests/test_mtrag.py | 268 ++++++++++ 25 files changed, 2596 insertions(+), 116 deletions(-) create mode 100644 evaluations/configs/mtrag_clapnq.yaml create mode 100644 evaluations/evaluations/datasets/mtrag.py create mode 100644 evaluations/evaluations/evaluators/conversation.py create mode 100644 evaluations/evaluations/evaluators/refusal.py create mode 100644 evaluations/evaluations/evaluators/retrieval.py create mode 100644 evaluations/evaluations/evaluators/transcript.py create mode 100644 evaluations/tests/test_conversation_evaluator.py create mode 100644 evaluations/tests/test_mtrag.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 5a549e05..d370a087 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. +- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, and per-turn tool-traffic attributes. ### Changed diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 69240049..419e6d7f 100644 --- a/docs/benchmarks.md +++ b/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` and `mtrag_clapnq_live` keys | ~2.8 GB | After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches): @@ -225,3 +226,29 @@ The reranker's contribution is larger here than on the single-doc datasets: hybr | `vllm:Gemma-4-26B-A4B-NVFP4` | none | 0.83 | 0.75 | *Measured on haiku.rag v0.66.0 with `qwen3-embedding:4b` (vLLM, dim 2560), judged by `vllm:Qwen3.6-35B-A3B-NVFP4`, 7,405 cases. The reranker lifts QA accuracy +2.7pts and `cited_map` +4.6pts. Without a reranker, `cited_map` (0.75) still exceeds the no-reranker retrieval MAP (0.70): the skill reformulates queries across search calls, partially recovering second-hop documents that a single query misses.* + +### MTRAG (ClapNQ) + +[MTRAG](https://github.com/IBM/mt-rag-benchmark) is IBM's multi-turn RAG benchmark (TACL 2025, SemEval-2026 Task 8): human-authored conversations with per-turn answerability labels and binary relevance judgments. We evaluate the ClapNQ (Wikipedia) domain: 183,408 passages, 29 conversations, 224 turns, 208 retrieval queries. + +Three dataset keys share one database. `mtrag_clapnq` retrieves with the raw last user turn and runs QA by replaying each task's reference conversation prefix as message history. `mtrag_clapnq_rewrite` retrieves with the human standalone rewrites. `mtrag_clapnq_live` replays whole conversations through a single capability session, carrying the model's own answers and tool history across turns. + +##### Retrieval (Recall@k / nDCG@k) + +Directly comparable with [IBM's published results](https://github.com/IBM/mt-rag-benchmark/tree/main/mtrag-human/retrieval_tasks). Elser is IBM's strongest reported retriever. + +| Retriever | Queries | R@5 | R@10 | nDCG@5 | nDCG@10 | +|-----------|---------|----:|-----:|-------:|--------:| +| Elser (IBM) | lastturn | 0.49 | 0.58 | 0.45 | 0.49 | +| `haiku.rag` | lastturn | 0.501 | 0.600 | 0.455 | 0.497 | +| Elser (IBM) | rewrite | 0.52 | 0.64 | 0.48 | 0.54 | +| `haiku.rag` | rewrite | 0.548 | 0.668 | 0.503 | 0.556 | + +##### QA accuracy + citation retrieval + +| Mode | Capability model | Turns | QA accuracy | Mean `cited_map` | +|------|------------------|------:|-------------|------------------| +| Gold-prefix (`mtrag_clapnq`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 223 | 0.68 | 0.35 | +| Live (`mtrag_clapnq_live`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 195/224 scored | 0.72 micro / 0.73 macro | 0.35 | + +*Measured on haiku.rag v0.67.1 with `qwen3-embedding:4b` (vLLM, dim 2560) and `Qwen3-Reranker-4B`, stock capability instructions, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` at temperature 0. The judge sampling has since been re-pinned repo-wide (0.6 with thinking), so future runs re-baseline. QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Live mode additionally reports refusal precision/recall against the per-turn answerability labels and per-turn pass rates; pass rate declines with conversation depth (93% at turn 1 to 38% at turn 9). The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.* diff --git a/evaluations/README.md b/evaluations/README.md index 954072f7..ff1c6f7d 100644 --- a/evaluations/README.md +++ b/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. diff --git a/evaluations/configs/mtrag_clapnq.yaml b/evaluations/configs/mtrag_clapnq.yaml new file mode 100644 index 00000000..654b1eb6 --- /dev/null +++ b/evaluations/configs/mtrag_clapnq.yaml @@ -0,0 +1,48 @@ +# Reference config for the `mtrag_clapnq` pre-built evaluation database. +# IBM MTRAG, ClapNQ (Wikipedia) domain: multi-turn retrieval and QA over +# 183,408 passages. Also serves mtrag_clapnq_rewrite and mtrag_clapnq_live. +# Run: evaluations run mtrag_clapnq --config configs/mtrag_clapnq.yaml +# base_url uses the `vllm` host serving each model over an OpenAI-compatible API. +# The corpus is text-only: no multimodal embedder, no vision paths. This eval +# cannot exercise image or vision turn-boundary behavior. + +environment: development + +storage: + auto_vacuum: false + +embeddings: + model: + provider: openai + name: qwen3-embedding-4b + vector_dim: 2560 + base_url: http://vllm:11431/v1 + +reranking: + model: + provider: vllm + name: Qwen/Qwen3-Reranker-4B + base_url: http://vllm:11455 + +qa: + model: + provider: openai + name: gemma4-26b + base_url: http://vllm:11432/v1 + # vLLM enforces input + max_tokens <= max_model_len, so a large output + # budget silently shrinks the input budget. MTRAG answers are sentences. + max_tokens: 8192 + +evaluations: + judge: + provider: openai + name: RedHatAI/Qwen3.6-35B-A3B-NVFP4 + base_url: http://vllm:11430/v1 + temperature: 0.6 + max_tokens: 16384 + extra_body: + top_p: 0.95 + top_k: 20 + min_p: 0 + chat_template_kwargs: + enable_thinking: true diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 36aaa70b..f254b67b 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -8,20 +8,28 @@ import typer from dotenv import find_dotenv, load_dotenv from huggingface_hub import HfApi, snapshot_download from pydantic_evals import Case, Dataset as EvalDataset, set_eval_attribute -from pydantic_evals.evaluators import Evaluator, LLMJudge +from pydantic_evals.evaluators import Evaluator from pydantic_evals.reporting import ReportCaseFailure from rich.console import Console from rich.progress import Progress -from evaluations.config import DatasetSpec +from evaluations.config import ConversationInput, DatasetSpec from evaluations.datasets import DATASETS from evaluations.evaluators import ( ANSWER_EQUIVALENCE_RUBRIC, - CitationMAPEvaluator, - MAPEvaluator, + REFUSAL_RUBRIC, + ConversationEvaluator, + RefusalJudge, + TranscriptLLMJudge, +) +from evaluations.capability_runner import ( + CapabilityFactory, + prefix_to_messages, + run_capability_conversation, + run_capability_question, ) -from evaluations.capability_runner import CapabilityFactory, run_capability_question from haiku.rag.client import HaikuRAG +from haiku.rag.client.documents import DocumentImport from haiku.rag.config import AppConfig, find_config_file, load_yaml_config from haiku.rag.config.models import ModelConfig from haiku.rag.logging import configure_cli_logging @@ -118,6 +126,59 @@ def build_experiment_metadata( return metadata +async def _ingest_batched( + rag: HaikuRAG, + spec: DatasetSpec, + corpus, + batch_size: int, + on_document: Callable[[], None] = lambda: None, +) -> None: + """Ingest inline-content documents via `import_documents` batches. + + Each batch writes the documents/chunks/document_items tables once and + embeds every chunk in one batched pass. A URI is skipped on resume only + when its document has chunks; a chunkless document (crash between the + document and chunk writes) is deleted and re-imported. + """ + uri_rows = await ( + rag.store.document_meta_table.query().select(["id", "uri"]).to_list() + ) + chunk_rows = await rag.store.chunks_table.query().select(["document_id"]).to_list() + chunked_ids = {row["document_id"] for row in chunk_rows} + complete = {row["uri"] for row in uri_rows if row["id"] in chunked_ids} + chunkless = { + row["uri"]: row["id"] for row in uri_rows if row["id"] not in chunked_ids + } + + batch: list[DocumentImport] = [] + for doc in corpus: + payload = spec.document_mapper(cast(Mapping[str, Any], doc)) + if payload is None or payload.uri in complete: + on_document() + continue + if payload.uri in chunkless: + await rag.delete_document(chunkless[payload.uri]) + assert payload.content is not None, "batched ingest requires inline content" + docling_document = await rag.convert(payload.content, format=payload.format) + chunks = await rag.chunk(docling_document) + batch.append( + DocumentImport( + docling_document=docling_document, + chunks=chunks, + uri=payload.uri, + title=payload.title, + metadata=payload.metadata or {}, + ) + ) + if len(batch) >= batch_size: + await rag.import_documents(batch) + batch = [] + on_document() + + if batch: + await rag.import_documents(batch) + + async def populate_db( spec: DatasetSpec, config: AppConfig, @@ -136,6 +197,17 @@ async def populate_db( with Progress() as progress: task = progress.add_task("[green]Populating database...", total=len(corpus)) async with HaikuRAG(db, config=config, create=True) as rag: + if spec.ingest_batch_size is not None: + await _ingest_batched( + rag, + spec, + corpus, + batch_size=spec.ingest_batch_size, + on_document=lambda: progress.advance(task), + ) + await rag.store.vacuum(retention_seconds=0) + return + docs_since_vacuum = 0 for doc in corpus: doc_mapping = cast(Mapping[str, Any], doc) @@ -233,16 +305,13 @@ async def run_retrieval_benchmark( console.print("No retrieval cases to evaluate.") return None - if spec.retrieval_evaluator is None: - raise ValueError(f"No retrieval evaluator configured for dataset: {spec.key}") - - evaluator = spec.retrieval_evaluator - metric_name = evaluator.__class__.__name__.replace("Evaluator", "").upper() + if not spec.retrieval_evaluators: + raise ValueError(f"No retrieval evaluators configured for dataset: {spec.key}") dataset = EvalDataset( name=f"{spec.key}-retrieval", cases=cases, - evaluators=[evaluator], + evaluators=list(spec.retrieval_evaluators), ) db = spec.db_path(db_path) @@ -250,7 +319,10 @@ async def run_retrieval_benchmark( async def retrieval_target(question: str) -> list[str]: chunks = await rag.search( - query=question, limit=5, include_images=False, filter=document_filter + query=question, + limit=spec.retrieval_limit, + include_images=False, + filter=document_filter, ) seen = set() @@ -280,25 +352,22 @@ async def run_retrieval_benchmark( metadata=experiment_metadata, ) - total_score = 0.0 - total_cases = 0 + per_metric: dict[str, list[float]] = {} for case in report.cases: - if case.scores: - for score_result in case.scores.values(): - total_score += score_result.value - total_cases += 1 - - mean_score = total_score / total_cases if total_cases > 0 else 0.0 + for key, score_result in case.scores.items(): + per_metric.setdefault(key, []).append(score_result.value) console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan") console.print(f"Dataset: {spec.key}") console.print(f"Total queries: {len(cases)}") - console.print(f"{metric_name}: {mean_score:.4f}") + results: dict[str, float] = {"queries": len(cases)} + for key, values in per_metric.items(): + mean_score = sum(values) / len(values) + metric_name = key.replace("Evaluator", "").upper() + console.print(f"{metric_name}: {mean_score:.4f}") + results[metric_name.lower()] = mean_score - return { - metric_name.lower(): mean_score, - "queries": len(cases), - } + return results def _capability_factory_for_target(target: Target) -> CapabilityFactory: @@ -313,13 +382,6 @@ def _capability_factory_for_target(target: Target) -> CapabilityFactory: raise ValueError(f"target {target!r} is not a capability target") -def _citation_evaluator_for(retrieval_evaluator: Evaluator | None) -> Evaluator | None: - """Return the citation-scoring twin of the dataset's retrieval evaluator.""" - if isinstance(retrieval_evaluator, MAPEvaluator): - return CitationMAPEvaluator() - return None - - def _attach_relevant_uris( cases: list[Case[str, str, dict[str, Any]]], spec: DatasetSpec, @@ -342,6 +404,8 @@ def _attach_relevant_uris( continue expected_by_question[sample.question] = sample.expected_uris for case in cases: + if not isinstance(case.inputs, str): + continue uris = expected_by_question.get(case.inputs) if uris is None: continue @@ -350,6 +414,102 @@ def _attach_relevant_uris( case.metadata = metadata +def _resolve_capability_config( + target: Target, config: AppConfig, capability_model: ModelConfig | None +) -> ModelConfig: + if target == "analysis-capability": + # Mirror the capability-code resolver: explicit analysis.model wins, + # else fall back to qa.model. + return capability_model or config.analysis.model or config.qa.model + return capability_model or config.qa.model + + +def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | None: + """Aggregate ConversationEvaluator scores across conversations. + + Micro rates weight every turn equally (sums across conversations); macro + rates average per-conversation means, so short conversations don't get + overweighted by micro nor long ones by macro. Failed conversations are + operational exclusions: they count toward the attempted coverage figures + but never toward the rates. + """ + + def _score(case, key: str): + result = case.scores.get(key) + return result.value if result is not None else None + + scored = [case for case in report_cases if _score(case, "turns_total") is not None] + if not scored: + return None + + failed_turns = sum( + len(failure.inputs) if isinstance(failure.inputs, list) else 0 + for failure in report_failures + ) + turns_total = sum(_score(case, "turns_total") for case in scored) + turns_judged = sum(_score(case, "turns_judged") or 0 for case in scored) + turns_passed = sum(_score(case, "turns_passed") for case in scored) + summary: dict[str, float | int] = { + "conversations": len(scored), + "conversations_attempted": len(report_cases) + len(report_failures), + "turns_total": turns_total, + "turns_judged": turns_judged, + "turns_attempted": turns_total + failed_turns, + "micro_pass_rate": turns_passed / turns_judged if turns_judged else 0.0, + "macro_pass_rate": sum(_score(case, "turn_pass_rate") for case in scored) + / len(scored), + } + + cited = [case for case in scored if _score(case, "cited_map") is not None] + eligible = sum(_score(case, "cited_eligible") for case in scored) + if cited and eligible: + summary["cited_eligible"] = eligible + summary["cited_map_micro"] = ( + sum( + _score(case, "cited_map") * _score(case, "cited_eligible") + for case in cited + ) + / eligible + ) + summary["cited_map_macro"] = sum( + _score(case, "cited_map") for case in cited + ) / len(cited) + + true_refusals = sum(_score(case, "true_refusals") or 0 for case in scored) + false_refusals = sum(_score(case, "false_refusals") or 0 for case in scored) + unanswerable = sum(_score(case, "unanswerable_turns") or 0 for case in scored) + refusals = true_refusals + false_refusals + summary["unanswerable_turns"] = unanswerable + summary["refusals"] = refusals + summary["refusal_precision"] = true_refusals / refusals if refusals else 0.0 + summary["refusal_recall"] = true_refusals / unanswerable if unanswerable else 0.0 + return summary + + +def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None: + """Refusal precision/recall against answerability labels. + + Uses cases the refusal judge scored (ANSWERABLE/UNANSWERABLE turns). + Returns (precision, recall, unanswerable_count, refusal_count), or None + when no case was judged. + """ + outcomes: list[tuple[str, bool]] = [] + for case in report_cases: + refused = case.assertions.get("refused") + label = (case.metadata or {}).get("answerability") + if refused is None or label not in ("ANSWERABLE", "UNANSWERABLE"): + continue + outcomes.append((label, bool(refused.value))) + if not outcomes: + return None + refusals = [(label, r) for label, r in outcomes if r] + true_refusals = sum(1 for label, _ in refusals if label == "UNANSWERABLE") + unanswerable = sum(1 for label, _ in outcomes if label == "UNANSWERABLE") + precision = true_refusals / len(refusals) if refusals else 0.0 + recall = true_refusals / unanswerable if unanswerable else 0.0 + return precision, recall, unanswerable, len(refusals) + + def _filter_qa_corpus(corpus, case_ids: set[str] | None): """Keep only rows whose ``id`` is in ``case_ids`` (failure-subset reruns). @@ -383,16 +543,11 @@ async def run_qa_benchmark( ] judge_config = judge_model or DEFAULT_JUDGE_MODEL - if target == "analysis-capability": - # Mirror the capability-code resolver: explicit analysis.model wins, - # else fall back to qa.model. - capability_config = capability_model or config.analysis.model or config.qa.model - else: - capability_config = capability_model or config.qa.model + capability_config = _resolve_capability_config(target, config, capability_model) db = spec.db_path(db_path) _attach_relevant_uris(cases, spec, limit) - citation_evaluator = _citation_evaluator_for(spec.retrieval_evaluator) + citation_evaluator = spec.citation_evaluator qa_evaluator = spec.qa_evaluator evaluators: list[Evaluator] @@ -400,7 +555,7 @@ async def run_qa_benchmark( evaluators = [qa_evaluator] else: evaluators = [ - LLMJudge( + TranscriptLLMJudge( rubric=ANSWER_EQUIVALENCE_RUBRIC, include_input=True, include_expected_output=True, @@ -413,8 +568,16 @@ async def run_qa_benchmark( ] if citation_evaluator is not None: evaluators.append(citation_evaluator) + if spec.evaluate_refusal: + evaluators.append( + RefusalJudge( + rubric=REFUSAL_RUBRIC, + model=get_model(judge_config, config), + assertion={"evaluation_name": "refused", "include_reason": False}, + ) + ) - evaluation_dataset = EvalDataset[str, str, dict[str, str]]( + evaluation_dataset = EvalDataset[Any, str, dict[str, Any]]( name=spec.key, cases=cases, evaluators=evaluators ) @@ -428,8 +591,9 @@ async def run_qa_benchmark( capability_config=capability_config, document_filter=document_filter, ) + experiment_metadata.update(spec.experiment_metadata or {}) - async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]): + async def _evaluate(answer_fn: Callable[[Any], Awaitable[str]]): return await evaluation_dataset.evaluate( answer_fn, name=eval_name, @@ -441,7 +605,13 @@ async def run_qa_benchmark( capability_factory = _capability_factory_for_target(target) resolved_capability_model = get_model(capability_config, config) - async def answer_question(question: str) -> str: + async def answer_question(inputs: str | ConversationInput) -> str: + if isinstance(inputs, ConversationInput): + question = inputs.question + message_history = prefix_to_messages(inputs.prefix) + else: + question = inputs + message_history = None result = await run_capability_question( capability_factory=capability_factory, db_path=db, @@ -449,6 +619,7 @@ async def run_qa_benchmark( question=question, capability_model=resolved_capability_model, document_filter=document_filter, + message_history=message_history, ) set_eval_attribute("cited_uris", result.cited_uris) set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids) @@ -488,6 +659,11 @@ async def run_qa_benchmark( console.print(f"Total questions: {total_processed}") console.print(f"Correct answers: {passing_cases}") console.print(f"QA Accuracy: {accuracy:.4f} ({accuracy * 100:.2f}%)") + if report.cases: + mean_task_time = sum(case.task_duration for case in report.cases) / len( + report.cases + ) + console.print(f"Avg task time per case: {mean_task_time:.2f}s") if citation_evaluator is not None: score_key = citation_evaluator.get_default_evaluation_name() @@ -508,11 +684,27 @@ async def run_qa_benchmark( f"\n=== Citation Retrieval ({score_key}) ===", style="bold cyan" ) console.print(f"Mean {score_key}: {mean_score:.4f}") + console.print( + f"Eligible cases (gold passages known): {len(scores)}/{len(report.cases)}" + ) console.print( f"Cite rate (≥1 citation): {cited_count / len(report.cases):.2%}" ) console.print(f"Mean citations per case: {mean_citations:.2f}") + if spec.evaluate_refusal: + metrics = _refusal_metrics(report.cases) + if metrics is not None: + precision, recall, unanswerable, refusals = metrics + console.print( + "\n=== Refusal vs answerability labels ===", style="bold cyan" + ) + console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}") + console.print( + f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} " + "(PARTIAL excluded)" + ) + if failures: console.print("[red]\nSummary of failures:[/red]") for failure in failures: @@ -524,6 +716,144 @@ async def run_qa_benchmark( return failures[0] if failures else None +async def run_live_qa_benchmark( + spec: DatasetSpec, + config: AppConfig, + limit: int | None = None, + name: str | None = None, + db_path: Path | None = None, + judge_model: ModelConfig | None = None, + target: Target = "rag-capability", + capability_model: ModelConfig | None = None, + case_ids: set[str] | None = None, +) -> None: + """Replay conversations turn by turn through one capability session. + + One case per conversation; ``limit`` counts conversations. Answers carry + forward as real message history, so prior-turn compaction is exercised. + """ + corpus = spec.qa_loader() + corpus = _filter_qa_corpus(corpus, case_ids) + if limit is not None: + corpus = corpus.select(range(min(limit, len(corpus)))) + + cases = [ + spec.qa_case_builder(index, cast(Mapping[str, Any], doc)) + for index, doc in enumerate(corpus, start=1) + ] + + judge_config = judge_model or DEFAULT_JUDGE_MODEL + capability_config = _resolve_capability_config(target, config, capability_model) + db = spec.db_path(db_path) + + evaluation_dataset = EvalDataset[Any, Any, dict[str, Any]]( + name=spec.key, + cases=cases, + evaluators=[ + ConversationEvaluator( + rubric=ANSWER_EQUIVALENCE_RUBRIC, + model=get_model(judge_config, config), + ) + ], + ) + + eval_name = name if name is not None else f"{spec.key}_qa_evaluation" + experiment_metadata = build_experiment_metadata( + dataset_key=spec.key, + test_cases=len(cases), + config=config, + judge_config=judge_config, + target=target, + capability_config=capability_config, + ) + experiment_metadata.update(spec.experiment_metadata or {}) + + capability_factory = _capability_factory_for_target(target) + resolved_capability_model = get_model(capability_config, config) + + async def answer_conversation(questions: list[str]) -> list[str]: + results = await run_capability_conversation( + capability_factory=capability_factory, + db_path=db, + config=config, + questions=list(questions), + capability_model=resolved_capability_model, + ) + set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results]) + set_eval_attribute("turn_n_search_calls", [r.n_search_calls for r in results]) + set_eval_attribute( + "turn_n_rejected_searches", [r.n_rejected_searches for r in results] + ) + set_eval_attribute("turn_n_failed_tools", [r.n_failed_tools for r in results]) + set_eval_attribute("turn_n_requests", [r.n_requests for r in results]) + return [r.answer for r in results] + + report = await evaluation_dataset.evaluate( + answer_conversation, + name=eval_name, + max_concurrency=1, + progress=True, + metadata=experiment_metadata, + ) + + summary = _live_summary(report.cases, report.failures) + console.print("\n=== Live Conversation Results ===", style="bold cyan") + if summary is None: + attempted = len(report.cases) + len(report.failures) + console.print(f"No conversations were scored ({attempted} attempted).") + else: + console.print( + f"Conversations scored: {summary['conversations']}" + f"/{summary['conversations_attempted']} | turns scored: " + f"{summary['turns_total']}/{summary['turns_attempted']}" + ) + if summary["turns_judged"] < summary["turns_total"]: + console.print( + f"Turns judged: {summary['turns_judged']}/{summary['turns_total']} " + "(per-turn judge errors excluded from rates)" + ) + if report.failures: + console.print( + "Failed conversations are operational exclusions — " + "not counted as wrong answers." + ) + console.print( + f"Answer pass rate — micro (per turn): {summary['micro_pass_rate']:.4f} | " + f"macro (per conversation): {summary['macro_pass_rate']:.4f}" + ) + if "cited_map_micro" in summary: + console.print( + f"cited_map — micro: {summary['cited_map_micro']:.4f} | " + f"macro: {summary['cited_map_macro']:.4f} " + f"(eligible turns: {summary['cited_eligible']})" + ) + console.print( + f"Refusal precision: {summary['refusal_precision']:.2%} | " + f"recall: {summary['refusal_recall']:.2%} " + f"(UNANSWERABLE turns: {summary['unanswerable_turns']}, " + f"refusals: {summary['refusals']})" + ) + if report.cases: + mean_task_time = sum(case.task_duration for case in report.cases) / len( + report.cases + ) + turns = sum(len(case.output or []) for case in report.cases) + per_turn = ( + sum(case.task_duration for case in report.cases) / turns if turns else 0.0 + ) + console.print( + f"Avg task time: {mean_task_time:.2f}s per conversation | " + f"{per_turn:.2f}s per turn" + ) + + if report.failures: + console.print("[red]\nSummary of failures:[/red]") + for failure in report.failures: + console.print(f"Case: {failure.name}") + console.print(f"Error: {failure.error_message}") + console.print("") + + async def evaluate_dataset( spec: DatasetSpec, config: AppConfig, @@ -566,7 +896,8 @@ async def evaluate_dataset( console.print( f"\nRunning QA benchmarks (target={target})...", style="bold yellow" ) - await run_qa_benchmark( + qa_benchmark = run_live_qa_benchmark if spec.live else run_qa_benchmark + await qa_benchmark( spec, config, limit=limit, @@ -621,9 +952,20 @@ def _resolve_dataset(dataset: str) -> DatasetSpec: def _resolve_datasets(dataset: str) -> list[DatasetSpec]: - """Resolve 'all' or a single dataset key to a list of DatasetSpecs.""" + """Resolve 'all' or a single dataset key to a list of DatasetSpecs. + + 'all' yields one spec per database: query variants sharing a db_filename + would otherwise be downloaded/uploaded twice. + """ if dataset.lower() == "all": - return list(DATASETS.values()) + seen: set[str] = set() + specs: list[DatasetSpec] = [] + for spec in DATASETS.values(): + if spec.db_filename in seen: + continue + seen.add(spec.db_filename) + specs.append(spec) + return specs return [_resolve_dataset(dataset)] diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index b87a6c4c..04ff75d1 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/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,13 +6,17 @@ from typing import Any, NamedTuple, Protocol, cast from pydantic_ai import Agent from pydantic_ai.messages import ( ModelMessage, + ModelRequest, ModelResponse, RetryPromptPart, + TextPart, ToolCallPart, ToolReturnPart, + UserPromptPart, ) from pydantic_ai.models import Model +from evaluations.config import Turn from haiku.rag.capabilities import RAGCapabilityBase from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult @@ -21,6 +25,17 @@ from haiku.rag.store.models.citation import Citation CapabilityFactory = Callable[..., RAGCapabilityBase[Any]] +def prefix_to_messages(turns: Iterable[Turn]) -> list[ModelMessage]: + """Render a conversation prefix as pydantic-ai message history.""" + messages: list[ModelMessage] = [] + for turn in turns: + if turn.speaker == "user": + messages.append(ModelRequest(parts=[UserPromptPart(content=turn.text)])) + else: + messages.append(ModelResponse(parts=[TextPart(content=turn.text)])) + return messages + + class _RagLikeState(Protocol): document_filter: str | None citation_index: dict[str, Citation] @@ -107,25 +122,14 @@ class _EvalDeps: state: dict[str, Any] = field(default_factory=dict) -async def run_capability_question( +def _prepare_agent( capability_factory: CapabilityFactory, db_path: Path, config: AppConfig, - question: str, capability_model: str | Model, - document_filter: str | None = None, - request_limit: int | None = None, -) -> CapabilityRunResult: - """Run a single question through a capability and return answer + retrieval data. - - Builds a native capability via ``capability_factory(db_path=..., config=...)``. - After the run, citations and searched documents - are extracted from the state for downstream eval scoring. - - The capability must produce a state with RAG-capability-shaped fields (citation - index, searches, optional document filter) — i.e. ``RAGState`` or - ``AnalysisState`` from ``haiku.rag.capabilities``. - """ + document_filter: str | None, + request_limit: int | None, +) -> tuple[RAGCapabilityBase[Any], _EvalDeps, Agent[_EvalDeps, str]]: capability = capability_factory( db_path=db_path, config=config, @@ -144,10 +148,98 @@ async def run_capability_question( deps_type=_EvalDeps, capabilities=[capability], ) - agent_result = await agent.run(question, deps=deps) - state = capability.state_type.model_validate(deps.state[capability.state_namespace]) - typed = cast(_RagLikeState, state) + return capability, deps, agent + +def _state_after_run( + capability: RAGCapabilityBase[Any], deps: _EvalDeps +) -> _RagLikeState: + state = capability.state_type.model_validate(deps.state[capability.state_namespace]) + return cast(_RagLikeState, state) + + +async def run_capability_question( + capability_factory: CapabilityFactory, + db_path: Path, + config: AppConfig, + question: str, + capability_model: str | Model, + document_filter: str | None = None, + request_limit: int | None = None, + message_history: list[ModelMessage] | None = None, +) -> CapabilityRunResult: + """Run a single question through a capability and return answer + retrieval data. + + Builds a native capability via ``capability_factory(db_path=..., config=...)``. + After the run, citations and searched documents + are extracted from the state for downstream eval scoring. + + The capability must produce a state with RAG-capability-shaped fields (citation + index, searches, optional document filter) — i.e. ``RAGState`` or + ``AnalysisState`` from ``haiku.rag.capabilities``. + """ + capability, deps, agent = _prepare_agent( + capability_factory, + db_path, + config, + capability_model, + document_filter, + request_limit, + ) + agent_result = await agent.run(question, deps=deps, message_history=message_history) + traffic = _count_tool_traffic( + agent_result.new_messages(), capability.state_namespace, capability.tool_names + ) + return _result_from_run( + agent_result.output, _state_after_run(capability, deps), traffic + ) + + +async def run_capability_conversation( + capability_factory: CapabilityFactory, + db_path: Path, + config: AppConfig, + questions: list[str], + capability_model: str | Model, + document_filter: str | None = None, + request_limit: int | None = None, +) -> list[CapabilityRunResult]: + """Run a conversation's user turns sequentially through one capability. + + Each turn runs with the previous turn's full ``all_messages()`` as history + (tool calls and returns included), so prior-turn compaction operates on + real evidence. Per-invocation state (citations, searches) is cleared by the + capability on every run, so each returned result reflects only its turn. + """ + capability, deps, agent = _prepare_agent( + capability_factory, + db_path, + config, + capability_model, + document_filter, + request_limit, + ) + history: list[ModelMessage] | None = None + results: list[CapabilityRunResult] = [] + for question in questions: + agent_result = await agent.run(question, deps=deps, message_history=history) + history = agent_result.all_messages() + traffic = _count_tool_traffic( + agent_result.new_messages(), + capability.state_namespace, + capability.tool_names, + ) + results.append( + _result_from_run( + agent_result.output, _state_after_run(capability, deps), traffic + ) + ) + return results + + +def _result_from_run( + answer: str, typed: _RagLikeState, traffic: ToolTraffic +) -> CapabilityRunResult: cited_chunk_ids: list[str] = list(typed.citations) seen_cited: set[str] = set() cited_uris: list[str] = [] @@ -168,17 +260,11 @@ async def run_capability_question( seen_searched.add(uri) searched_uris.append(uri) - executions = getattr(state, "executions", None) + executions = getattr(typed, "executions", None) n_executions = len(executions) if executions is not None else 0 - traffic = _count_tool_traffic( - agent_result.all_messages(), - capability.state_namespace, - capability.tool_names, - ) - return CapabilityRunResult( - answer=agent_result.output, + answer=answer, cited_uris=cited_uris, cited_chunk_ids=cited_chunk_ids, searched_uris=searched_uris, diff --git a/evaluations/evaluations/config.py b/evaluations/evaluations/config.py index c5ecfd17..d8b0ebe0 100644 --- a/evaluations/evaluations/config.py +++ b/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,8 @@ DocumentLoader = Callable[[], Dataset] DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None] RetrievalLoader = Callable[[], Dataset] RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None] -CaseBuilder = Callable[[int, Mapping[str, Any]], Case[str, str, dict[str, str]]] +QAInput = str | ConversationInput +CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]] @dataclass @@ -43,9 +74,15 @@ class DatasetSpec: qa_case_builder: CaseBuilder retrieval_loader: RetrievalLoader | None = None retrieval_mapper: RetrievalMapper | None = None - retrieval_evaluator: Evaluator | None = None + retrieval_evaluators: list[Evaluator] | None = None + citation_evaluator: Evaluator | None = None qa_evaluator: Evaluator | None = None document_limit: int | None = None + retrieval_limit: int = 5 + ingest_batch_size: int | None = None + evaluate_refusal: bool = False + live: bool = False + experiment_metadata: dict[str, Any] | None = None def db_path(self, override_path: Path | None = None) -> Path: """Get the database path. diff --git a/evaluations/evaluations/datasets/__init__.py b/evaluations/evaluations/datasets/__init__.py index ee3f425a..e3532304 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/evaluations/evaluations/datasets/__init__.py @@ -1,6 +1,11 @@ from evaluations.config import DatasetSpec from .hotpotqa import HOTPOTQA_SPEC +from .mtrag import ( + MTRAG_CLAPNQ_LIVE_SPEC, + MTRAG_CLAPNQ_REWRITE_SPEC, + MTRAG_CLAPNQ_SPEC, +) from .open_rag_bench import ( ORB_MULTIMODAL_NEMOTRON_SPEC, ORB_MULTIMODAL_SPEC, @@ -12,6 +17,9 @@ DATASETS: dict[str, DatasetSpec] = { spec.key: spec for spec in ( HOTPOTQA_SPEC, + MTRAG_CLAPNQ_SPEC, + MTRAG_CLAPNQ_REWRITE_SPEC, + MTRAG_CLAPNQ_LIVE_SPEC, ORB_TEXT_SPEC, ORB_MULTIMODAL_SPEC, ORB_MULTIMODAL_NEMOTRON_SPEC, diff --git a/evaluations/evaluations/datasets/hotpotqa.py b/evaluations/evaluations/datasets/hotpotqa.py index d47f939e..f38d596f 100644 --- a/evaluations/evaluations/datasets/hotpotqa.py +++ b/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(), ) diff --git a/evaluations/evaluations/datasets/mtrag.py b/evaluations/evaluations/datasets/mtrag.py new file mode 100644 index 00000000..205bef4b --- /dev/null +++ b/evaluations/evaluations/datasets/mtrag.py @@ -0,0 +1,298 @@ +import json +import zipfile +from collections.abc import Iterable, Mapping +from functools import partial +from pathlib import Path +from typing import Any + +import httpx +from datasets import Dataset +from pydantic_evals import Case + +from evaluations.config import ( + ConversationInput, + DatasetSpec, + DocumentPayload, + RetrievalSample, + Turn, +) +from evaluations.evaluators import ( + CitationMAPEvaluator, + MAPEvaluator, + NDCGEvaluator, + RecallEvaluator, +) + +REPO_SHA = "cc5b1d481b391181b89f7ced860308482e785463" +_BASE_URL = f"https://raw.githubusercontent.com/IBM/mt-rag-benchmark/{REPO_SHA}" + +_CORPUS_FILE = "corpora/passage_level/clapnq.jsonl.zip" +_QRELS_FILE = "mtrag-human/retrieval_tasks/clapnq/qrels/dev.tsv" +_QUERY_FILES = { + "lastturn": "mtrag-human/retrieval_tasks/clapnq/clapnq_lastturn.jsonl", + "rewrite": "mtrag-human/retrieval_tasks/clapnq/clapnq_rewrite.jsonl", +} +_GEN_TASKS_FILE = "mtrag-human/generation_tasks/reference.jsonl" +_CLAPNQ_COLLECTION = "mt-rag-clapnq-elser-512-100-20240503" + + +def get_cache_dir() -> Path: + cache_dir = Path.home() / ".cache" / "haiku.rag" / "evaluations" / "mtrag" + cache_dir.mkdir(parents=True, exist_ok=True) + return cache_dir + + +def _download(rel_path: str) -> Path: + dest = get_cache_dir() / rel_path.replace("/", "_") + if dest.exists(): + return dest + + with httpx.stream( + "GET", f"{_BASE_URL}/{rel_path}", timeout=120.0, follow_redirects=True + ) as response: + response.raise_for_status() + tmp = dest.with_suffix(dest.suffix + ".part") + with tmp.open("wb") as fh: + for data in response.iter_bytes(): + fh.write(data) + tmp.rename(dest) + return dest + + +def _parse_qrels(lines: Iterable[str]) -> dict[str, list[str]]: + """Group qrel corpus-ids by query-id, preserving file order.""" + qrels: dict[str, list[str]] = {} + rows = iter(lines) + next(rows) # header: query-id / corpus-id / score + for line in rows: + if not line.strip(): + continue + query_id, corpus_id, _score = line.rstrip("\n").split("\t") + qrels.setdefault(query_id, []).append(corpus_id) + return qrels + + +def _validate_qrels_resolve( + corpus_ids: set[str], qrels: Mapping[str, list[str]] +) -> None: + unresolved = sorted( + {cid for ids in qrels.values() for cid in ids if cid not in corpus_ids} + ) + if unresolved: + raise ValueError( + f"{len(unresolved)} qrel corpus-ids do not resolve to corpus " + f"passages, e.g. {unresolved[:3]}" + ) + + +def _join_queries_qrels( + queries: Iterable[Mapping[str, Any]], qrels: Mapping[str, list[str]] +) -> list[dict[str, Any]]: + records = [] + for query in queries: + query_id = query["_id"] + expected = qrels.get(query_id) + if expected is None: + raise ValueError(f"query {query_id} has no qrels") + records.append( + { + "query_id": query_id, + "question": query["text"], + "expected_uris": expected, + } + ) + return records + + +def _load_qrels() -> dict[str, list[str]]: + path = _download(_QRELS_FILE) + return _parse_qrels(path.read_text().splitlines()) + + +def load_clapnq_corpus() -> Dataset: + path = _download(_CORPUS_FILE) + records: list[dict[str, str]] = [] + with zipfile.ZipFile(path) as zf: + with zf.open(zf.namelist()[0]) as fh: + for line in fh: + rec = json.loads(line) + records.append( + {"_id": rec["_id"], "title": rec["title"], "text": rec["text"]} + ) + _validate_qrels_resolve({rec["_id"] for rec in records}, _load_qrels()) + return Dataset.from_list(records) + + +def map_mtrag_document(doc: Mapping[str, Any]) -> DocumentPayload: + return DocumentPayload(uri=doc["_id"], content=doc["text"], title=doc["title"]) + + +def load_clapnq_retrieval(variant: str) -> Dataset: + path = _download(_QUERY_FILES[variant]) + queries = [json.loads(line) for line in path.read_text().splitlines() if line] + return Dataset.from_list(_join_queries_qrels(queries, _load_qrels())) + + +def map_mtrag_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None: + return RetrievalSample( + question=doc["question"], + expected_uris=tuple(doc["expected_uris"]), + ) + + +def _task_to_record( + task: Mapping[str, Any], qrels: Mapping[str, list[str]] +) -> dict[str, Any] | None: + """Reduce a reference.jsonl generation task to the fields QA cases need. + + Task `contexts` are the original system's retrievals, never gold relevance; + gold passages come from the qrels keyed by task_id. + """ + if task["Collection"] != _CLAPNQ_COLLECTION: + return None + return { + "id": task["task_id"], + "turn": task["turn"], + "turns": [ + {"speaker": message["speaker"], "text": message["text"]} + for message in task["input"] + ], + "answer": task["targets"][0]["text"], + "answerability": task["Answerability"][0], + "multi_turn_type": task["Multi-Turn"][0], + "question_type": list(task["Question Type"]), + "relevant_uris": qrels.get(task["task_id"]), + } + + +def load_clapnq_qa() -> Dataset: + path = _download(_GEN_TASKS_FILE) + qrels = _load_qrels() + records = [] + for line in path.read_text().splitlines(): + if not line.strip(): + continue + record = _task_to_record(json.loads(line), qrels) + if record is not None: + records.append(record) + return Dataset.from_list(records) + + +def build_mtrag_case( + index: int, doc: Mapping[str, Any] +) -> Case[ConversationInput, str, dict[str, Any]]: + metadata: dict[str, Any] = { + "task_id": doc["id"], + "turn": doc["turn"], + "answerability": doc["answerability"], + "multi_turn_type": doc["multi_turn_type"], + "question_type": list(doc["question_type"]), + } + if doc["relevant_uris"]: + metadata["relevant_uris"] = list(doc["relevant_uris"]) + return Case( + name=f"{index}_{doc['id']}", + inputs=ConversationInput( + turns=[Turn(**turn) for turn in doc["turns"]], + ), + expected_output=doc["answer"], + metadata=metadata, + ) + + +def _group_conversations(records: list[dict[str, Any]]) -> list[dict[str, Any]]: + """Group per-turn generation records into full conversations. + + Turns are ordered numerically within each conversation; each turn carries + its user question, reference answer, answerability label, and gold + passages when the turn has qrels. + """ + grouped: dict[str, list[dict[str, Any]]] = {} + for record in records: + conversation_id = record["id"].split("<::>")[0] + grouped.setdefault(conversation_id, []).append(record) + + conversations = [] + for conversation_id, tasks in grouped.items(): + tasks.sort(key=lambda record: int(record["turn"])) + turns = [] + for task in tasks: + turn: dict[str, Any] = { + "task_id": task["id"], + "turn": task["turn"], + "question": task["turns"][-1]["text"], + "reference": task["answer"], + "answerability": task["answerability"], + "multi_turn_type": task["multi_turn_type"], + "question_type": list(task["question_type"]), + } + if task["relevant_uris"]: + turn["relevant_uris"] = list(task["relevant_uris"]) + turns.append(turn) + conversations.append({"id": conversation_id, "turns": turns}) + return conversations + + +def load_clapnq_conversations() -> Dataset: + corpus = load_clapnq_qa() + return Dataset.from_list(_group_conversations([dict(row) for row in corpus])) + + +def build_mtrag_live_case( + index: int, doc: Mapping[str, Any] +) -> Case[list[str], list[str], dict[str, Any]]: + questions = [turn["question"] for turn in doc["turns"]] + metadata_turns = [ + { + key: value + for key, value in turn.items() + if key != "question" and value is not None + } + for turn in doc["turns"] + ] + return Case( + name=f"{index}_{doc['id']}", + inputs=questions, + metadata={"conversation_id": doc["id"], "turns": metadata_turns}, + ) + + +def _mtrag_spec(key: str, variant: str) -> DatasetSpec: + return DatasetSpec( + key=key, + db_filename="mtrag_clapnq.lancedb", + document_loader=load_clapnq_corpus, + document_mapper=map_mtrag_document, + qa_loader=load_clapnq_qa, + qa_case_builder=build_mtrag_case, + retrieval_loader=partial(load_clapnq_retrieval, variant), + retrieval_mapper=map_mtrag_retrieval, + retrieval_evaluators=[ + RecallEvaluator(k=5), + RecallEvaluator(k=10), + NDCGEvaluator(k=5), + NDCGEvaluator(k=10), + MAPEvaluator(), + ], + citation_evaluator=CitationMAPEvaluator(), + retrieval_limit=10, + ingest_batch_size=512, + evaluate_refusal=True, + experiment_metadata={"mtrag_mode": "gold_prefix"}, + ) + + +MTRAG_CLAPNQ_SPEC = _mtrag_spec("mtrag_clapnq", "lastturn") +MTRAG_CLAPNQ_REWRITE_SPEC = _mtrag_spec("mtrag_clapnq_rewrite", "rewrite") + +MTRAG_CLAPNQ_LIVE_SPEC = DatasetSpec( + key="mtrag_clapnq_live", + db_filename="mtrag_clapnq.lancedb", + document_loader=load_clapnq_corpus, + document_mapper=map_mtrag_document, + qa_loader=load_clapnq_conversations, + qa_case_builder=build_mtrag_live_case, + ingest_batch_size=512, + live=True, + experiment_metadata={"mtrag_mode": "live_session"}, +) diff --git a/evaluations/evaluations/datasets/open_rag_bench.py b/evaluations/evaluations/datasets/open_rag_bench.py index f20f48d9..8ffbf6df 100644 --- a/evaluations/evaluations/datasets/open_rag_bench.py +++ b/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(), ) diff --git a/evaluations/evaluations/datasets/t2_ragbench.py b/evaluations/evaluations/datasets/t2_ragbench.py index 5f172e9c..7fe64319 100644 --- a/evaluations/evaluations/datasets/t2_ragbench.py +++ b/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(), ) diff --git a/evaluations/evaluations/evaluators/__init__.py b/evaluations/evaluations/evaluators/__init__.py index d90f7449..fd3fe169 100644 --- a/evaluations/evaluations/evaluators/__init__.py +++ b/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,21 @@ from evaluations.evaluators.judge import ( ) from evaluations.evaluators.map import MAPEvaluator from evaluations.evaluators.number_match import NumberMatchEvaluator +from evaluations.evaluators.refusal import REFUSAL_RUBRIC, RefusalJudge +from evaluations.evaluators.retrieval import NDCGEvaluator, RecallEvaluator +from evaluations.evaluators.transcript import TranscriptLLMJudge __all__ = [ "ANSWER_EQUIVALENCE_RUBRIC", + "REFUSAL_RUBRIC", "CitationMAPEvaluator", + "ConversationEvaluator", "LLMJudge", "LLMJudgeResponseSchema", "MAPEvaluator", + "NDCGEvaluator", "NumberMatchEvaluator", + "RecallEvaluator", + "RefusalJudge", + "TranscriptLLMJudge", ] diff --git a/evaluations/evaluations/evaluators/citation.py b/evaluations/evaluations/evaluators/citation.py index 7a528288..f1f0cbf0 100644 --- a/evaluations/evaluations/evaluators/citation.py +++ b/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) diff --git a/evaluations/evaluations/evaluators/conversation.py b/evaluations/evaluations/evaluators/conversation.py new file mode 100644 index 00000000..8b97c60f --- /dev/null +++ b/evaluations/evaluations/evaluators/conversation.py @@ -0,0 +1,120 @@ +from dataclasses import dataclass + +from pydantic_ai import models +from pydantic_evals.evaluators import Evaluator, EvaluatorContext +from pydantic_evals.evaluators.evaluator import EvaluationReason, EvaluatorOutput +from pydantic_evals.evaluators.llm_as_a_judge import ( + judge_input_output_expected, + judge_output, +) + +from evaluations.evaluators.citation import average_precision +from evaluations.evaluators.refusal import REFUSAL_RUBRIC + +_REFUSAL_LABELS = ("ANSWERABLE", "UNANSWERABLE") + + +@dataclass +class ConversationEvaluator(Evaluator): + """Score a live-session conversation turn by turn. + + Expects the case output to be the list of per-turn answers, case inputs + the list of user questions, ``metadata["turns"]`` the per-turn reference, + answerability label, and optional gold ``relevant_uris``, and the + ``turn_cited_uris`` attribute the per-turn cited URIs. + + Each turn's answer is judged against the reference with the conversation + so far — including the model's own earlier answers — as context. Citation + AP is computed on turns with gold passages; refusal on ANSWERABLE and + UNANSWERABLE turns. Returned counts allow micro aggregation across + conversations; ``turn_pass_rate`` is the per-conversation (macro) rate. + Per-turn verdicts are returned as ``turn_{n}_pass`` (with the judge's + reason), ``turn_{n}_refused``, and ``turn_{n}_cited_ap`` for diagnosis. + """ + + rubric: str + model: models.Model | models.KnownModelName | str | None = None + + async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput: + questions: list[str] = list(ctx.inputs) + answers: list[str] = list(ctx.output) + turns: list[dict] = (ctx.metadata or {}).get("turns", []) + turn_cited: list[list[str]] = list( + ctx.attributes.get("turn_cited_uris") or [[] for _ in answers] + ) + if not (len(questions) == len(answers) == len(turns) == len(turn_cited)): + raise ValueError( + f"conversation arrays disagree: {len(questions)} questions, " + f"{len(answers)} answers, {len(turns)} turn annotations, " + f"{len(turn_cited)} citation lists" + ) + + passed = 0 + judged = 0 + citation_scores: list[float] = [] + true_refusals = 0 + false_refusals = 0 + unanswerable = 0 + per_turn: dict[str, EvaluationReason | bool | float | str] = {} + + transcript_lines: list[str] = [] + for index, (question, answer, turn) in enumerate( + zip(questions, answers, turns) + ): + number = index + 1 + transcript_lines.append(f"user: {question}") + transcript = "\n".join(transcript_lines) + transcript_lines.append(f"agent: {answer}") + + try: + grading = await judge_input_output_expected( + transcript, answer, turn["reference"], self.rubric, self.model + ) + except Exception as error: + per_turn[f"turn_{number}_judge_error"] = str(error)[:200] + else: + judged += 1 + if grading.pass_: + passed += 1 + per_turn[f"turn_{number}_pass"] = EvaluationReason( + value=grading.pass_, reason=grading.reason + ) + + label = turn.get("answerability") + if label in _REFUSAL_LABELS: + try: + refused = ( + await judge_output(answer, REFUSAL_RUBRIC, self.model) + ).pass_ + except Exception as error: + per_turn[f"turn_{number}_judge_error"] = str(error)[:200] + else: + per_turn[f"turn_{number}_refused"] = refused + if label == "UNANSWERABLE": + unanswerable += 1 + if refused: + true_refusals += 1 + elif refused: + false_refusals += 1 + + relevant = set(turn.get("relevant_uris") or []) + if relevant: + turn_ap = average_precision(turn_cited[index], relevant) + citation_scores.append(turn_ap) + per_turn[f"turn_{number}_cited_ap"] = turn_ap + + total = len(answers) + result: dict[str, EvaluationReason | bool | float | int | str] = { + "turn_pass_rate": passed / judged if judged else 0.0, + "turns_passed": passed, + "turns_judged": judged, + "turns_total": total, + "cited_eligible": len(citation_scores), + "true_refusals": true_refusals, + "false_refusals": false_refusals, + "unanswerable_turns": unanswerable, + } + if citation_scores: + result["cited_map"] = sum(citation_scores) / len(citation_scores) + result.update(per_turn) + return result diff --git a/evaluations/evaluations/evaluators/refusal.py b/evaluations/evaluations/evaluators/refusal.py new file mode 100644 index 00000000..ede0b4a5 --- /dev/null +++ b/evaluations/evaluations/evaluators/refusal.py @@ -0,0 +1,28 @@ +from dataclasses import dataclass + +from pydantic_evals.evaluators import EvaluatorContext, LLMJudge +from pydantic_evals.evaluators.evaluator import EvaluatorOutput + +REFUSAL_RUBRIC = ( + "The output declines to answer: it states that the information is not " + "available, cannot be found, or is not known, rather than providing a " + "substantive answer." +) + +_ELIGIBLE_LABELS = ("ANSWERABLE", "UNANSWERABLE") + + +@dataclass +class RefusalJudge(LLMJudge): + """Judges whether the output declines to answer. + + Runs only on turns whose answerability label makes refusal meaningful + (ANSWERABLE/UNANSWERABLE); other turns produce no evaluation and cost no + judge call. + """ + + async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput: + label = (ctx.metadata or {}).get("answerability") + if label not in _ELIGIBLE_LABELS: + return {} + return await super().evaluate(ctx) diff --git a/evaluations/evaluations/evaluators/retrieval.py b/evaluations/evaluations/evaluators/retrieval.py new file mode 100644 index 00000000..11684627 --- /dev/null +++ b/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 diff --git a/evaluations/evaluations/evaluators/transcript.py b/evaluations/evaluations/evaluators/transcript.py new file mode 100644 index 00000000..ec0d7d8d --- /dev/null +++ b/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) diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 97ff0cc8..1edad597 100644 --- a/evaluations/tests/test_benchmark.py +++ b/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,378 @@ class TestResolveDataset: _resolve_dataset("nonexistent") +class TestConversationInputDispatch: + @pytest.mark.asyncio + async def test_prefix_rides_as_message_history(self, tmp_path: Path) -> None: + """A ConversationInput case reaches the capability as final question + plus the prefix converted to message history.""" + from dataclasses import dataclass + + from pydantic_evals import Case + from pydantic_evals.evaluators import Evaluator, EvaluatorContext + + from evaluations.capability_runner import CapabilityRunResult + from evaluations.config import ConversationInput, Turn + + @dataclass + class AlwaysOne(Evaluator): + def evaluate(self, ctx: EvaluatorContext) -> float: + return 1.0 + + def build_case(idx: int, doc) -> Case: + return Case( + name="c1", + inputs=ConversationInput( + turns=[ + Turn(speaker="user", text="q1"), + Turn(speaker="agent", text="a1"), + Turn(speaker="user", text="q2"), + ] + ), + expected_output="ref", + ) + + spec = DatasetSpec( + key="test", + db_filename="test.lancedb", + document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + document_mapper=lambda doc: None, + qa_loader=lambda: [{"id": "t1"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + qa_case_builder=build_case, + qa_evaluator=AlwaysOne(), + ) + + with ( + patch("evaluations.benchmark.get_model", return_value="fake-model"), + patch( + "evaluations.benchmark.run_capability_question", + new_callable=AsyncMock, + return_value=CapabilityRunResult(answer="answer"), + ) as run_question, + ): + await run_qa_benchmark(spec, AppConfig(), db_path=tmp_path / "test.lancedb") + + assert run_question.await_args is not None + kwargs = run_question.await_args.kwargs + assert kwargs["question"] == "q2" + history = kwargs["message_history"] + assert len(history) == 2 + assert history[0].parts[0].content == "q1" + assert history[1].parts[0].content == "a1" + + +class TestRefusalMetrics: + def _case(self, label: str | None, refused: bool | None) -> MagicMock: + case = MagicMock() + case.metadata = {"answerability": label} if label is not None else {} + case.assertions = ( + {"refused": MagicMock(value=refused)} if refused is not None else {} + ) + return case + + def test_precision_and_recall(self) -> None: + from evaluations.benchmark import _refusal_metrics + + cases = [ + self._case("UNANSWERABLE", True), # true refusal + self._case("UNANSWERABLE", False), # missed refusal + self._case("ANSWERABLE", True), # false refusal + self._case("ANSWERABLE", False), # answered correctly + self._case("PARTIAL", None), # skipped by the judge, no assertion + self._case(None, None), # no label + ] + + metrics = _refusal_metrics(cases) + + assert metrics is not None + precision, recall, unanswerable, refusals = metrics + assert precision == 0.5 # 1 true refusal of 2 refusals + assert recall == 0.5 # 1 of 2 unanswerable turns refused + assert unanswerable == 2 + assert refusals == 2 + + def test_none_when_no_judged_cases(self) -> None: + from evaluations.benchmark import _refusal_metrics + + assert _refusal_metrics([self._case("PARTIAL", None)]) is None + + +class TestLiveSummary: + def _case(self, scores: dict[str, float | int]) -> MagicMock: + case = MagicMock() + case.scores = {key: MagicMock(value=value) for key, value in scores.items()} + return case + + def test_micro_and_macro_aggregation(self) -> None: + from evaluations.benchmark import _live_summary + + # Conversation A: 1/4 turns pass; B: 2/2 pass. Micro weights turns + # (3/6); macro averages conversations ((0.25 + 1.0) / 2). + cases = [ + self._case( + { + "turn_pass_rate": 0.25, + "turns_passed": 1, + "turns_judged": 4, + "turns_total": 4, + "cited_map": 0.5, + "cited_eligible": 3, + "true_refusals": 1, + "false_refusals": 1, + "unanswerable_turns": 2, + } + ), + self._case( + { + "turn_pass_rate": 1.0, + "turns_passed": 2, + "turns_judged": 2, + "turns_total": 2, + "cited_map": 1.0, + "cited_eligible": 1, + "true_refusals": 0, + "false_refusals": 0, + "unanswerable_turns": 0, + } + ), + ] + + failure = MagicMock() + failure.inputs = ["fq1", "fq2", "fq3"] + summary = _live_summary(cases, [failure]) + + assert summary is not None + assert summary["conversations"] == 2 + assert summary["conversations_attempted"] == 3 + assert summary["turns_total"] == 6 + assert summary["turns_judged"] == 6 + assert summary["turns_attempted"] == 9 + assert summary["micro_pass_rate"] == pytest.approx(0.5) + assert summary["macro_pass_rate"] == pytest.approx(0.625) + assert summary["cited_eligible"] == 4 + assert summary["cited_map_micro"] == pytest.approx((0.5 * 3 + 1.0 * 1) / 4) + assert summary["cited_map_macro"] == pytest.approx(0.75) + assert summary["refusal_precision"] == pytest.approx(0.5) + assert summary["refusal_recall"] == pytest.approx(0.5) + + def test_none_without_scored_cases(self) -> None: + from evaluations.benchmark import _live_summary + + assert _live_summary([self._case({})]) is None + + def test_micro_rate_uses_judged_turns(self) -> None: + from evaluations.benchmark import _live_summary + + cases = [ + self._case( + { + "turn_pass_rate": 1.0, + "turns_passed": 3, + "turns_judged": 3, + "turns_total": 4, # one turn's judge errored + "cited_eligible": 0, + "true_refusals": 0, + "false_refusals": 0, + "unanswerable_turns": 0, + } + ) + ] + + summary = _live_summary(cases) + + assert summary is not None + assert summary["micro_pass_rate"] == 1.0 + assert summary["turns_judged"] == 3 + assert summary["turns_total"] == 4 + + def test_failed_conversations_do_not_affect_rates(self) -> None: + from evaluations.benchmark import _live_summary + + cases = [ + self._case( + { + "turn_pass_rate": 1.0, + "turns_passed": 2, + "turns_judged": 2, + "turns_total": 2, + "cited_eligible": 0, + "true_refusals": 0, + "false_refusals": 0, + "unanswerable_turns": 0, + } + ) + ] + failure = MagicMock() + failure.inputs = ["fq1", "fq2"] + + summary = _live_summary(cases, [failure]) + + assert summary is not None + assert summary["micro_pass_rate"] == 1.0 + assert summary["macro_pass_rate"] == 1.0 + assert summary["conversations_attempted"] == 2 + assert summary["turns_attempted"] == 4 + + +class TestLiveConversationDispatch: + @pytest.mark.asyncio + async def test_live_spec_replays_conversation(self, tmp_path: Path) -> None: + from pydantic_evals import Case + + from evaluations.benchmark import run_live_qa_benchmark + from evaluations.capability_runner import CapabilityRunResult + + def build_case(idx: int, doc) -> Case: + return Case( + name="conv1", + inputs=["q1", "q2"], + metadata={ + "conversation_id": "conv1", + "turns": [ + {"reference": "r1", "answerability": "ANSWERABLE"}, + {"reference": "r2", "answerability": "ANSWERABLE"}, + ], + }, + ) + + spec = DatasetSpec( + key="test_live", + db_filename="test.lancedb", + document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + document_mapper=lambda doc: None, + qa_loader=lambda: [{"id": "conv1"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + qa_case_builder=build_case, + live=True, + ) + + turn_results = [ + CapabilityRunResult(answer="a1", cited_uris=["u1"]), + CapabilityRunResult(answer="a2", cited_uris=[]), + ] + with ( + patch("evaluations.benchmark.get_model", return_value="fake-model"), + patch( + "evaluations.benchmark.run_capability_conversation", + new_callable=AsyncMock, + return_value=turn_results, + ) as run_conversation, + patch( + "evaluations.evaluators.conversation.judge_input_output_expected", + new_callable=AsyncMock, + return_value=MagicMock(score=None, pass_=True, reason=None), + ), + patch( + "evaluations.evaluators.conversation.judge_output", + new_callable=AsyncMock, + return_value=MagicMock(score=None, pass_=False, reason=None), + ), + ): + await run_live_qa_benchmark( + spec, AppConfig(), db_path=tmp_path / "test.lancedb" + ) + + assert run_conversation.await_args is not None + assert run_conversation.await_args.kwargs["questions"] == ["q1", "q2"] + + @pytest.mark.asyncio + async def test_live_records_per_turn_traffic_arrays(self, tmp_path: Path) -> None: + """Per-turn tool traffic is recorded as question-length arrays, in the + same list-indexed-by-turn shape as turn_cited_uris.""" + from pydantic_evals import Case + + from evaluations.benchmark import run_live_qa_benchmark + from evaluations.capability_runner import CapabilityRunResult + + def build_case(idx: int, doc) -> Case: + return Case( + name="conv1", + inputs=["q1", "q2"], + metadata={ + "conversation_id": "conv1", + "turns": [ + {"reference": "r1", "answerability": "ANSWERABLE"}, + {"reference": "r2", "answerability": "ANSWERABLE"}, + ], + }, + ) + + spec = DatasetSpec( + key="test_live", + db_filename="test.lancedb", + document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + document_mapper=lambda doc: None, + qa_loader=lambda: [{"id": "conv1"}], # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + qa_case_builder=build_case, + live=True, + ) + + turn_results = [ + CapabilityRunResult( + answer="a1", + cited_uris=["u1"], + n_search_calls=2, + n_rejected_searches=1, + n_failed_tools=1, + n_requests=4, + ), + CapabilityRunResult(answer="a2"), + ] + recorded: dict[str, object] = {} + + with ( + patch("evaluations.benchmark.get_model", return_value="fake-model"), + patch( + "evaluations.benchmark.set_eval_attribute", + side_effect=lambda key, value: recorded.__setitem__(key, value), + ), + patch( + "evaluations.benchmark.run_capability_conversation", + new_callable=AsyncMock, + return_value=turn_results, + ), + patch( + "evaluations.evaluators.conversation.judge_input_output_expected", + new_callable=AsyncMock, + return_value=MagicMock(score=None, pass_=True, reason=None), + ), + patch( + "evaluations.evaluators.conversation.judge_output", + new_callable=AsyncMock, + return_value=MagicMock(score=None, pass_=False, reason=None), + ), + ): + await run_live_qa_benchmark( + spec, AppConfig(), db_path=tmp_path / "test.lancedb" + ) + + assert recorded["turn_n_search_calls"] == [2, 0] + assert recorded["turn_n_rejected_searches"] == [1, 0] + assert recorded["turn_n_failed_tools"] == [1, 0] + assert recorded["turn_n_requests"] == [4, 0] + questions = 2 + for key, value in recorded.items(): + if key.startswith("turn_"): + assert isinstance(value, list) and len(value) == questions, key + + +class TestResolveDatasets: + def test_all_dedupes_shared_databases(self) -> None: + """Specs sharing a db_filename (mtrag query variants) appear once, so + `download all`/`upload all` do not process the same DB twice.""" + from evaluations.benchmark import _resolve_datasets + + specs = _resolve_datasets("all") + filenames = [spec.db_filename for spec in specs] + assert len(filenames) == len(set(filenames)) + assert "mtrag_clapnq.lancedb" in filenames + + def test_single_key_not_deduped(self) -> None: + from evaluations.benchmark import _resolve_datasets + + specs = _resolve_datasets("mtrag_clapnq_rewrite") + assert [spec.key for spec in specs] == ["mtrag_clapnq_rewrite"] + + class TestLoadConfig: def test_explicit_path(self, tmp_path: Path) -> None: config_file = tmp_path / "test.yaml" @@ -383,17 +755,117 @@ class TestRunQaBenchmarkCapabilityTarget: class TestCitationEvaluatorWiring: - def test_returns_map_twin_for_map_evaluator(self) -> None: - from evaluations.benchmark import _citation_evaluator_for - from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator + def test_specs_with_retrieval_declare_citation_evaluator(self) -> None: + """Citation scoring is declared per spec, not inferred: every dataset + that scores retrieval also scores citations.""" + from evaluations.datasets import DATASETS + from evaluations.evaluators import CitationMAPEvaluator - result = _citation_evaluator_for(MAPEvaluator()) - assert isinstance(result, CitationMAPEvaluator) + for spec in DATASETS.values(): + if spec.retrieval_evaluators: + assert isinstance(spec.citation_evaluator, CitationMAPEvaluator), ( + spec.key + ) - def test_returns_none_for_no_evaluator(self) -> None: - from evaluations.benchmark import _citation_evaluator_for - assert _citation_evaluator_for(None) is None +class TestBatchedIngest: + def _rag( + self, + complete_uris: list[str] | None = None, + chunkless_uris: list[str] | None = None, + ) -> MagicMock: + complete_uris = complete_uris or [] + chunkless_uris = chunkless_uris or [] + + def _table(rows: list[dict]) -> MagicMock: + table = MagicMock() + table.query.return_value.select.return_value.to_list = AsyncMock( + return_value=rows + ) + return table + + rag = MagicMock() + rag.store.document_meta_table = _table( + [{"id": f"id-{uri}", "uri": uri} for uri in complete_uris + chunkless_uris] + ) + rag.store.chunks_table = _table( + [{"document_id": f"id-{uri}"} for uri in complete_uris] + ) + rag.convert = AsyncMock(side_effect=lambda content, **kw: f"docling:{content}") + rag.chunk = AsyncMock(return_value=[]) + rag.import_documents = AsyncMock() + rag.delete_document = AsyncMock() + return rag + + def _spec(self) -> DatasetSpec: + return DatasetSpec( + key="test", + db_filename="test.lancedb", + document_loader=lambda: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + document_mapper=lambda doc: ( + None + if doc["uri"] == "bad" + else DocumentPayload(uri=doc["uri"], content=f"text {doc['uri']}") + ), + qa_loader=lambda: [], # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + qa_case_builder=lambda idx, doc: None, # type: ignore[arg-type] # ty: ignore[invalid-argument-type] + ) + + @pytest.mark.asyncio + async def test_imports_in_bounded_batches(self) -> None: + from evaluations.benchmark import _ingest_batched + + rag = self._rag() + corpus = [{"uri": f"u{i}"} for i in range(5)] + + await _ingest_batched(rag, self._spec(), corpus, batch_size=2) + + batch_uris = [ + [imp.uri for imp in call.args[0]] + for call in rag.import_documents.call_args_list + ] + assert batch_uris == [["u0", "u1"], ["u2", "u3"], ["u4"]] + + @pytest.mark.asyncio + async def test_resume_skips_complete_uris(self) -> None: + from evaluations.benchmark import _ingest_batched + + rag = self._rag(complete_uris=["u0", "u2"]) + corpus = [{"uri": f"u{i}"} for i in range(4)] + + await _ingest_batched(rag, self._spec(), corpus, batch_size=10) + + (batch,), _ = rag.import_documents.call_args + assert [imp.uri for imp in batch] == ["u1", "u3"] + assert rag.convert.await_count == 2 + rag.delete_document.assert_not_awaited() + + @pytest.mark.asyncio + async def test_resume_reimports_chunkless_documents(self) -> None: + """A crash between the document and chunk writes leaves a document + without chunks; resume must delete and re-import it, not skip it.""" + from evaluations.benchmark import _ingest_batched + + rag = self._rag(complete_uris=["u0"], chunkless_uris=["u1"]) + corpus = [{"uri": "u0"}, {"uri": "u1"}] + + await _ingest_batched(rag, self._spec(), corpus, batch_size=10) + + rag.delete_document.assert_awaited_once_with("id-u1") + (batch,), _ = rag.import_documents.call_args + assert [imp.uri for imp in batch] == ["u1"] + + @pytest.mark.asyncio + async def test_unmapped_documents_skipped(self) -> None: + from evaluations.benchmark import _ingest_batched + + rag = self._rag() + corpus = [{"uri": "u0"}, {"uri": "bad"}, {"uri": "u1"}] + + await _ingest_batched(rag, self._spec(), corpus, batch_size=10) + + (batch,), _ = rag.import_documents.call_args + assert [imp.uri for imp in batch] == ["u0", "u1"] class TestAttachRelevantUris: @@ -433,7 +905,7 @@ class TestAttachRelevantUris: retrieval_mapper=lambda d: RetrievalSample( question=d["q"], expected_uris=d["uris"] ), - retrieval_evaluator=MAPEvaluator(), + retrieval_evaluators=[MAPEvaluator()], ) _attach_relevant_uris(cases, spec, limit=None) @@ -511,7 +983,7 @@ class TestRetrievalTarget: retrieval_mapper=lambda d: RetrievalSample( question=d["q"], expected_uris=d["uris"] ), - retrieval_evaluator=MAPEvaluator(), + retrieval_evaluators=[MAPEvaluator()], ) @pytest.mark.asyncio diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index 9256d3d3..058ac734 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/evaluations/tests/test_capability_runner.py @@ -144,7 +144,9 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp with patch( "evaluations.capability_runner.Agent.run", new_callable=AsyncMock ) as run: - run.return_value = SimpleNamespace(output="done", all_messages=lambda: []) + run.return_value = SimpleNamespace( + output="done", all_messages=lambda: [], new_messages=lambda: [] + ) await run_capability_question( lambda **_kwargs: capability, @@ -157,3 +159,139 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp assert capability.request_limit == expected assert "usage_limits" not in run.call_args.kwargs + + +class TestPrefixToMessages: + def test_maps_turns_to_model_messages(self) -> None: + from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + TextPart, + UserPromptPart, + ) + + from evaluations.capability_runner import prefix_to_messages + from evaluations.config import Turn + + messages = prefix_to_messages( + [ + Turn(speaker="user", text="who takes photos of planes?"), + Turn(speaker="agent", text="Ground-to-air photographers."), + ] + ) + + assert len(messages) == 2 + assert isinstance(messages[0], ModelRequest) + assert isinstance(messages[0].parts[0], UserPromptPart) + assert messages[0].parts[0].content == "who takes photos of planes?" + assert isinstance(messages[1], ModelResponse) + assert isinstance(messages[1].parts[0], TextPart) + assert messages[1].parts[0].content == "Ground-to-air photographers." + + def test_empty_prefix(self) -> None: + from evaluations.capability_runner import prefix_to_messages + + assert prefix_to_messages([]) == [] + + +async def test_message_history_passed_to_agent_run(tmp_path): + from evaluations.capability_runner import prefix_to_messages + from evaluations.config import Turn + + history = prefix_to_messages([Turn(speaker="user", text="earlier question")]) + capability = create_rag( + db_path=tmp_path / "rag.lancedb", + config=AppConfig(), + defer_loading=False, + ) + with patch( + "evaluations.capability_runner.Agent.run", new_callable=AsyncMock + ) as run: + run.return_value = SimpleNamespace(output="done", new_messages=lambda: []) + + await run_capability_question( + lambda **_kwargs: capability, + tmp_path / "rag.lancedb", + AppConfig(), + "follow-up question", + TestModel(call_tools=[]), + message_history=history, + ) + + assert run.call_args.kwargs["message_history"] is history + + +async def test_conversation_threads_own_messages_across_turns(tmp_path): + """Each turn runs with the previous turn's full message history (including + tool traffic), so prior-turn compaction operates on real history.""" + from evaluations.capability_runner import run_capability_conversation + + capability = create_rag( + db_path=tmp_path / "rag.lancedb", + config=AppConfig(), + defer_loading=False, + ) + histories: list[object] = [] + + async def _run(question, deps=None, message_history=None): + histories.append(message_history) + return SimpleNamespace( + output=f"answer to {question}", + all_messages=lambda: [f"history after {question}"], + new_messages=lambda: [], + ) + + with patch("evaluations.capability_runner.Agent.run", side_effect=_run): + result = await run_capability_conversation( + lambda **_kwargs: capability, + tmp_path / "rag.lancedb", + AppConfig(), + ["q1", "q2", "q3"], + TestModel(call_tools=[]), + ) + + assert [t.answer for t in result] == [ + "answer to q1", + "answer to q2", + "answer to q3", + ] + assert histories == [None, ["history after q1"], ["history after q2"]] + + +async def test_conversation_end_to_end_with_test_model(tmp_path): + from evaluations.capability_runner import run_capability_conversation + + result = await run_capability_conversation( + create_rag, + tmp_path / "rag.lancedb", + AppConfig(), + ["first question", "follow-up"], + TestModel(call_tools=[]), + ) + + assert len(result) == 2 + assert all(turn.answer == "success (no tool calls)" for turn in result) + assert all(turn.cited_uris == [] for turn in result) + + +async def test_gold_prefix_run_answers_with_history(tmp_path): + """End-to-end through a real Agent: the prefix rides along as history.""" + from evaluations.capability_runner import prefix_to_messages + from evaluations.config import Turn + + history = prefix_to_messages( + [ + Turn(speaker="user", text="who takes photos of planes?"), + Turn(speaker="agent", text="Ground-to-air photographers."), + ] + ) + result = await run_capability_question( + create_rag, + tmp_path / "rag.lancedb", + AppConfig(), + "No, I meant photos in the air.", + TestModel(call_tools=[]), + message_history=history, + ) + + assert result.answer == "success (no tool calls)" diff --git a/evaluations/tests/test_citation_evaluators.py b/evaluations/tests/test_citation_evaluators.py index 31d0d6be..76d75a37 100644 --- a/evaluations/tests/test_citation_evaluators.py +++ b/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" diff --git a/evaluations/tests/test_config.py b/evaluations/tests/test_config.py index 37342701..38ad1c5c 100644 --- a/evaluations/tests/test_config.py +++ b/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: diff --git a/evaluations/tests/test_conversation_evaluator.py b/evaluations/tests/test_conversation_evaluator.py new file mode 100644 index 00000000..f485995f --- /dev/null +++ b/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", + } diff --git a/evaluations/tests/test_evaluators.py b/evaluations/tests/test_evaluators.py index c0adecd3..ed443e42 100644 --- a/evaluations/tests/test_evaluators.py +++ b/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() diff --git a/evaluations/tests/test_mtrag.py b/evaluations/tests/test_mtrag.py new file mode 100644 index 00000000..94a9cde4 --- /dev/null +++ b/evaluations/tests/test_mtrag.py @@ -0,0 +1,268 @@ +import pytest + +from evaluations.config import ConversationInput +from evaluations.datasets import DATASETS +from evaluations.datasets.mtrag import ( + MTRAG_CLAPNQ_LIVE_SPEC, + MTRAG_CLAPNQ_REWRITE_SPEC, + MTRAG_CLAPNQ_SPEC, + _group_conversations, + _join_queries_qrels, + _parse_qrels, + _task_to_record, + _validate_qrels_resolve, + build_mtrag_case, + build_mtrag_live_case, + map_mtrag_document, + map_mtrag_retrieval, +) +from evaluations.evaluators import ( + CitationMAPEvaluator, + MAPEvaluator, + NDCGEvaluator, + RecallEvaluator, +) + +GENERATION_TASK = { + "task_id": "conv1<::>2", + "conversation_id": "conv1", + "turn": "2", + "Collection": "mt-rag-clapnq-elser-512-100-20240503", + "Answerability": ["ANSWERABLE"], + "Multi-Turn": ["Follow-up"], + "Question Type": ["Factoid"], + "input": [ + {"speaker": "user", "text": "q1", "metadata": {}}, + {"speaker": "agent", "text": "a1", "metadata": {}}, + {"speaker": "user", "text": "q2", "metadata": {}}, + ], + "targets": [{"text": "reference answer"}], + "contexts": [{"document_id": "retrieved-not-gold"}], +} + + +class TestDocumentMapper: + def test_maps_passage_to_payload(self) -> None: + payload = map_mtrag_document( + {"_id": "837799097_6931-7548-0-617", "title": "T", "text": "body"} + ) + assert payload.uri == "837799097_6931-7548-0-617" + assert payload.title == "T" + assert payload.content == "body" + + +class TestQrels: + QRELS_TSV = ( + "query-id\tcorpus-id\tscore\n" + "conv1<::>2\tdoc1_0-10-0-10\t1\n" + "conv1<::>2\tdoc2_5-20-0-15\t1\n" + "conv2<::>1\tdoc3_0-9-0-9\t1\n" + ) + + def test_parse_groups_by_query_preserving_order(self) -> None: + qrels = _parse_qrels(self.QRELS_TSV.splitlines()) + assert qrels == { + "conv1<::>2": ["doc1_0-10-0-10", "doc2_5-20-0-15"], + "conv2<::>1": ["doc3_0-9-0-9"], + } + + def test_join_builds_records(self) -> None: + qrels = _parse_qrels(self.QRELS_TSV.splitlines()) + queries = [ + {"_id": "conv1<::>2", "text": "q one"}, + {"_id": "conv2<::>1", "text": "q two"}, + ] + records = _join_queries_qrels(queries, qrels) + assert records == [ + { + "query_id": "conv1<::>2", + "question": "q one", + "expected_uris": ["doc1_0-10-0-10", "doc2_5-20-0-15"], + }, + { + "query_id": "conv2<::>1", + "question": "q two", + "expected_uris": ["doc3_0-9-0-9"], + }, + ] + + def test_join_raises_on_query_without_qrels(self) -> None: + with pytest.raises(ValueError, match="no qrels"): + _join_queries_qrels([{"_id": "missing<::>1", "text": "q"}], {}) + + def test_validation_passes_when_all_resolve(self) -> None: + qrels = {"q1": ["a", "b"]} + _validate_qrels_resolve({"a", "b", "c"}, qrels) + + def test_validation_raises_on_unresolved_id(self) -> None: + qrels = {"q1": ["a", "ghost"]} + with pytest.raises(ValueError, match="ghost"): + _validate_qrels_resolve({"a"}, qrels) + + +class TestRetrievalMapper: + def test_maps_joined_record(self) -> None: + sample = map_mtrag_retrieval( + { + "query_id": "conv1<::>2", + "question": "who?", + "expected_uris": ["u1", "u2"], + } + ) + assert sample is not None + assert sample.question == "who?" + assert sample.expected_uris == ("u1", "u2") + + +class TestSpecs: + def test_registered(self) -> None: + assert DATASETS["mtrag_clapnq"] is MTRAG_CLAPNQ_SPEC + assert DATASETS["mtrag_clapnq_rewrite"] is MTRAG_CLAPNQ_REWRITE_SPEC + + def test_variants_share_db(self) -> None: + assert MTRAG_CLAPNQ_SPEC.db_filename == MTRAG_CLAPNQ_REWRITE_SPEC.db_filename + + def test_retrieval_configuration(self) -> None: + for spec in (MTRAG_CLAPNQ_SPEC, MTRAG_CLAPNQ_REWRITE_SPEC): + assert spec.retrieval_limit == 10 + assert spec.ingest_batch_size == 512 + assert spec.retrieval_evaluators is not None + kinds = { + (type(e), getattr(e, "k", None)) for e in spec.retrieval_evaluators + } + assert kinds == { + (RecallEvaluator, 5), + (RecallEvaluator, 10), + (NDCGEvaluator, 5), + (NDCGEvaluator, 10), + (MAPEvaluator, None), + } + assert isinstance(spec.citation_evaluator, CitationMAPEvaluator) + + def test_refusal_evaluation_enabled(self) -> None: + assert MTRAG_CLAPNQ_SPEC.evaluate_refusal is True + + +class TestGenerationTasks: + def test_task_to_record(self) -> None: + record = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p1", "p2"]}) + assert record == { + "id": "conv1<::>2", + "turn": "2", + "turns": [ + {"speaker": "user", "text": "q1"}, + {"speaker": "agent", "text": "a1"}, + {"speaker": "user", "text": "q2"}, + ], + "answer": "reference answer", + "answerability": "ANSWERABLE", + "multi_turn_type": "Follow-up", + "question_type": ["Factoid"], + "relevant_uris": ["p1", "p2"], + } + + def test_task_without_qrels_has_no_relevant_uris(self) -> None: + record = _task_to_record(GENERATION_TASK, {}) + assert record is not None + assert record["relevant_uris"] is None + + def test_other_collections_excluded(self) -> None: + task = {**GENERATION_TASK, "Collection": "mt-rag-govt-elser-512-100-20240611"} + assert _task_to_record(task, {}) is None + + def test_build_case_conversation_and_metadata(self) -> None: + record = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p1"]}) + assert record is not None + case = build_mtrag_case(3, record) + + assert isinstance(case.inputs, ConversationInput) + assert case.inputs.question == "q2" + assert [t.speaker for t in case.inputs.turns] == ["user", "agent", "user"] + assert case.expected_output == "reference answer" + assert case.metadata == { + "task_id": "conv1<::>2", + "turn": "2", + "answerability": "ANSWERABLE", + "multi_turn_type": "Follow-up", + "question_type": ["Factoid"], + "relevant_uris": ["p1"], + } + + def test_build_case_omits_relevant_uris_when_absent(self) -> None: + record = _task_to_record( + {**GENERATION_TASK, "Answerability": ["UNANSWERABLE"]}, {} + ) + assert record is not None + case = build_mtrag_case(1, record) + + assert case.metadata is not None + assert "relevant_uris" not in case.metadata + assert case.metadata["answerability"] == "UNANSWERABLE" + + +class TestLiveConversations: + def _records(self) -> list[dict]: + turn1 = _task_to_record( + { + **GENERATION_TASK, + "task_id": "conv1<::>1", + "turn": "1", + "input": [{"speaker": "user", "text": "q1", "metadata": {}}], + "targets": [{"text": "r1"}], + }, + {"conv1<::>1": ["p1"]}, + ) + turn2 = _task_to_record(GENERATION_TASK, {"conv1<::>2": ["p2", "p3"]}) + other = _task_to_record( + { + **GENERATION_TASK, + "task_id": "conv2<::>1", + "turn": "1", + "input": [{"speaker": "user", "text": "other q", "metadata": {}}], + "targets": [{"text": "other r"}], + "Answerability": ["UNANSWERABLE"], + }, + {}, + ) + assert turn1 and turn2 and other + # turn 2 first: grouping must sort turns numerically within a conversation + return [turn2, turn1, other] + + def test_grouping_sorts_turns_within_conversations(self) -> None: + conversations = _group_conversations(self._records()) + assert [c["id"] for c in conversations] == ["conv1", "conv2"] + conv1 = conversations[0] + assert [t["question"] for t in conv1["turns"]] == ["q1", "q2"] + assert [t["reference"] for t in conv1["turns"]] == ["r1", "reference answer"] + assert conv1["turns"][1]["relevant_uris"] == ["p2", "p3"] + + def test_build_live_case(self) -> None: + conversations = _group_conversations(self._records()) + case = build_mtrag_live_case(1, conversations[0]) + + assert case.inputs == ["q1", "q2"] + assert case.metadata is not None + assert case.metadata["conversation_id"] == "conv1" + turns = case.metadata["turns"] + assert turns[0] == { + "task_id": "conv1<::>1", + "turn": "1", + "reference": "r1", + "answerability": "ANSWERABLE", + "multi_turn_type": "Follow-up", + "question_type": ["Factoid"], + "relevant_uris": ["p1"], + } + other_case = build_mtrag_live_case(2, conversations[1]) + assert other_case.metadata is not None + assert "relevant_uris" not in other_case.metadata["turns"][0] + + def test_live_spec(self) -> None: + assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC + assert MTRAG_CLAPNQ_LIVE_SPEC.db_filename == MTRAG_CLAPNQ_SPEC.db_filename + assert MTRAG_CLAPNQ_LIVE_SPEC.live is True + assert MTRAG_CLAPNQ_LIVE_SPEC.retrieval_loader is None + assert MTRAG_CLAPNQ_LIVE_SPEC.experiment_metadata == { + "mtrag_mode": "live_session" + } + assert MTRAG_CLAPNQ_SPEC.experiment_metadata == {"mtrag_mode": "gold_prefix"} From db2b8fb883655f157e058eb9100409ab941b6970 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 22:52:36 +0300 Subject: [PATCH 3/9] Add compaction arms and grounding status to MTRAG live runs --- CHANGELOG.md | 2 +- docs/benchmarks.md | 6 +- evaluations/evaluations/benchmark.py | 3 + evaluations/evaluations/capability_runner.py | 30 ++++- evaluations/evaluations/config.py | 1 + evaluations/evaluations/datasets/__init__.py | 2 + evaluations/evaluations/datasets/mtrag.py | 29 ++-- evaluations/tests/test_benchmark.py | 50 +++++++ evaluations/tests/test_capability_runner.py | 133 ++++++++++++++++++- evaluations/tests/test_mtrag.py | 15 ++- 10 files changed, 251 insertions(+), 20 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index d370a087..5bcc7be9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added - `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata. -- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, and per-turn tool-traffic attributes. +- `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 diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 419e6d7f..dd73d320 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -37,7 +37,7 @@ Active datasets: | `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB | | `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB | | `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB | -| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite` and `mtrag_clapnq_live` keys | ~2.8 GB | +| `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): @@ -231,7 +231,7 @@ The reranker's contribution is larger here than on the single-doc datasets: hybr [MTRAG](https://github.com/IBM/mt-rag-benchmark) is IBM's multi-turn RAG benchmark (TACL 2025, SemEval-2026 Task 8): human-authored conversations with per-turn answerability labels and binary relevance judgments. We evaluate the ClapNQ (Wikipedia) domain: 183,408 passages, 29 conversations, 224 turns, 208 retrieval queries. -Three dataset keys share one database. `mtrag_clapnq` retrieves with the raw last user turn and runs QA by replaying each task's reference conversation prefix as message history. `mtrag_clapnq_rewrite` retrieves with the human standalone rewrites. `mtrag_clapnq_live` replays whole conversations through a single capability session, carrying the model's own answers and tool history across turns. +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) @@ -251,4 +251,4 @@ Directly comparable with [IBM's published results](https://github.com/IBM/mt-rag | Gold-prefix (`mtrag_clapnq`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 223 | 0.68 | 0.35 | | Live (`mtrag_clapnq_live`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 195/224 scored | 0.72 micro / 0.73 macro | 0.35 | -*Measured on haiku.rag v0.67.1 with `qwen3-embedding:4b` (vLLM, dim 2560) and `Qwen3-Reranker-4B`, stock capability instructions, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` at temperature 0. The judge sampling has since been re-pinned repo-wide (0.6 with thinking), so future runs re-baseline. QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Live mode additionally reports refusal precision/recall against the per-turn answerability labels and per-turn pass rates; pass rate declines with conversation depth (93% at turn 1 to 38% at turn 9). The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.* +*Measured on haiku.rag v0.67.1 with `qwen3-embedding:4b` (vLLM, dim 2560) and `Qwen3-Reranker-4B`, stock capability instructions, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` at temperature 0. Two changes since re-baseline these numbers: the judge sampling was re-pinned repo-wide (0.6 with thinking), and 0.74.0 rewrote the citation instructions (refusals now declare an empty citation list instead of being exempt). Retrieval is unaffected by both and remains the control. QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Live mode additionally reports refusal precision/recall against the per-turn answerability labels and per-turn pass rates; pass rate declines with conversation depth (93% at turn 1 to 38% at turn 9). The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.* diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index f254b67b..9b18f962 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -630,6 +630,7 @@ 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) @@ -778,6 +779,7 @@ async def run_live_qa_benchmark( config=config, questions=list(questions), capability_model=resolved_capability_model, + 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]) @@ -786,6 +788,7 @@ async def run_live_qa_benchmark( ) 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( diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index 04ff75d1..89e0e334 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -16,8 +16,12 @@ from pydantic_ai.messages import ( ) 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 @@ -40,6 +44,7 @@ class _RagLikeState(Protocol): document_filter: str | None citation_index: dict[str, Citation] citations: list[str] + evidence: CapabilityEvidenceRecord searches: dict[str, list[SearchResult]] @@ -55,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): @@ -129,6 +135,7 @@ def _prepare_agent( capability_model: str | Model, 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, @@ -142,11 +149,14 @@ def _prepare_agent( 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, ) return capability, deps, agent @@ -203,13 +213,16 @@ async def run_capability_conversation( capability_model: str | Model, document_filter: str | None = None, request_limit: int | 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), so prior-turn compaction operates on - real evidence. Per-invocation state (citations, searches) is cleared by the - capability on every run, so each returned result reflects only its turn. + (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, @@ -218,6 +231,7 @@ async def run_capability_conversation( capability_model, document_filter, request_limit, + compaction=compaction, ) history: list[ModelMessage] | None = None results: list[CapabilityRunResult] = [] @@ -263,6 +277,13 @@ def _result_from_run( executions = getattr(typed, "executions", None) n_executions = len(executions) if executions is not None else 0 + record = typed.evidence + status = ( + citation_status([record], question=record.question) + if record.question is not None + else None + ) + return CapabilityRunResult( answer=answer, cited_uris=cited_uris, @@ -278,4 +299,5 @@ def _result_from_run( n_rejected_searches=traffic.n_rejected_searches, n_failed_tools=traffic.n_failed_tools, n_requests=traffic.n_requests, + citation_status=status, ) diff --git a/evaluations/evaluations/config.py b/evaluations/evaluations/config.py index d8b0ebe0..23faef87 100644 --- a/evaluations/evaluations/config.py +++ b/evaluations/evaluations/config.py @@ -82,6 +82,7 @@ class DatasetSpec: ingest_batch_size: int | None = None evaluate_refusal: bool = False live: bool = False + compaction: bool = False experiment_metadata: dict[str, Any] | None = None def db_path(self, override_path: Path | None = None) -> Path: diff --git a/evaluations/evaluations/datasets/__init__.py b/evaluations/evaluations/datasets/__init__.py index e3532304..45c29adb 100644 --- a/evaluations/evaluations/datasets/__init__.py +++ b/evaluations/evaluations/datasets/__init__.py @@ -3,6 +3,7 @@ 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, ) @@ -20,6 +21,7 @@ DATASETS: dict[str, DatasetSpec] = { 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, diff --git a/evaluations/evaluations/datasets/mtrag.py b/evaluations/evaluations/datasets/mtrag.py index 205bef4b..420913fd 100644 --- a/evaluations/evaluations/datasets/mtrag.py +++ b/evaluations/evaluations/datasets/mtrag.py @@ -285,14 +285,23 @@ def _mtrag_spec(key: str, variant: str) -> DatasetSpec: MTRAG_CLAPNQ_SPEC = _mtrag_spec("mtrag_clapnq", "lastturn") MTRAG_CLAPNQ_REWRITE_SPEC = _mtrag_spec("mtrag_clapnq_rewrite", "rewrite") -MTRAG_CLAPNQ_LIVE_SPEC = DatasetSpec( - key="mtrag_clapnq_live", - db_filename="mtrag_clapnq.lancedb", - document_loader=load_clapnq_corpus, - document_mapper=map_mtrag_document, - qa_loader=load_clapnq_conversations, - qa_case_builder=build_mtrag_live_case, - ingest_batch_size=512, - live=True, - experiment_metadata={"mtrag_mode": "live_session"}, + +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 ) diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 1edad597..a9aa6454 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -188,6 +188,52 @@ class TestConversationInputDispatch: 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: @@ -371,6 +417,7 @@ class TestLiveConversationDispatch: 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 = [ @@ -401,6 +448,7 @@ class TestLiveConversationDispatch: 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 @pytest.mark.asyncio async def test_live_records_per_turn_traffic_arrays(self, tmp_path: Path) -> None: @@ -442,6 +490,7 @@ class TestLiveConversationDispatch: n_rejected_searches=1, n_failed_tools=1, n_requests=4, + citation_status="grounded", ), CapabilityRunResult(answer="a2"), ] @@ -477,6 +526,7 @@ class TestLiveConversationDispatch: 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_"): diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index 058ac734..f6e0fd67 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/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): @@ -161,6 +166,62 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp 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 ( @@ -258,6 +319,76 @@ async def test_conversation_threads_own_messages_across_turns(tmp_path): assert histories == [None, ["history after q1"], ["history after q2"]] +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 diff --git a/evaluations/tests/test_mtrag.py b/evaluations/tests/test_mtrag.py index 94a9cde4..0304a2fa 100644 --- a/evaluations/tests/test_mtrag.py +++ b/evaluations/tests/test_mtrag.py @@ -263,6 +263,19 @@ class TestLiveConversations: 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" + "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, + } From f1d5918d43fba2fbd3039ae70ec3174ce1f9bd97 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 14 Aug 2026 15:11:41 +0300 Subject: [PATCH 4/9] Update MTRAG benchmark numbers to the glimmer baseline --- docs/benchmarks.md | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index dd73d320..7d603b89 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -248,7 +248,8 @@ Directly comparable with [IBM's published results](https://github.com/IBM/mt-rag | Mode | Capability model | Turns | QA accuracy | Mean `cited_map` | |------|------------------|------:|-------------|------------------| -| Gold-prefix (`mtrag_clapnq`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 223 | 0.68 | 0.35 | -| Live (`mtrag_clapnq_live`) | `vllm:Gemma-4-26B-A4B-NVFP4` | 195/224 scored | 0.72 micro / 0.73 macro | 0.35 | +| 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.67.1 with `qwen3-embedding:4b` (vLLM, dim 2560) and `Qwen3-Reranker-4B`, stock capability instructions, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` at temperature 0. Two changes since re-baseline these numbers: the judge sampling was re-pinned repo-wide (0.6 with thinking), and 0.74.0 rewrote the citation instructions (refusals now declare an empty citation list instead of being exempt). Retrieval is unaffected by both and remains the control. QA numbers are internal (our judge and rubric) and not comparable with IBM's published generation metrics. Live mode additionally reports refusal precision/recall against the per-turn answerability labels and per-turn pass rates; pass rate declines with conversation depth (93% at turn 1 to 38% at turn 9). The dataset is text-only (ClapNQ passages), so it exercises no multimodal paths.* +*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). The two live arms differ only in registering `EvidenceCompactionCapability`. Compaction reduces input tokens per request 1.8x (7.5k vs 13.5k), leaves citation retrieval unchanged, and improves answer pass rate (paired McNemar p = 0.031, 95% CI +0.8 to +8.2pp). 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.* From bf94efc6551300220adf471cd85626cab938b914 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 14 Aug 2026 15:03:41 +0300 Subject: [PATCH 5/9] Simplify the eval harness and share the embed-fill path. --- evaluations/evaluations/benchmark.py | 267 ++++++++++-------- evaluations/evaluations/capability_runner.py | 6 +- evaluations/evaluations/config.py | 2 - evaluations/evaluations/datasets/mtrag.py | 21 +- .../evaluations/evaluators/__init__.py | 7 +- .../evaluations/evaluators/conversation.py | 6 +- evaluations/evaluations/evaluators/map.py | 20 +- evaluations/evaluations/evaluators/refusal.py | 5 +- evaluations/tests/test_benchmark.py | 9 +- evaluations/tests/test_mtrag.py | 5 +- haiku_rag_slim/haiku/rag/client/documents.py | 25 +- haiku_rag_slim/haiku/rag/client/processing.py | 14 +- 12 files changed, 194 insertions(+), 193 deletions(-) diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 9b18f962..0d150284 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -1,8 +1,8 @@ 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 @@ -17,6 +17,7 @@ from evaluations.config import ConversationInput, DatasetSpec from evaluations.datasets import DATASETS from evaluations.evaluators import ( ANSWER_EQUIVALENCE_RUBRIC, + REFUSAL_ELIGIBLE_LABELS, REFUSAL_RUBRIC, ConversationEvaluator, RefusalJudge, @@ -394,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)))) @@ -424,7 +427,7 @@ def _resolve_capability_config( return capability_model or config.qa.model -def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | None: +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 @@ -447,7 +450,7 @@ def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | for failure in report_failures ) turns_total = sum(_score(case, "turns_total") for case in scored) - turns_judged = sum(_score(case, "turns_judged") or 0 for case in scored) + turns_judged = sum(_score(case, "turns_judged") for case in scored) turns_passed = sum(_score(case, "turns_passed") for case in scored) summary: dict[str, float | int] = { "conversations": len(scored), @@ -475,9 +478,9 @@ def _live_summary(report_cases, report_failures=()) -> dict[str, float | int] | _score(case, "cited_map") for case in cited ) / len(cited) - true_refusals = sum(_score(case, "true_refusals") or 0 for case in scored) - false_refusals = sum(_score(case, "false_refusals") or 0 for case in scored) - unanswerable = sum(_score(case, "unanswerable_turns") or 0 for case in scored) + 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 @@ -497,7 +500,7 @@ def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None: for case in report_cases: refused = case.assertions.get("refused") label = (case.metadata or {}).get("answerability") - if refused is None or label not in ("ANSWERABLE", "UNANSWERABLE"): + if refused is None or label not in REFUSAL_ELIGIBLE_LABELS: continue outcomes.append((label, bool(refused.value))) if not outcomes: @@ -520,18 +523,29 @@ def _filter_qa_corpus(corpus, case_ids: set[str] | None): return corpus.filter(lambda row: row.get("id") in case_ids) -async def run_qa_benchmark( +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 = None, - name: str | None = None, - db_path: Path | None = None, - judge_model: ModelConfig | None = None, - target: Target = "rag-capability", - capability_model: ModelConfig | None = None, - case_ids: set[str] | None = None, - document_filter: str | None = None, -) -> ReportCaseFailure[str, str, dict[str, str]] | None: + 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: @@ -544,7 +558,74 @@ async def run_qa_benchmark( judge_config = judge_model or DEFAULT_JUDGE_MODEL capability_config = _resolve_capability_config(target, config, capability_model) - db = spec.db_path(db_path) + + eval_name = name if name is not None else f"{spec.key}_qa_evaluation" + experiment_metadata = build_experiment_metadata( + dataset_key=spec.key, + test_cases=len(cases), + config=config, + judge_config=judge_config, + target=target, + capability_config=capability_config, + document_filter=document_filter, + ) + experiment_metadata.update(spec.experiment_metadata or {}) + + return _QARun( + cases=cases, + db=spec.db_path(db_path), + judge_config=judge_config, + eval_name=eval_name, + experiment_metadata=experiment_metadata, + capability_factory=_capability_factory_for_target(target), + capability_model=get_model(capability_config, config), + ) + + +def _print_mean_task_time(report_cases, unit: str = "case") -> None: + if not report_cases: + return + mean = sum(case.task_duration for case in report_cases) / len(report_cases) + console.print(f"Avg task time per {unit}: {mean:.2f}s") + + +def _print_failures(failures, show_question: bool = False) -> None: + if not failures: + return + console.print("[red]\nSummary of failures:[/red]") + for failure in failures: + console.print(f"Case: {failure.name}") + if show_question: + console.print(f"Question: {failure.inputs}") + console.print(f"Error: {failure.error_message}") + console.print("") + + +async def run_qa_benchmark( + spec: DatasetSpec, + config: AppConfig, + limit: int | None = None, + name: str | None = None, + db_path: Path | None = None, + judge_model: ModelConfig | None = None, + target: Target = "rag-capability", + capability_model: ModelConfig | None = None, + case_ids: set[str] | None = None, + document_filter: str | None = None, +) -> ReportCaseFailure[str, str, dict[str, str]] | None: + run = _prepare_qa_run( + spec, + config, + limit, + name, + db_path, + judge_model, + target, + capability_model, + case_ids, + document_filter, + ) + cases, judge_config = run.cases, run.judge_config _attach_relevant_uris(cases, spec, limit) citation_evaluator = spec.citation_evaluator @@ -568,43 +649,20 @@ async def run_qa_benchmark( ] if citation_evaluator is not None: evaluators.append(citation_evaluator) - if spec.evaluate_refusal: - evaluators.append( - RefusalJudge( - rubric=REFUSAL_RUBRIC, - model=get_model(judge_config, config), - assertion={"evaluation_name": "refused", "include_reason": False}, - ) + # RefusalJudge scores only cases whose metadata carries an answerability + # label; on unlabeled datasets it returns no score without a judge call. + evaluators.append( + RefusalJudge( + rubric=REFUSAL_RUBRIC, + model=get_model(judge_config, config), + assertion={"evaluation_name": "refused", "include_reason": False}, ) + ) evaluation_dataset = EvalDataset[Any, str, dict[str, Any]]( name=spec.key, cases=cases, evaluators=evaluators ) - 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 {}) - - async def _evaluate(answer_fn: Callable[[Any], 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(inputs: str | ConversationInput) -> str: if isinstance(inputs, ConversationInput): question = inputs.question @@ -613,11 +671,11 @@ async def run_qa_benchmark( 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, ) @@ -633,7 +691,13 @@ async def run_qa_benchmark( 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 @@ -660,11 +724,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}%)") - if report.cases: - mean_task_time = sum(case.task_duration for case in report.cases) / len( - report.cases - ) - console.print(f"Avg task time per case: {mean_task_time:.2f}s") + _print_mean_task_time(report.cases) if citation_evaluator is not None: score_key = citation_evaluator.get_default_evaluation_name() @@ -693,26 +753,16 @@ async def run_qa_benchmark( ) console.print(f"Mean citations per case: {mean_citations:.2f}") - if spec.evaluate_refusal: - metrics = _refusal_metrics(report.cases) - if metrics is not None: - precision, recall, unanswerable, refusals = metrics - console.print( - "\n=== Refusal vs answerability labels ===", style="bold cyan" - ) - console.print(f"Refusal precision: {precision:.2%} | recall: {recall:.2%}") - console.print( - f"UNANSWERABLE turns: {unanswerable} | refusals: {refusals} " - "(PARTIAL excluded)" - ) + if (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)" + ) - 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("") + _print_failures(failures, show_question=True) return failures[0] if failures else None @@ -727,58 +777,44 @@ async def run_live_qa_benchmark( 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. """ - corpus = spec.qa_loader() - corpus = _filter_qa_corpus(corpus, case_ids) - if limit is not None: - corpus = corpus.select(range(min(limit, len(corpus)))) - - cases = [ - spec.qa_case_builder(index, cast(Mapping[str, Any], doc)) - for index, doc in enumerate(corpus, start=1) - ] - - judge_config = judge_model or DEFAULT_JUDGE_MODEL - capability_config = _resolve_capability_config(target, config, capability_model) - db = spec.db_path(db_path) + 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=cases, + cases=run.cases, evaluators=[ ConversationEvaluator( rubric=ANSWER_EQUIVALENCE_RUBRIC, - model=get_model(judge_config, config), + model=get_model(run.judge_config, config), ) ], ) - eval_name = name if name is not None else f"{spec.key}_qa_evaluation" - experiment_metadata = build_experiment_metadata( - dataset_key=spec.key, - test_cases=len(cases), - config=config, - judge_config=judge_config, - target=target, - capability_config=capability_config, - ) - experiment_metadata.update(spec.experiment_metadata or {}) - - capability_factory = _capability_factory_for_target(target) - resolved_capability_model = get_model(capability_config, config) - async def answer_conversation(questions: list[str]) -> list[str]: results = await run_capability_conversation( - capability_factory=capability_factory, - db_path=db, + capability_factory=run.capability_factory, + db_path=run.db, config=config, questions=list(questions), - capability_model=resolved_capability_model, + capability_model=run.capability_model, compaction=spec.compaction, ) set_eval_attribute("turn_cited_uris", [r.cited_uris for r in results]) @@ -793,10 +829,10 @@ async def run_live_qa_benchmark( report = await evaluation_dataset.evaluate( answer_conversation, - name=eval_name, + name=run.eval_name, max_concurrency=1, progress=True, - metadata=experiment_metadata, + metadata=run.experiment_metadata, ) summary = _live_summary(report.cases, report.failures) @@ -849,12 +885,7 @@ async def run_live_qa_benchmark( f"{per_turn:.2f}s per turn" ) - if report.failures: - console.print("[red]\nSummary of failures:[/red]") - for failure in report.failures: - console.print(f"Case: {failure.name}") - console.print(f"Error: {failure.error_message}") - console.print("") + _print_failures(report.failures) async def evaluate_dataset( diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index 89e0e334..2ab4cef6 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -211,8 +211,6 @@ async def run_capability_conversation( config: AppConfig, questions: list[str], capability_model: str | Model, - document_filter: str | None = None, - request_limit: int | None = None, compaction: bool = False, ) -> list[CapabilityRunResult]: """Run a conversation's user turns sequentially through one capability. @@ -229,8 +227,8 @@ async def run_capability_conversation( db_path, config, capability_model, - document_filter, - request_limit, + document_filter=None, + request_limit=None, compaction=compaction, ) history: list[ModelMessage] | None = None diff --git a/evaluations/evaluations/config.py b/evaluations/evaluations/config.py index 23faef87..205905d2 100644 --- a/evaluations/evaluations/config.py +++ b/evaluations/evaluations/config.py @@ -60,7 +60,6 @@ DocumentLoader = Callable[[], Dataset] DocumentMapper = Callable[[Mapping[str, Any]], DocumentPayload | None] RetrievalLoader = Callable[[], Dataset] RetrievalMapper = Callable[[Mapping[str, Any]], RetrievalSample | None] -QAInput = str | ConversationInput CaseBuilder = Callable[[int, Mapping[str, Any]], Case[Any, Any, dict[str, Any]]] @@ -80,7 +79,6 @@ class DatasetSpec: document_limit: int | None = None retrieval_limit: int = 5 ingest_batch_size: int | None = None - evaluate_refusal: bool = False live: bool = False compaction: bool = False experiment_metadata: dict[str, Any] | None = None diff --git a/evaluations/evaluations/datasets/mtrag.py b/evaluations/evaluations/datasets/mtrag.py index 420913fd..61cc29b0 100644 --- a/evaluations/evaluations/datasets/mtrag.py +++ b/evaluations/evaluations/datasets/mtrag.py @@ -165,7 +165,7 @@ def _task_to_record( } -def load_clapnq_qa() -> Dataset: +def _qa_records() -> list[dict[str, Any]]: path = _download(_GEN_TASKS_FILE) qrels = _load_qrels() records = [] @@ -175,7 +175,11 @@ def load_clapnq_qa() -> Dataset: record = _task_to_record(json.loads(line), qrels) if record is not None: records.append(record) - return Dataset.from_list(records) + return records + + +def load_clapnq_qa() -> Dataset: + return Dataset.from_list(_qa_records()) def build_mtrag_case( @@ -225,17 +229,15 @@ def _group_conversations(records: list[dict[str, Any]]) -> list[dict[str, Any]]: "answerability": task["answerability"], "multi_turn_type": task["multi_turn_type"], "question_type": list(task["question_type"]), + "relevant_uris": list(task["relevant_uris"] or []), } - if task["relevant_uris"]: - turn["relevant_uris"] = list(task["relevant_uris"]) turns.append(turn) conversations.append({"id": conversation_id, "turns": turns}) return conversations def load_clapnq_conversations() -> Dataset: - corpus = load_clapnq_qa() - return Dataset.from_list(_group_conversations([dict(row) for row in corpus])) + return Dataset.from_list(_group_conversations(_qa_records())) def build_mtrag_live_case( @@ -243,11 +245,7 @@ def build_mtrag_live_case( ) -> Case[list[str], list[str], dict[str, Any]]: questions = [turn["question"] for turn in doc["turns"]] metadata_turns = [ - { - key: value - for key, value in turn.items() - if key != "question" and value is not None - } + {key: value for key, value in turn.items() if key != "question"} for turn in doc["turns"] ] return Case( @@ -277,7 +275,6 @@ def _mtrag_spec(key: str, variant: str) -> DatasetSpec: citation_evaluator=CitationMAPEvaluator(), retrieval_limit=10, ingest_batch_size=512, - evaluate_refusal=True, experiment_metadata={"mtrag_mode": "gold_prefix"}, ) diff --git a/evaluations/evaluations/evaluators/__init__.py b/evaluations/evaluations/evaluators/__init__.py index fd3fe169..97fb771b 100644 --- a/evaluations/evaluations/evaluators/__init__.py +++ b/evaluations/evaluations/evaluators/__init__.py @@ -7,12 +7,17 @@ from evaluations.evaluators.judge import ( ) from evaluations.evaluators.map import MAPEvaluator from evaluations.evaluators.number_match import NumberMatchEvaluator -from evaluations.evaluators.refusal import REFUSAL_RUBRIC, RefusalJudge +from evaluations.evaluators.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", diff --git a/evaluations/evaluations/evaluators/conversation.py b/evaluations/evaluations/evaluators/conversation.py index 8b97c60f..78ca74e4 100644 --- a/evaluations/evaluations/evaluators/conversation.py +++ b/evaluations/evaluations/evaluators/conversation.py @@ -9,9 +9,7 @@ from pydantic_evals.evaluators.llm_as_a_judge import ( ) from evaluations.evaluators.citation import average_precision -from evaluations.evaluators.refusal import REFUSAL_RUBRIC - -_REFUSAL_LABELS = ("ANSWERABLE", "UNANSWERABLE") +from evaluations.evaluators.refusal import REFUSAL_ELIGIBLE_LABELS, REFUSAL_RUBRIC @dataclass @@ -81,7 +79,7 @@ class ConversationEvaluator(Evaluator): ) label = turn.get("answerability") - if label in _REFUSAL_LABELS: + if label in REFUSAL_ELIGIBLE_LABELS: try: refused = ( await judge_output(answer, REFUSAL_RUBRIC, self.model) diff --git a/evaluations/evaluations/evaluators/map.py b/evaluations/evaluations/evaluators/map.py index 1746f6ad..e4fe5a32 100644 --- a/evaluations/evaluations/evaluators/map.py +++ b/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) diff --git a/evaluations/evaluations/evaluators/refusal.py b/evaluations/evaluations/evaluators/refusal.py index ede0b4a5..435720fc 100644 --- a/evaluations/evaluations/evaluators/refusal.py +++ b/evaluations/evaluations/evaluators/refusal.py @@ -1,4 +1,5 @@ from dataclasses import dataclass +from typing import Final from pydantic_evals.evaluators import EvaluatorContext, LLMJudge from pydantic_evals.evaluators.evaluator import EvaluatorOutput @@ -9,7 +10,7 @@ REFUSAL_RUBRIC = ( "substantive answer." ) -_ELIGIBLE_LABELS = ("ANSWERABLE", "UNANSWERABLE") +REFUSAL_ELIGIBLE_LABELS: Final = ("ANSWERABLE", "UNANSWERABLE") @dataclass @@ -23,6 +24,6 @@ class RefusalJudge(LLMJudge): async def evaluate(self, ctx: EvaluatorContext) -> EvaluatorOutput: label = (ctx.metadata or {}).get("answerability") - if label not in _ELIGIBLE_LABELS: + if label not in REFUSAL_ELIGIBLE_LABELS: return {} return await super().evaluate(ctx) diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index a9aa6454..5345f571 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -332,7 +332,7 @@ class TestLiveSummary: def test_none_without_scored_cases(self) -> None: from evaluations.benchmark import _live_summary - assert _live_summary([self._case({})]) is None + assert _live_summary([self._case({})], []) is None def test_micro_rate_uses_judged_turns(self) -> None: from evaluations.benchmark import _live_summary @@ -352,7 +352,7 @@ class TestLiveSummary: ) ] - summary = _live_summary(cases) + summary = _live_summary(cases, []) assert summary is not None assert summary["micro_pass_rate"] == 1.0 @@ -784,8 +784,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 diff --git a/evaluations/tests/test_mtrag.py b/evaluations/tests/test_mtrag.py index 0304a2fa..90a7cc7b 100644 --- a/evaluations/tests/test_mtrag.py +++ b/evaluations/tests/test_mtrag.py @@ -139,9 +139,6 @@ class TestSpecs: } assert isinstance(spec.citation_evaluator, CitationMAPEvaluator) - def test_refusal_evaluation_enabled(self) -> None: - assert MTRAG_CLAPNQ_SPEC.evaluate_refusal is True - class TestGenerationTasks: def test_task_to_record(self) -> None: @@ -255,7 +252,7 @@ class TestLiveConversations: } other_case = build_mtrag_live_case(2, conversations[1]) assert other_case.metadata is not None - assert "relevant_uris" not in other_case.metadata["turns"][0] + assert other_case.metadata["turns"][0]["relevant_uris"] == [] def test_live_spec(self) -> None: assert DATASETS["mtrag_clapnq_live"] is MTRAG_CLAPNQ_LIVE_SPEC diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 1f5c4aae..7257fcbd 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -299,21 +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. """ - missing = [ - chunk - for _, chunks, _ in prepared - for chunk in chunks - if chunk.embedding is None - ] - if missing: - from haiku.rag.embeddings import embed_chunks - - embedded_flat = await embed_chunks(missing, client.embedder, client._config) - # Assign positionally: duplicate chunk texts across documents make a - # content-keyed lookup ambiguous. - for chunk, with_embedding in zip(missing, embedded_flat): - chunk.embedding = with_embedding.embedding - embedded: list[list[Chunk]] = [chunks 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] diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 87f59ff3..4d541658 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/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: From 0634964e64e9ab2a05282eae9510d5b914a648cb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 14 Aug 2026 15:26:07 +0300 Subject: [PATCH 6/9] Point the mtrag reference config at the measured baseline model --- evaluations/configs/mtrag_clapnq.yaml | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/evaluations/configs/mtrag_clapnq.yaml b/evaluations/configs/mtrag_clapnq.yaml index 654b1eb6..7fa34d84 100644 --- a/evaluations/configs/mtrag_clapnq.yaml +++ b/evaluations/configs/mtrag_clapnq.yaml @@ -1,6 +1,7 @@ # Reference config for the `mtrag_clapnq` pre-built evaluation database. # IBM MTRAG, ClapNQ (Wikipedia) domain: multi-turn retrieval and QA over -# 183,408 passages. Also serves mtrag_clapnq_rewrite and mtrag_clapnq_live. +# 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 @@ -27,11 +28,18 @@ reranking: qa: model: provider: openai - name: gemma4-26b - base_url: http://vllm:11432/v1 + 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: From 3b9d6ae2c6dd4d3abe1b00edf578fdd1d548d41b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 14 Aug 2026 17:05:45 +0300 Subject: [PATCH 7/9] Publish the paired statistics behind the MTRAG compaction claims --- docs/benchmarks.md | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 7d603b89..68340bb0 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -252,4 +252,11 @@ Directly comparable with [IBM's published results](https://github.com/IBM/mt-rag | 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). The two live arms differ only in registering `EvidenceCompactionCapability`. Compaction reduces input tokens per request 1.8x (7.5k vs 13.5k), leaves citation retrieval unchanged, and improves answer pass rate (paired McNemar p = 0.031, 95% CI +0.8 to +8.2pp). 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.* +*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). From 587ba75a623a77ded961fe5f779086efa560d1a9 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 17 Aug 2026 11:03:52 +0300 Subject: [PATCH 8/9] Thread the document filter through the live QA runner --- evaluations/evaluations/benchmark.py | 1 + evaluations/evaluations/capability_runner.py | 3 ++- evaluations/tests/test_benchmark.py | 11 +++++++-- evaluations/tests/test_capability_runner.py | 26 ++++++++++++++++++++ 4 files changed, 38 insertions(+), 3 deletions(-) diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 0d150284..0f0796e6 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -815,6 +815,7 @@ async def run_live_qa_benchmark( 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]) diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index 2ab4cef6..613ae439 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -211,6 +211,7 @@ async def run_capability_conversation( 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. @@ -227,7 +228,7 @@ async def run_capability_conversation( db_path, config, capability_model, - document_filter=None, + document_filter=document_filter, request_limit=None, compaction=compaction, ) diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 5345f571..5961a51e 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -443,12 +443,19 @@ class TestLiveConversationDispatch: ), ): await run_live_qa_benchmark( - spec, AppConfig(), db_path=tmp_path / "test.lancedb" + 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: @@ -1143,7 +1150,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: diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index f6e0fd67..d7adeba5 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/evaluations/tests/test_capability_runner.py @@ -319,6 +319,32 @@ async def test_conversation_threads_own_messages_across_turns(tmp_path): 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' From 73dfc62f33e7f94cba187162e45c903bd9efc0de Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 17 Aug 2026 11:19:59 +0300 Subject: [PATCH 9/9] Exclude fully unjudged conversations from the macro pass rate --- evaluations/evaluations/benchmark.py | 10 +++++-- evaluations/tests/test_benchmark.py | 41 ++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 0f0796e6..cb9f534c 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -452,6 +452,10 @@ def _live_summary(report_cases, report_failures) -> dict[str, float | int] | Non 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), @@ -459,8 +463,10 @@ def _live_summary(report_cases, report_failures) -> dict[str, float | int] | Non "turns_judged": turns_judged, "turns_attempted": turns_total + failed_turns, "micro_pass_rate": turns_passed / turns_judged if turns_judged else 0.0, - "macro_pass_rate": sum(_score(case, "turn_pass_rate") for case in scored) - / len(scored), + "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] diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 5961a51e..16b43013 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -359,6 +359,47 @@ class TestLiveSummary: 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