Merge pull request #389 from ggozad/chore/clean-up-evaluations

Clean up and update evaluations.
This commit is contained in:
Yiorgis Gozadinos 2026-06-01 11:07:32 +03:00 committed by GitHub
commit d5e5733f67
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 362 additions and 796 deletions

View file

@ -36,13 +36,6 @@ Active datasets:
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
Inactive (kept downloadable, not currently maintained):
| Dataset | Size |
|---------|------|
| `repliqa` | ~30MB |
| `hotpotqa` | ~331MB |
After downloading, run benchmarks with `--skip-db` to use the pre-built database:
```bash
@ -85,20 +78,13 @@ evaluations:
### Retrieval Metrics
**Mean Reciprocal Rank (MRR)** - Used when each query has exactly one relevant document.
- For each query, find the rank (position) of the first relevant document in top-K results
- Reciprocal rank = `1/rank` (e.g., rank 3 → 1/3 ≈ 0.333)
- If not found in top-K, score is 0
- MRR is the mean across all queries
- Range: 0 (never found) to 1 (always at rank 1)
**Mean Average Precision (MAP)** - Used when queries have multiple relevant documents.
**Mean Average Precision (MAP)** scores ranked retrieval results against the gold `expected_uris`.
- For each relevant document at position k, calculate precision@k = (relevant docs in top k) / k
- Average Precision (AP) = mean of these precision values / total relevant documents
- Average Precision (AP) = sum of these precision values / total relevant documents
- MAP is the mean of AP scores across all queries
- Range: 0 to 1. Rewards ranking relevant documents higher
- For single-doc queries this collapses to `1/rank` (i.e. reciprocal rank)
### QA Accuracy
@ -108,7 +94,7 @@ We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibrat
### Citation Retrieval
Alongside QA accuracy, a second metric scores the URIs the skill registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MRR / MAP math as raw retrieval. The score key is `cited_mrr` for single-doc datasets and `cited_map` for multi-doc. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
Alongside QA accuracy, a second metric scores the URIs the skill registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
This is computed alongside QA accuracy from the same skill run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the skill grounded its answer on it.
@ -116,18 +102,6 @@ This is computed alongside QA accuracy from the same skill run, no extra invocat
Numbers measured under the current pinned judge (`ollama:qwen3.6`) on a recent `haiku.rag` version.
### Wix
[WixQA](https://huggingface.co/datasets/Wix/WixQA) is real customer support questions paired with curated answers. 200 cases.
`evaluations run wix --target rag-skill` runs the RAG skill end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the skill registered via the `cite` tool against the gold `expected_uris`.
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
|------------------------------|------------------------|-------------|------------------|
| `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | 0.87 | 0.38 |
*Measured on haiku.rag v0.48.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `chunk_size=256`, `search.limit=5`. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4` (qwen3.6 family, NVFP4 quant served via vLLM rather than the default Ollama). 172 / 198 completed cases (2 errored).*
### OpenRAG Bench (ORB)
[OpenRAG Bench](https://huggingface.co/datasets/vectara/open_ragbench) contains ArXiv research papers with multimodal question-answering pairs. Queries include both text-based and image-based questions, testing retrieval and reasoning over visual content like figures, charts, and diagrams. Each query maps to one relevant document.
@ -141,114 +115,49 @@ Two approaches are benchmarked separately:
##### Retrieval (MAP)
| Embedding Model | Source bucket | Cases | MAP |
|------------------------------|--------------------|------:|-------:|
| `Qwen/Qwen3-VL-Embedding-8B` | text only | 1914 | 0.9801 |
| `Qwen/Qwen3-VL-Embedding-8B` | text + image | 763 | 0.9720 |
| `Qwen/Qwen3-VL-Embedding-8B` | text + table | 148 | 0.9786 |
| `Qwen/Qwen3-VL-Embedding-8B` | text + table+image | 220 | 0.9720 |
| `Qwen/Qwen3-VL-Embedding-8B` | **all** | 3045 | **0.9774** |
| Embedding Model | Cases | MAP |
|------------------------------------------|------:|-------:|
| `Qwen/Qwen3-VL-Embedding-8B` | 3045 | 0.9774 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | 3045 | 0.9709 |
##### QA Accuracy
##### QA accuracy + citation retrieval
| Embedding Model | Skill model | Reranker | Source bucket | Cases | Accuracy |
|------------------------------|-----------------------------------|------------------------|---------------|------:|---------:|
| `Qwen/Qwen3-VL-Embedding-8B` | `ollama:qwen3.6` (vision) | none | text only | 682 | 96.9 % |
| `Qwen/Qwen3-VL-Embedding-8B` | `ollama:qwen3.6` (vision) | none | with image | 299 | 91.3 % |
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text | 894 | 88.5 % |
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text+image | 341 | 88.0 % |
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text+table | 72 | 88.9 % |
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | text+table+image | 102 | 93.1 % |
| `Qwen/Qwen3-VL-Embedding-8B` | `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | **all** | 1409 | **88.7 %** |
| Embedding Model | Target | Skill model | Cases | QA accuracy | Mean `cited_map` |
|------------------------------------------|-----------------|------------------------------|------:|-------------|------------------|
| `Qwen/Qwen3-VL-Embedding-8B` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-skill` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-skill`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.35 |
*Measured on haiku.rag v0.50.0, no reranker (no multimodal reranker available), judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases and predates the `cited_map` evaluator. The analysis skill answers via Python sandbox over the document VFS and only cites in ~35 % of cases, which depresses `cited_map`.*
#### Text embedder + VLM picture descriptions
##### Retrieval (MAP)
| Embedding Model | VLM | Source bucket | Cases | MAP |
|----------------------|----------------------|---------------|------:|-----------:|
| `qwen3-embedding:4b` | Ollama / ministral-3 | **all** | 3045 | **0.9722** |
| Embedding Model | VLM | Reranker | Cases | MAP |
|------------------------------------------|----------------------|------------------------|------:|-------:|
| `qwen3-embedding:4b` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9834 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `mxbai-rerank-base-v2` | 3045 | 0.9863 |
*Measured on haiku.rag v0.45.0.*
*Measured on haiku.rag v0.50.0.*
##### QA accuracy + citation retrieval
| Embedding Model | VLM | Skill model | QA accuracy | Mean `cited_map` |
|------------------------|----------------------|------------------------------|-------------|------------------|
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 0.88 | 0.89 |
| Embedding Model | VLM | Skill model | Cases | QA accuracy | Mean `cited_map` |
|------------------------------------------|----------------------|------------------------------|------:|-------------|------------------|
| `qwen3-embedding:4b` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.80 |
| `nvidia/llama-nemotron-embed-vl-1b-v2` | Ollama / ministral-3 | `vllm:Gemma-4-26B-A4B-NVFP4` | 2836 | 0.96 | 0.81 |
*Measured on haiku.rag v0.48.0, with `mxbai-rerank-base-v2`, on all 3045 cases. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4`.*
*Measured on haiku.rag v0.50.0 with `mxbai-rerank-base-v2`, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Nemotron ran 2836 / 3045 cases (run stopped early; numbers stable from ~12% onward).*
## Inactive datasets
### Wix
The benchmarks below are not currently maintained. Numbers were measured against earlier `haiku.rag` versions and an older pinned judge (`ollama:gpt-oss`), before the skill workflow became the only path. Retrieval tables don't depend on the judge, but the QA tables aren't reproducible against the current skill-only setup. We may revive them.
[WixQA](https://huggingface.co/datasets/Wix/WixQA) is real customer support questions paired with curated answers. 200 cases.
### RepliQA
`evaluations run wix --target rag-skill` runs the RAG skill end-to-end and produces both QA accuracy and a citation retrieval metric (`cited_map`) computed from the URIs the skill registered via the `cite` tool against the gold `expected_uris`.
[RepliQA](https://huggingface.co/datasets/ServiceNow/repliqa) contains synthetic news stories with question-answer pairs. We use `News Stories` from `repliqa_3` (1035 documents). Each question has exactly one relevant document, so we use MRR for retrieval evaluation.
| Skill model | Reranker | QA accuracy | Mean `cited_map` |
|------------------------------|------------------------|-------------|------------------|
| `vllm:Gemma-4-26B-A4B-NVFP4` | `mxbai-rerank-base-v2` | 0.87 | 0.38 |
#### Retrieval (MRR)
| Embedding Model | MRR | Reranker |
|-------------------------------|------|----------|
| Ollama / `qwen3-embedding:8b` | 0.91 | - |
*Measured on haiku.rag v0.19.6.*
#### QA Accuracy
| Embedding Model | QA Model | Accuracy | Reranker |
|------------------------------|----------------------------------|----------|------------------------|
| Ollama / `qwen3-embedding:4b` | Ollama / `gpt-oss` - no thinking | 0.82 | None |
| Ollama / `qwen3-embedding:8b` | Ollama / `gpt-oss` - thinking | 0.89 | None |
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` - thinking | 0.85 | None |
| Ollama / `mxbai-embed-large` | Ollama / `qwen3` - thinking | 0.87 | `mxbai-rerank-base-v2` |
| Ollama / `mxbai-embed-large` | Ollama / `qwen3:0.6b` | 0.28 | None |
*Measured on haiku.rag v0.19.6, judged by `ollama:gpt-oss`.*
Note the significant degradation when very small models are used such as `qwen3:0.6b`.
### HotpotQA
[HotpotQA](https://huggingface.co/datasets/hotpotqa/hotpot_qa) is a multi-hop question answering dataset requiring reasoning over multiple Wikipedia paragraphs. Each question requires evidence from 2+ documents, making it ideal for testing retrieval and reasoning capabilities. We use MAP for retrieval evaluation since queries have multiple relevant documents.
#### Retrieval (MAP)
| Embedding Model | MAP | Reranker |
|----------------------|------|----------|
| `qwen3-embedding:4b` | 0.69 | none |
*Measured on haiku.rag v0.20.2.*
#### QA Accuracy
| Embedding Model | QA Model | Accuracy |
|----------------------|--------------------------|----------|
| `qwen3-embedding:4b` | `gpt-oss:20b` - thinking | 0.86 |
*Measured on haiku.rag v0.20.2, judged by `ollama:gpt-oss`.*
### Wix (historical, plain text and HTML)
Earlier Wix runs measured under different chunk settings and reranker combinations, against the older `gpt-oss` judge.
#### Retrieval (MAP)
| Embedding Model | Chunk size | MAP | Reranker | Notes |
|------------------------|------------|------|------------------------|------------------------------|
| `qwen3-embedding:4b` | 256 | 0.34 | None | html, `chunk-radius=2` |
| `qwen3-embedding:4b` | 256 | 0.39 | `mxbai-rerank-base-v2` | html, `chunk-radius=2` |
| `qwen3-embedding:4b` | 256 | 0.43 | None | plain text, `chunk-radius=0` |
| `qwen3-embedding:4b` | 512 | 0.45 | None | plain text, `chunk-radius=0` |
*Measured on haiku.rag v0.27.2.*
#### QA Accuracy
| Embedding Model | Chunk size | QA Model | Accuracy | Notes |
|----------------------|------------|-----------------------------|----------|------------------------------|
| `qwen3-embedding:4b` | 256 | `gpt-oss:20b` - no thinking | 0.80 | html, `chunk-radius=2` |
| `qwen3-embedding:4b` | 256 | `gpt-oss:20b` - no thinking | 0.83 | html, `chunk-radius=2`, `jinaai/jina-reranker-v3` |
*Measured on haiku.rag v0.27.2, judged by `ollama:gpt-oss`.*
*Measured on haiku.rag v0.48.0 with `qwen3-embedding:4b` (vLLM, dim 2560), `chunk_size=256`, `search.limit=5`. Judged by `vllm:Qwen3.6-35B-A3B-NVFP4` (qwen3.6 family, NVFP4 quant served via vLLM rather than the default Ollama). 172 / 198 completed cases (2 errored).*

View file

@ -8,9 +8,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:
- RepliQA (`repliqa`)
- WiX (`wix`)
- HotpotQA (`hotpotqa`)
- 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.
@ -21,24 +19,24 @@ After installing the package, you can run evaluations using the `evaluations` co
```bash
# Run retrieval + QA benchmarks
evaluations run repliqa
evaluations run wix
evaluations run orb_text
# Use a custom config file
evaluations run repliqa --config /path/to/haiku.rag.yaml
evaluations run wix --config /path/to/haiku.rag.yaml
# Override the database path
evaluations run repliqa --db /path/to/custom.lancedb
evaluations run wix --db /path/to/custom.lancedb
# Skip database population and run only benchmarks
evaluations run repliqa --skip-db
evaluations run wix --skip-db
# Skip specific benchmarks
evaluations run repliqa --skip-retrieval
evaluations run repliqa --skip-qa
evaluations run wix --skip-retrieval
evaluations run wix --skip-qa
# Limit the number of test cases
evaluations run repliqa --limit 100
evaluations run wix --limit 100
```
### Choosing the target
@ -54,24 +52,23 @@ evaluations run wix --target analysis-skill --skill-model ollama:gpt-oss
`--skill-model "provider:name"` overrides the skill model independently from
the judge (defaults to `qa.model`, or `analysis.model` when set for the
analysis-skill target). A citation retrieval metric (`cited_mrr` / `cited_map`)
is computed alongside QA accuracy from the URIs the skill registered via the
`cite` tool.
analysis-skill target). A citation retrieval metric (`cited_map`) is computed
alongside QA accuracy from the URIs the skill registered via the `cite` tool.
### Pre-built Databases
Download pre-built evaluation databases from HuggingFace:
```bash
evaluations download repliqa
evaluations download wix
evaluations download all
evaluations download repliqa --force
evaluations download wix --force
```
Upload databases (maintainer only):
```bash
evaluations upload repliqa
evaluations upload wix
evaluations upload all
```

