Merge pull request #543 from Cwiesen/feat-evaluations-search-filter

feat: add search_filter to evaluations

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### 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.
+
 ### Removed
 
 - `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`.

docs/benchmarks.md

@@ -65,6 +65,7 @@ evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.l
 - `--name NAME` - Override the evaluation name
 - `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers.
 - `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`).
+- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)).
 
 If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
 
@@ -86,6 +87,25 @@ evaluations:
         enable_thinking: true
 ```
 
+### Restricting the corpus
+
+When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`). Each dataset writes its own URIs: `orb_text` uses bare arXiv ids such as `2407.01528v3`, `hotpotqa` uses page titles.
+
+```bash
+evaluations run orb_text --skip-db --config haiku.rag.s3.yaml \
+  --filter "uri LIKE '2407%'"
+```
+
+If the corpora are distinguished by a tag rather than by URI, attach it at ingest time as document metadata and match it with `LIKE`. `metadata` is stored as a `json.dumps` string, so there is no JSON subfield access — match the serialized key/value, including the space after the colon:
+
+```bash
+evaluations run orb_text --skip-db --filter "metadata LIKE '%\"corpus\": \"orb_text\"%'"
+```
+
+The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `document_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results.
+
+Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus.
+
 ## Methodology
 
 ### Retrieval Metrics

evaluations/evaluations/benchmark.py

@@ -65,6 +65,7 @@ def build_experiment_metadata(
     judge_config: ModelConfig | None = None,
     target: Target = "rag-capability",
     capability_config: ModelConfig | None = None,
+    document_filter: str | None = None,
 ) -> dict[str, Any]:
     """Build experiment metadata for Logfire tracking."""
     metadata: dict[str, Any] = {
@@ -88,6 +89,7 @@ def build_experiment_metadata(
         "qa_enable_thinking": config.qa.model.enable_thinking,
         "qa_extra_body": config.qa.model.extra_body,
         "qa_max_searches": config.qa.max_searches,
+        "document_filter": document_filter,
     }
     if judge_config is not None:
         metadata.update(
@@ -191,6 +193,7 @@ async def run_retrieval_benchmark(
     name: str | None = None,
     db_path: Path | None = None,
     multimodal_only: bool = False,
+    document_filter: str | None = None,
 ) -> dict[str, float] | None:
     if spec.retrieval_loader is None or spec.retrieval_mapper is None:
         console.print("Skipping retrieval benchmark; no retrieval config.")
@@ -246,7 +249,9 @@ async def run_retrieval_benchmark(
     async with HaikuRAG(db, config=config, read_only=True) as rag:
 
         async def retrieval_target(question: str) -> list[str]:
-            chunks = await rag.search(query=question, limit=5, include_images=False)
+            chunks = await rag.search(
+                query=question, limit=5, include_images=False, filter=document_filter
+            )
 
             seen = set()
             identifiers = []
@@ -264,6 +269,7 @@ async def run_retrieval_benchmark(
             dataset_key=spec.key,
             test_cases=len(cases),
             config=config,
+            document_filter=document_filter,
         )
 
         report = await dataset.evaluate(
@@ -364,6 +370,7 @@ async def run_qa_benchmark(
     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:
     corpus = spec.qa_loader()
     corpus = _filter_qa_corpus(corpus, case_ids)
@@ -419,6 +426,7 @@ async def run_qa_benchmark(
         judge_config=judge_config,
         target=target,
         capability_config=capability_config,
+        document_filter=document_filter,
     )
 
     async def _evaluate(answer_fn: Callable[[str], Awaitable[str]]):
@@ -440,6 +448,7 @@ async def run_qa_benchmark(
             config=config,
             question=question,
             capability_model=resolved_capability_model,
+            document_filter=document_filter,
         )
         set_eval_attribute("cited_uris", result.cited_uris)
         set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids)
@@ -530,7 +539,11 @@ async def evaluate_dataset(
     target: Target = "rag-capability",
     capability_model: ModelConfig | None = None,
     case_ids: set[str] | None = None,
+    document_filter: str | None = None,
 ) -> None:
+    if document_filter is not None:
+        console.print(f"Document filter: {document_filter}", style="dim")
+
     if not skip_db:
         console.print(f"Using dataset: {spec.key}", style="bold magenta")
         await populate_db(
@@ -546,6 +559,7 @@ async def evaluate_dataset(
             name=name,
             db_path=db_path,
             multimodal_only=multimodal_only,
+            document_filter=document_filter,
         )
 
     if not skip_qa:
@@ -562,6 +576,7 @@ async def evaluate_dataset(
             target=target,
             capability_model=capability_model,
             case_ids=case_ids,
+            document_filter=document_filter,
         )
 
 
@@ -651,6 +666,17 @@ def run(
             "analysis.model when --target is analysis-capability) from the config."
         ),
     ),
+    document_filter: str | None = typer.Option(
+        None,
+        "--filter",
+        "-f",
+        help=(
+            "SQL WHERE clause over document columns (id, uri, title, "
+            "created_at, updated_at, metadata) restricting every benchmark "
+            "search, e.g. \"uri LIKE '%arxiv%'\". metadata is stored as a "
+            "string, so match it with LIKE."
+        ),
+    ),
     filter_ids: Path | None = typer.Option(
         None,
         "--filter-ids",
@@ -688,6 +714,7 @@ def run(
             target=target_value,
             capability_model=capability_model_config,
             case_ids=_load_case_ids(filter_ids),
+            document_filter=document_filter,
         )
     )
 

evaluations/tests/test_benchmark.py

@@ -17,16 +17,16 @@ from haiku.rag.config.models import AppConfig, ModelConfig
 
 def _stub_spec(**overrides) -> DatasetSpec:
     """A DatasetSpec whose loaders/mappers are inert, for tests that only
-    exercise the surrounding plumbing."""
-    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,
-        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]
-        **overrides,
-    )
+    exercise the surrounding plumbing. Any field can be overridden."""
+    fields: dict = {
+        "key": "test",
+        "db_filename": "test.lancedb",
+        "document_loader": lambda: None,
+        "document_mapper": lambda doc: None,
+        "qa_loader": lambda: [],
+        "qa_case_builder": lambda idx, doc: None,
+    }
+    return DatasetSpec(**{**fields, **overrides})
 
 
 class TestBuildExperimentMetadata:
@@ -581,6 +581,120 @@ class TestRetrievalTarget:
         assert result["map"] == 0.5
 
 
+class TestDocumentFilterThreading:
+    """The filter must reach both benchmark phases, so retrieval and QA score
+    the same subset of the database."""
+
+    def test_metadata_records_filter(self) -> None:
+        result = build_experiment_metadata(
+            dataset_key="test",
+            test_cases=1,
+            config=AppConfig(),
+            document_filter="uri LIKE '%arxiv%'",
+        )
+        assert result["document_filter"] == "uri LIKE '%arxiv%'"
+
+    def test_metadata_filter_is_none_when_unset(self) -> None:
+        result = build_experiment_metadata(
+            dataset_key="test", test_cases=1, config=AppConfig()
+        )
+        assert result["document_filter"] is None
+
+    @pytest.mark.asyncio
+    async def test_retrieval_search_receives_filter(self, tmp_path: Path) -> None:
+        from haiku.rag.store.models.chunk import SearchResult
+
+        from evaluations.benchmark import run_retrieval_benchmark
+        from evaluations.config import RetrievalSample
+        from evaluations.evaluators import MAPEvaluator
+
+        searches: list[dict] = []
+
+        class FakeRag:
+            async def search(self, **kwargs) -> list[SearchResult]:
+                searches.append(kwargs)
+                return [SearchResult(content="x", score=1.0, document_uri="uri-x")]
+
+        spec = _stub_spec(
+            retrieval_loader=lambda: [{"q": "What is X?", "uris": ("uri-x",)}],
+            retrieval_mapper=lambda d: RetrievalSample(
+                question=d["q"], expected_uris=d["uris"]
+            ),
+            retrieval_evaluator=MAPEvaluator(),
+        )
+
+        with patch("evaluations.benchmark.HaikuRAG") as mock_haiku:
+            mock_haiku.return_value.__aenter__.return_value = FakeRag()
+            await run_retrieval_benchmark(
+                spec,
+                AppConfig(),
+                db_path=tmp_path / "test.lancedb",
+                document_filter="uri LIKE '%arxiv%'",
+            )
+
+        assert searches[0]["filter"] == "uri LIKE '%arxiv%'"
+
+    @pytest.mark.asyncio
+    async def test_qa_capability_run_receives_filter(self, tmp_path: Path) -> None:
+        from pydantic_evals import Case
+
+        from evaluations.capability_runner import CapabilityRunResult
+        from evaluations.evaluators import NumberMatchEvaluator
+
+        # A deterministic evaluator, so no judge model is constructed.
+        spec = _stub_spec(
+            qa_loader=lambda: [{"question": "What is X?", "answer": "42"}],
+            qa_case_builder=lambda idx, doc: Case(
+                name=f"case-{idx}",
+                inputs=doc["question"],
+                expected_output=doc["answer"],
+            ),
+            qa_evaluator=NumberMatchEvaluator(),
+        )
+
+        with patch(
+            "evaluations.benchmark.run_capability_question",
+            new_callable=AsyncMock,
+            return_value=CapabilityRunResult(answer="ANSWER: 42"),
+        ) as mock_run:
+            await run_qa_benchmark(
+                spec,
+                AppConfig(),
+                db_path=tmp_path / "test.lancedb",
+                document_filter="uri LIKE '%arxiv%'",
+            )
+
+        mock_run.assert_awaited_once()
+        assert mock_run.call_args[1]["document_filter"] == "uri LIKE '%arxiv%'"
+
+    @pytest.mark.asyncio
+    async def test_evaluate_dataset_passes_filter_to_both_phases(self) -> None:
+        expected = """metadata LIKE '%"corpus": "orb_text"%'"""
+
+        with (
+            patch(
+                "evaluations.benchmark.run_retrieval_benchmark", new_callable=AsyncMock
+            ) as mock_retrieval,
+            patch(
+                "evaluations.benchmark.run_qa_benchmark", new_callable=AsyncMock
+            ) as mock_qa,
+        ):
+            await evaluate_dataset(
+                spec=_stub_spec(),
+                config=AppConfig(),
+                skip_db=True,
+                skip_retrieval=False,
+                skip_qa=False,
+                limit=None,
+                name=None,
+                db_path=None,
+                document_filter=expected,
+            )
+
+        assert mock_retrieval.call_args[1]["document_filter"] == expected
+        assert mock_qa.call_args[1]["document_filter"] == expected
+
+
 class TestEvaluateDatasetCaseIds:
     def _spec(self) -> DatasetSpec:
         return _stub_spec()