View file

@ -19,9 +19,7 @@ from evaluations.datasets import DATASETS
from evaluations.evaluators import (
ANSWER_EQUIVALENCE_RUBRIC,
CitationMAPEvaluator,
CitationMRREvaluator,
MAPEvaluator,
MRREvaluator,
)
from evaluations.skill_runner import SkillFactory, run_skill_question
from haiku.rag.client import HaikuRAG
@ -30,11 +28,6 @@ from haiku.rag.config.models import ModelConfig
from haiku.rag.logging import configure_cli_logging
from haiku.rag.utils import get_model, parse_model_option
_CITATION_EVALUATORS: dict[type[Evaluator], type[Evaluator]] = {
MRREvaluator: CitationMRREvaluator,
MAPEvaluator: CitationMAPEvaluator,
}
Target = Literal["rag-skill", "analysis-skill"]
TARGETS: tuple[Target, ...] = ("rag-skill", "analysis-skill")
@ -306,10 +299,9 @@ def _skill_factory_for_target(target: Target) -> SkillFactory:
def _citation_evaluator_for(retrieval_evaluator: Evaluator | None) -> Evaluator | None:
"""Return the citation-scoring twin of the dataset's retrieval evaluator."""
if retrieval_evaluator is None:
return None
twin = _CITATION_EVALUATORS.get(type(retrieval_evaluator))
return twin() if twin is not None else None
if isinstance(retrieval_evaluator, MAPEvaluator):
return CitationMAPEvaluator()
return None
def _attach_relevant_uris(

View file

@ -1,16 +1,12 @@
from evaluations.config import DatasetSpec
from .hotpotqa import HOTPOTQA_SPEC
from .open_rag_bench import ORB_MULTIMODAL_SPEC, ORB_TEXT_SPEC
from .repliqa import REPLIQA_SPEC
from .wix import WIX_SPEC
DATASETS: dict[str, DatasetSpec] = {
spec.key: spec
for spec in (
REPLIQA_SPEC,
WIX_SPEC,
HOTPOTQA_SPEC,
ORB_TEXT_SPEC,
ORB_MULTIMODAL_SPEC,
)

View file

@ -1,108 +0,0 @@
from collections.abc import Mapping
from typing import Any, cast
from datasets import Dataset, load_dataset
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MAPEvaluator
def load_hotpotqa_validation() -> Dataset:
dataset_dict = load_dataset("hotpotqa/hotpot_qa", "distractor")
return dataset_dict["validation"]
def extract_unique_documents(dataset: Dataset) -> list[dict[str, Any]]:
"""Extract unique documents from all context paragraphs, deduplicated by title."""
seen_titles: set[str] = set()
documents: list[dict[str, Any]] = []
for sample in dataset:
sample = cast(Mapping[str, Any], sample)
context = sample["context"]
titles = context["title"]
sentences_list = context["sentences"]
for title, sentences in zip(titles, sentences_list):
if title in seen_titles:
continue
seen_titles.add(title)
content = " ".join(sentences)
documents.append({"title": title, "content": content})
return documents
_cached_documents: list[dict[str, Any]] | None = None
def load_hotpotqa_documents() -> list[dict[str, Any]]:
"""Load and cache unique documents from HotpotQA."""
global _cached_documents
if _cached_documents is None:
dataset = load_hotpotqa_validation()
_cached_documents = extract_unique_documents(dataset)
return _cached_documents
def document_loader() -> Dataset:
"""Return documents as a Dataset-like iterable."""
docs = load_hotpotqa_documents()
return Dataset.from_list(docs)
def map_hotpotqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
return DocumentPayload(
uri=doc["title"],
content=doc["content"],
title=doc["title"],
)
def map_hotpotqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
supporting_facts = doc["supporting_facts"]
titles = supporting_facts["title"]
if not titles:
return None
unique_titles = tuple(dict.fromkeys(titles))
return RetrievalSample(
question=doc["question"],
expected_uris=unique_titles,
)
def build_hotpotqa_case(
index: int, doc: Mapping[str, Any]
) -> Case[str, str, dict[str, str]]:
question_id = doc["id"]
question_type = doc["type"]
level = doc["level"]
case_name = f"{index}_{question_id}"
return Case(
name=case_name,
inputs=doc["question"],
expected_output=doc["answer"],
metadata={
"question_id": str(question_id),
"type": str(question_type),
"level": str(level),
"case_index": str(index),
},
)
HOTPOTQA_SPEC = DatasetSpec(
key="hotpotqa",
db_filename="hotpotqa.lancedb",
document_loader=document_loader,
document_mapper=map_hotpotqa_document,
qa_loader=load_hotpotqa_validation,
qa_case_builder=build_hotpotqa_case,
retrieval_loader=load_hotpotqa_validation,
retrieval_mapper=map_hotpotqa_retrieval,
retrieval_evaluator=MAPEvaluator(),
)

View file

@ -1,60 +0,0 @@
from collections.abc import Mapping
from typing import Any
from datasets import Dataset, load_dataset
from pydantic_evals import Case
from evaluations.config import DatasetSpec, DocumentPayload, RetrievalSample
from evaluations.evaluators import MRREvaluator
def load_repliqa_corpus() -> Dataset:
dataset_dict = load_dataset("ServiceNow/repliqa")
dataset = dataset_dict["repliqa_3"]
return dataset.filter(lambda doc: doc["document_topic"] == "News Stories")
def map_repliqa_document(doc: Mapping[str, Any]) -> DocumentPayload:
return DocumentPayload(
uri=str(doc["document_id"]),
content=doc["document_extracted"],
)
def map_repliqa_retrieval(doc: Mapping[str, Any]) -> RetrievalSample | None:
expected_answer = doc["answer"]
if expected_answer == "The answer is not found in the document.":
return None
return RetrievalSample(
question=doc["question"],
expected_uris=(str(doc["document_id"]),),
)
def build_repliqa_case(
index: int, doc: Mapping[str, Any]
) -> Case[str, str, dict[str, str]]:
document_id = doc["document_id"]
case_name = f"{index}_{document_id}" if document_id is not None else f"case_{index}"
return Case(
name=case_name,
inputs=doc["question"],
expected_output=doc["answer"],
metadata={
"document_id": str(document_id),
"case_index": str(index),
},
)
REPLIQA_SPEC = DatasetSpec(
key="repliqa",
db_filename="repliqa.lancedb",
document_loader=load_repliqa_corpus,
document_mapper=map_repliqa_document,
qa_loader=load_repliqa_corpus,
qa_case_builder=build_repliqa_case,
retrieval_loader=load_repliqa_corpus,
retrieval_mapper=map_repliqa_retrieval,
retrieval_evaluator=MRREvaluator(),
)

View file

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

View file

@ -13,35 +13,13 @@ def _relevant_uris(ctx: EvaluatorContext) -> set[str]:
return set(ctx.metadata.get("relevant_uris", []))
@dataclass
class CitationMRREvaluator(Evaluator):
"""Reciprocal rank over the URIs the skill 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``. Returns ``1.0/rank`` of the first cited URI that is in
the relevant set, or ``0.0`` if none match.
Use for single-document datasets, mirroring :class:`MRREvaluator`.
"""
def get_default_evaluation_name(self) -> str:
return "cited_mrr"
def evaluate(self, ctx: EvaluatorContext) -> float:
relevant = _relevant_uris(ctx)
for rank, uri in enumerate(_cited_uris(ctx), start=1):
if uri in relevant:
return 1.0 / rank
return 0.0
@dataclass
class CitationMAPEvaluator(Evaluator):
"""Average precision over the URIs the skill cited via the `cite` tool.
Same input shape as :class:`CitationMRREvaluator`; use for multi-document
datasets, mirroring :class:`MAPEvaluator`.
Reads ``cited_uris`` from ``ctx.attributes`` (recorded during the task run
via :func:`pydantic_evals.set_eval_attribute`) and ``relevant_uris`` from
``ctx.metadata``.
"""
def get_default_evaluation_name(self) -> str:

View file

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

View file

@ -74,19 +74,19 @@ class TestBuildExperimentMetadata:
class TestResolveDataset:
def test_valid_dataset(self) -> None:
spec = _resolve_dataset("repliqa")
assert spec.key == "repliqa"
spec = _resolve_dataset("wix")
assert spec.key == "wix"
def test_case_insensitive(self) -> None:
spec = _resolve_dataset("REPLIQA")
assert spec.key == "repliqa"
spec = _resolve_dataset("WIX")
assert spec.key == "wix"
def test_unknown_dataset_raises(self) -> None:
with pytest.raises(typer.BadParameter, match="Unknown dataset 'nonexistent'"):
_resolve_dataset("nonexistent")
def test_error_lists_valid_datasets(self) -> None:
with pytest.raises(typer.BadParameter, match="repliqa"):
with pytest.raises(typer.BadParameter, match="wix"):
_resolve_dataset("nonexistent")
@ -325,13 +325,6 @@ class TestRunQaBenchmarkSkillTarget:
class TestCitationEvaluatorWiring:
def test_returns_mrr_twin_for_mrr_evaluator(self) -> None:
from evaluations.benchmark import _citation_evaluator_for
from evaluations.evaluators import CitationMRREvaluator, MRREvaluator
result = _citation_evaluator_for(MRREvaluator())
assert isinstance(result, CitationMRREvaluator)
def test_returns_map_twin_for_map_evaluator(self) -> None:
from evaluations.benchmark import _citation_evaluator_for
from evaluations.evaluators import CitationMAPEvaluator, MAPEvaluator
@ -351,7 +344,7 @@ class TestAttachRelevantUris:
from evaluations.benchmark import _attach_relevant_uris
from evaluations.config import RetrievalSample
from evaluations.evaluators import MRREvaluator
from evaluations.evaluators import MAPEvaluator
cases: list[Case[str, str, dict]] = [
Case(name="c1", inputs="What is X?", expected_output="X is a thing"),
@ -382,7 +375,7 @@ class TestAttachRelevantUris:
retrieval_mapper=lambda d: RetrievalSample(
question=d["q"], expected_uris=d["uris"]
),
retrieval_evaluator=MRREvaluator(),
retrieval_evaluator=MAPEvaluator(),
)
_attach_relevant_uris(cases, spec, limit=None)

View file

@ -1,9 +1,6 @@
from unittest.mock import MagicMock
from evaluations.evaluators.citation import (
CitationMAPEvaluator,
CitationMRREvaluator,
)
from evaluations.evaluators.citation import CitationMAPEvaluator
def _ctx(cited: list[str], relevant: list[str]) -> MagicMock:
@ -13,41 +10,6 @@ def _ctx(cited: list[str], relevant: list[str]) -> MagicMock:
return ctx
class TestCitationMRREvaluator:
def setup_method(self) -> None:
self.evaluator = CitationMRREvaluator()
def test_first_citation_is_relevant(self) -> None:
assert self.evaluator.evaluate(_ctx(["a", "b"], ["a"])) == 1.0
def test_second_citation_is_relevant(self) -> None:
assert self.evaluator.evaluate(_ctx(["a", "b"], ["b"])) == 0.5
def test_no_citations(self) -> None:
assert self.evaluator.evaluate(_ctx([], ["a"])) == 0.0
def test_no_relevant(self) -> None:
assert self.evaluator.evaluate(_ctx(["a"], [])) == 0.0
def test_no_matches(self) -> None:
assert self.evaluator.evaluate(_ctx(["a", "b"], ["c"])) == 0.0
def test_metadata_none(self) -> None:
ctx = MagicMock()
ctx.metadata = None
ctx.attributes = {"cited_uris": ["a"]}
assert self.evaluator.evaluate(ctx) == 0.0
def test_attribute_missing(self) -> None:
ctx = MagicMock()
ctx.metadata = {"relevant_uris": ["a"]}
ctx.attributes = {}
assert self.evaluator.evaluate(ctx) == 0.0
def test_evaluation_name(self) -> None:
assert self.evaluator.get_default_evaluation_name() == "cited_mrr"
class TestCitationMAPEvaluator:
def setup_method(self) -> None:
self.evaluator = CitationMAPEvaluator()

View file

@ -1,11 +1,5 @@
from pathlib import Path
from evaluations.datasets.hotpotqa import (
build_hotpotqa_case,
extract_unique_documents,
map_hotpotqa_document,
map_hotpotqa_retrieval,
)
from evaluations.datasets.open_rag_bench import (
build_orb_case,
download_pdf,
@ -13,11 +7,6 @@ from evaluations.datasets.open_rag_bench import (
map_orb_document,
map_orb_retrieval,
)
from evaluations.datasets.repliqa import (
build_repliqa_case,
map_repliqa_document,
map_repliqa_retrieval,
)
from evaluations.datasets.wix import (
build_wix_case,
map_wix_document,
@ -25,50 +14,6 @@ from evaluations.datasets.wix import (
)
class TestRepliqa:
def test_map_document(self) -> None:
doc = {"document_id": "doc-42", "document_extracted": "Some content here."}
payload = map_repliqa_document(doc)
assert payload.uri == "doc-42"
assert payload.content == "Some content here."
def test_map_retrieval(self) -> None:
doc = {
"question": "What happened?",
"answer": "Something happened.",
"document_id": "doc-42",
}
sample = map_repliqa_retrieval(doc)
assert sample is not None
assert sample.question == "What happened?"
assert sample.expected_uris == ("doc-42",)
def test_map_retrieval_skips_unanswerable(self) -> None:
doc = {
"question": "What?",
"answer": "The answer is not found in the document.",
"document_id": "doc-1",
}
assert map_repliqa_retrieval(doc) is None
def test_build_case(self) -> None:
doc = {
"document_id": "doc-7",
"question": "Why?",
"answer": "Because.",
}
case = build_repliqa_case(3, doc)
assert case.name == "3_doc-7"
assert case.inputs == "Why?"
assert case.expected_output == "Because."
assert case.metadata == {"document_id": "doc-7", "case_index": "3"}
def test_build_case_none_document_id(self) -> None:
doc = {"document_id": None, "question": "Q?", "answer": "A."}
case = build_repliqa_case(1, doc)
assert case.name == "case_1"
class TestWix:
def test_map_document_with_all_fields(self) -> None:
doc = {
@ -136,78 +81,6 @@ class TestWix:
assert case.name == "case_1"
class TestHotpotQA:
def test_map_document(self) -> None:
doc = {"title": "Albert Einstein", "content": "Was a physicist."}
payload = map_hotpotqa_document(doc)
assert payload.uri == "Albert Einstein"
assert payload.content == "Was a physicist."
assert payload.title == "Albert Einstein"
def test_map_retrieval(self) -> None:
doc = {
"question": "Who was Einstein?",
"supporting_facts": {"title": ["Albert Einstein", "Physics"]},
}
sample = map_hotpotqa_retrieval(doc)
assert sample is not None
assert sample.expected_uris == ("Albert Einstein", "Physics")
def test_map_retrieval_deduplicates_titles(self) -> None:
doc = {
"question": "Q?",
"supporting_facts": {"title": ["A", "B", "A"]},
}
sample = map_hotpotqa_retrieval(doc)
assert sample is not None
assert sample.expected_uris == ("A", "B")
def test_map_retrieval_no_titles(self) -> None:
doc = {"question": "Q?", "supporting_facts": {"title": []}}
assert map_hotpotqa_retrieval(doc) is None
def test_build_case(self) -> None:
doc = {
"id": "abc123",
"question": "What is X?",
"answer": "X is Y.",
"type": "comparison",
"level": "hard",
}
case = build_hotpotqa_case(5, doc)
assert case.name == "5_abc123"
assert case.inputs == "What is X?"
assert case.expected_output == "X is Y."
assert case.metadata == {
"question_id": "abc123",
"type": "comparison",
"level": "hard",
"case_index": "5",
}
def test_extract_unique_documents(self) -> None:
# Simulate a minimal dataset with context
dataset = [
{
"context": {
"title": ["Doc A", "Doc B"],
"sentences": [["Sentence 1."], ["Sentence 2.", " More."]],
}
},
{
"context": {
"title": ["Doc A", "Doc C"],
"sentences": [["Dupe."], ["Sentence 3."]],
}
},
]
docs = extract_unique_documents(dataset) # type: ignore[arg-type] # ty: ignore[invalid-argument-type]
assert len(docs) == 3
titles = [d["title"] for d in docs]
assert titles == ["Doc A", "Doc B", "Doc C"]
assert docs[1]["content"] == "Sentence 2. More."
class TestOpenRAGBench:
def test_map_document(self, tmp_path: Path) -> None:
# Pre-create a cached PDF

View file

@ -3,54 +3,6 @@ from unittest.mock import MagicMock
import pytest
from evaluations.evaluators.map import MAPEvaluator
from evaluations.evaluators.mrr import MRREvaluator
class TestMRREvaluator:
def setup_method(self) -> None:
self.evaluator = MRREvaluator()
def _make_ctx(
self, relevant_uris: list[str], retrieved_uris: list[str]
) -> MagicMock:
ctx = MagicMock()
ctx.metadata = {"relevant_uris": relevant_uris}
ctx.output = retrieved_uris
return ctx
def test_first_result_relevant(self) -> None:
ctx = self._make_ctx(["doc1"], ["doc1", "doc2", "doc3"])
assert self.evaluator.evaluate(ctx) == 1.0
def test_second_result_relevant(self) -> None:
ctx = self._make_ctx(["doc2"], ["doc1", "doc2", "doc3"])
assert self.evaluator.evaluate(ctx) == 0.5
def test_third_result_relevant(self) -> None:
ctx = self._make_ctx(["doc3"], ["doc1", "doc2", "doc3"])
assert self.evaluator.evaluate(ctx) == pytest.approx(1 / 3)
def test_no_relevant_found(self) -> None:
ctx = self._make_ctx(["doc_x"], ["doc1", "doc2", "doc3"])
assert self.evaluator.evaluate(ctx) == 0.0
def test_empty_retrieved(self) -> None:
ctx = self._make_ctx(["doc1"], [])
assert self.evaluator.evaluate(ctx) == 0.0
def test_multiple_relevant_returns_first_match(self) -> None:
ctx = self._make_ctx(["doc2", "doc3"], ["doc1", "doc2", "doc3"])
assert self.evaluator.evaluate(ctx) == 0.5
def test_none_metadata(self) -> None:
ctx = MagicMock()
ctx.metadata = None
ctx.output = ["doc1"]
assert self.evaluator.evaluate(ctx) == 0.0
def test_empty_relevant_uris(self) -> None:
ctx = self._make_ctx([], ["doc1", "doc2"])
assert self.evaluator.evaluate(ctx) == 0.0
class TestMAPEvaluator:

View file

@ -71,7 +71,6 @@ members = ["haiku_rag_slim", "evaluations"]
dev = [
"haiku.rag-evals",
"haiku.rag-slim[ingester]",
"datasets>=4.8.4",
"zensical",
"pre-commit>=4.5.1",
"pydantic-ai-slim[anthropic]",

View file

@ -1,8 +1,9 @@
import json
import logging
import os
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from typing import TYPE_CHECKING, Any
# Prevent tests from loading user's local haiku.rag.yaml by setting env var
# to a test config file BEFORE any haiku.rag imports.
@ -17,14 +18,12 @@ embeddings:
vector_dim: 2560
""")
os.environ["HAIKU_RAG_CONFIG_PATH"] = str(_test_config_path)
os.environ["HF_HUB_OFFLINE"] = "1"
import pydantic_ai.models # noqa: E402
import pytest # noqa: E402
import yaml # noqa: E402
if TYPE_CHECKING:
from datasets import Dataset
from vcr import VCR
setattr(pydantic_ai.models, "ALLOW_MODEL_REQUESTS", False)
@ -32,18 +31,10 @@ logging.getLogger("vcr.cassette").setLevel(logging.WARNING)
@pytest.fixture(scope="session")
def qa_corpus() -> "Dataset":
from datasets import Dataset, load_dataset, load_from_disk
ds_path = Path(__file__).parent / "data" / "dataset"
ds_path.mkdir(parents=True, exist_ok=True)
try:
return cast(Dataset, load_from_disk(ds_path))
except FileNotFoundError:
ds: Dataset = load_dataset("ServiceNow/repliqa")["repliqa_3"]
corpus = ds.filter(lambda doc: doc["document_topic"] == "News Stories")
corpus.save_to_disk(ds_path)
return corpus
def qa_corpus() -> list[dict[str, str]]:
corpus_path = Path(__file__).parent / "data" / "qa_corpus.json"
with open(corpus_path) as f:
return json.load(f)
@pytest.fixture

View file

@ -1,106 +0,0 @@
{
"builder_name": "parquet",
"citation": "",
"config_name": "default",
"dataset_name": "repliqa",
"dataset_size": 648387648,
"description": "",
"download_checksums": {
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_0-00000-of-00001.parquet": {
"num_bytes": 20855947,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_1-00000-of-00001.parquet": {
"num_bytes": 20903799,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_2-00000-of-00001.parquet": {
"num_bytes": 20876034,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_3-00000-of-00001.parquet": {
"num_bytes": 20919214,
"checksum": null
},
"hf://datasets/ServiceNow/repliqa@bc880adc948fd3a70d5f8b2b3a1d1ee90d820dbd/data/repliqa_4-00000-of-00001.parquet": {
"num_bytes": 20866527,
"checksum": null
}
},
"download_size": 104421521,
"features": {
"document_id": {
"dtype": "string",
"_type": "Value"
},
"document_topic": {
"dtype": "string",
"_type": "Value"
},
"document_path": {
"dtype": "string",
"_type": "Value"
},
"document_extracted": {
"dtype": "string",
"_type": "Value"
},
"question_id": {
"dtype": "string",
"_type": "Value"
},
"question": {
"dtype": "string",
"_type": "Value"
},
"answer": {
"dtype": "string",
"_type": "Value"
},
"long_answer": {
"dtype": "string",
"_type": "Value"
}
},
"homepage": "",
"license": "",
"size_in_bytes": 752809169,
"splits": {
"repliqa_0": {
"name": "repliqa_0",
"num_bytes": 129579219,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_1": {
"name": "repliqa_1",
"num_bytes": 129844840,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_2": {
"name": "repliqa_2",
"num_bytes": 129338522,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_3": {
"name": "repliqa_3",
"num_bytes": 129981017,
"num_examples": 17955,
"dataset_name": "repliqa"
},
"repliqa_4": {
"name": "repliqa_4",
"num_bytes": 129644050,
"num_examples": 17950,
"dataset_name": "repliqa"
}
},
"version": {
"version_str": "0.0.0",
"major": 0,
"minor": 0,
"patch": 0
}
}

View file

@ -1,13 +0,0 @@
{
"_data_files": [
{
"filename": "data-00000-of-00001.arrow"
}
],
"_fingerprint": "2e29f63d782f12f0",
"_format_columns": null,
"_format_kwargs": {},
"_format_type": null,
"_output_all_columns": false,
"_split": "repliqa_3"
}

252
tests/data/qa_corpus.json Normal file

File diff suppressed because one or more lines are too long

View file

@ -1,5 +1,4 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
@ -7,7 +6,9 @@ from haiku.rag.store.models.chunk import Chunk, ChunkMetadata, SearchResult
@pytest.mark.vcr()
async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
async def test_chunk_repository_operations(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test ChunkRepository operations."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Get the first document from the corpus
@ -46,7 +47,9 @@ async def test_chunk_repository_operations(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path):
async def test_chunk_repository_pagination(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test ChunkRepository pagination with get_by_document_id and count_by_document_id."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Get the first document from the corpus (should produce multiple chunks)
@ -95,7 +98,7 @@ async def test_chunk_repository_pagination(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_chunking_pipeline(qa_corpus: Dataset, temp_db_path):
async def test_chunking_pipeline(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test document chunking using client primitives."""
from haiku.rag.client import HaikuRAG
from haiku.rag.embeddings import embed_chunks

View file

@ -2,7 +2,6 @@ from pathlib import Path
from unittest.mock import AsyncMock, Mock, patch
import pytest
from datasets import Dataset
from transformers import AutoTokenizer
from haiku.rag.chunkers import get_chunker
@ -18,7 +17,7 @@ def vcr_cassette_dir():
@pytest.mark.asyncio
async def test_local_chunker(qa_corpus: Dataset):
async def test_local_chunker(qa_corpus: list[dict[str, str]]):
"""Test DoclingLocalChunker with real document."""
chunker = DoclingLocalChunker()
doc_text = qa_corpus[0]["document_extracted"]
@ -106,7 +105,7 @@ def test_tokenizer_cached_across_chunker_instances():
@pytest.mark.asyncio
async def test_local_chunker_hierarchical(qa_corpus: Dataset):
async def test_local_chunker_hierarchical(qa_corpus: list[dict[str, str]]):
"""Test DoclingLocalChunker with hierarchical chunking."""
config = AppConfig()
config.processing.chunker_type = "hierarchical"

View file

@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, patch
import httpx
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
@ -20,7 +19,7 @@ def vcr_cassette_dir():
@pytest.mark.vcr()
async def test_client_document_crud(qa_corpus: Dataset, temp_db_path):
async def test_client_document_crud(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test HaikuRAG CRUD operations for documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Get test data
@ -123,7 +122,7 @@ async def test_client_resolve_document(temp_db_path):
@pytest.mark.vcr()
async def test_client_update_document(qa_corpus: Dataset, temp_db_path):
async def test_client_update_document(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test updating document with individual parameters."""
async with HaikuRAG(temp_db_path, create=True) as client:
# Get test data

View file

@ -1,5 +1,4 @@
import pytest
from datasets import Dataset
from haiku.rag.store.engine import Store
from haiku.rag.store.models.document import Document
@ -8,7 +7,7 @@ from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_document_list_excludes_content_by_default(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""list_all excludes content and docling_document by default."""
async with Store(temp_db_path, create=True) as store:
@ -34,7 +33,7 @@ async def test_document_list_excludes_content_by_default(
@pytest.mark.asyncio
async def test_document_list_includes_content_when_requested(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""list_all returns content when include_content=True."""
async with Store(temp_db_path, create=True) as store:
@ -51,7 +50,7 @@ async def test_document_list_includes_content_when_requested(
@pytest.mark.asyncio
async def test_document_list_with_filter(qa_corpus: Dataset, temp_db_path):
async def test_document_list_with_filter(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test listing documents with filter clause."""
async with Store(temp_db_path, create=True) as store:
doc_repo = DocumentRepository(store)
@ -229,7 +228,7 @@ def test_get_page_images():
@pytest.mark.asyncio
async def test_get_docling_data_loads_only_docling_columns(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""get_docling_data returns docling blob without loading content."""
import json
@ -276,7 +275,9 @@ async def test_get_docling_data_loads_only_docling_columns(
@pytest.mark.asyncio
async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_db_path):
async def test_get_pages_data_loads_only_pages_column(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""get_pages_data returns only page image data for a document."""
import json
@ -309,7 +310,7 @@ async def test_get_pages_data_loads_only_pages_column(qa_corpus: Dataset, temp_d
@pytest.mark.asyncio
async def test_get_pages_data_none_for_markdown_document(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Markdown documents have no page images — get_pages_data returns None pages."""
async with Store(temp_db_path, create=True) as store:
@ -330,7 +331,7 @@ async def test_get_pages_data_none_for_markdown_document(
@pytest.mark.asyncio
async def test_document_get_by_uri_with_special_characters(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test get_by_uri handles URIs with special characters like single quotes."""
async with Store(temp_db_path, create=True) as store:

View file

@ -4,7 +4,6 @@ from pathlib import Path
from typing import TypedDict
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG, RebuildMode
@ -19,10 +18,10 @@ class ChunkData(TypedDict):
@pytest.mark.vcr()
async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_full(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test full rebuild: converts, chunks, and embeds all documents."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
assert doc.docling_document is not None
@ -49,10 +48,10 @@ async def test_rebuild_full(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_embed_only(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test embed-only rebuild: keeps chunks, only regenerates embeddings."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
original_docling_json = doc.docling_document
@ -86,7 +85,7 @@ async def test_rebuild_embed_only(qa_corpus: Dataset, temp_db_path):
@pytest.mark.vcr()
async def test_rebuild_embed_only_multi_doc_streams_via_staging(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Embed-only rebuild with multiple documents preserves chunks via staging.
@ -100,8 +99,8 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
- the rebuild yields every document with chunks.
"""
async with HaikuRAG(temp_db_path, create=True) as client:
doc1 = await client.create_document(content=qa_corpus["document_extracted"][0])
doc2 = await client.create_document(content=qa_corpus["document_extracted"][1])
doc1 = await client.create_document(content=qa_corpus[0]["document_extracted"])
doc2 = await client.create_document(content=qa_corpus[1]["document_extracted"])
assert doc1.id is not None and doc2.id is not None
chunks_before_1 = await client.chunk_repository.get_by_document_id(doc1.id)
@ -133,7 +132,9 @@ async def test_rebuild_embed_only_multi_doc_streams_via_staging(
@pytest.mark.vcr()
async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_drops_leftover_staging_table(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Staging table without marker is treated as partial phase 1 and dropped.
Simulates a phase-1 interruption by creating only the staging table (no
@ -143,7 +144,7 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
from haiku.rag.client.rebuild import _StagingChunkRecord
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
# Simulate a partial phase 1 (staging exists, marker absent).
@ -167,7 +168,7 @@ async def test_rebuild_drops_leftover_staging_table(qa_corpus: Dataset, temp_db_
@pytest.mark.vcr()
async def test_rebuild_resumes_phase2_from_staging_after_crash(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Marker + staging present → phase 2 resumes from staging instead of
redoing phase 1.
@ -184,7 +185,7 @@ async def test_rebuild_resumes_phase2_from_staging_after_crash(
)
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
original_chunks = await client.chunk_repository.get_by_document_id(doc.id)
assert original_chunks
@ -280,7 +281,7 @@ async def test_rebuild_drops_orphan_marker(temp_db_path):
@pytest.mark.vcr()
async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Staging + marker from a prior embed-only crash → dropped on RECHUNK.
@ -295,7 +296,7 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
)
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
await client.store.db.create_table(
@ -317,10 +318,12 @@ async def test_rebuild_non_embed_mode_drops_staging_recovery_state(
@pytest.mark.vcr()
async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_embed_only_skips_unchanged(
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test embed-only rebuild skips chunks with unchanged embeddings."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
# Get embeddings before rebuild
@ -354,7 +357,7 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
@pytest.mark.vcr()
async def test_rebuild_embed_only_with_changed_vector_dim(
qa_corpus: Dataset, temp_db_path
qa_corpus: list[dict[str, str]], temp_db_path
):
"""Test embed-only rebuild when vector dimension changes.
@ -373,7 +376,7 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
# Step 1: Create a database with normal 2560-dim embeddings
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
chunks_before = await client.chunk_repository.get_by_document_id(doc.id)
@ -463,10 +466,10 @@ async def test_rebuild_embed_only_with_changed_vector_dim(
@pytest.mark.vcr()
async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
async def test_rebuild_rechunk(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test rechunk rebuild: re-chunks from content without accessing source files."""
async with HaikuRAG(temp_db_path, create=True) as client:
doc = await client.create_document(content=qa_corpus["document_extracted"][0])
doc = await client.create_document(content=qa_corpus[0]["document_extracted"])
assert doc.id is not None
assert doc.docling_document is not None

View file

@ -1,5 +1,4 @@
import pytest
from datasets import Dataset
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
@ -7,7 +6,7 @@ from haiku.rag.store.models import SearchResult
@pytest.mark.vcr()
async def test_search_qa_corpus(qa_corpus: Dataset, temp_db_path):
async def test_search_qa_corpus(qa_corpus: list[dict[str, str]], temp_db_path):
"""Test that documents can be found by searching with their associated questions."""
async with HaikuRAG(db_path=temp_db_path, config=Config, create=True) as client:
# Load unique documents (limited to 10)

View file

@ -1479,7 +1479,6 @@ tui = [
[package.dev-dependencies]
dev = [
{ name = "datasets" },
{ name = "haiku-rag-evals" },
{ name = "haiku-rag-slim", extra = ["ingester"] },
{ name = "pre-commit" },
@ -1506,7 +1505,6 @@ provides-extras = ["tui", "s3", "cross-encoder", "ingester"]
[package.metadata.requires-dev]
dev = [
{ name = "datasets", specifier = ">=4.8.4" },
{ name = "haiku-rag-evals", editable = "evaluations" },
{ name = "haiku-rag-slim", extras = ["ingester"], editable = "haiku_rag_slim" },
{ name = "pre-commit", specifier = ">=4.5.1" },