From 645b6ede236e1b75a826e32364039ffd73b20dfa Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 31 Aug 2026 13:38:30 +0300 Subject: [PATCH] Make retrieval fetch depth a run option Hybrid search inside one database fuses its own vector and FTS rankings with lancedb's RRFReranker over exactly the requested limit, and both branch queries derive from the same inner query, so there is no branch-depth knob. Below roughly 50 candidates the two rankings stop overlapping, nothing sums, and the fusion degenerates: measured recall@5 on a single database was 0.000 at fetch 5, 10 and 20, then 0.267 at 50 and 0.350 at 100. A dataset's retrieval_limit therefore fixes which regime it measures, and comparing regimes would otherwise need one dataset per depth. Claude-Session: https://claude.ai/code/session_01WhudUtZm6qqiuv8Y1sbwSc --- CHANGELOG.md | 1 + evaluations/evaluations/benchmark.py | 11 +++++ evaluations/evaluations/retrieval.py | 4 +- evaluations/tests/test_mtrag_federated.py | 55 +++++++++++++++++++++++ 4 files changed, 70 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 10a71f47..fbefbdb9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added +- `evaluations run --retrieval-limit N`: candidates each database fetches during the retrieval benchmark, overriding the dataset's `retrieval_limit`. - `mtrag_federated` evaluation dataset and its reference config `evaluations/configs/mtrag_federated.yaml`: MTRAG ClapNQ partitioned by article title into `n` collections, scored on retrieval only with Recall@5/@10, nDCG@5 and MAP. `python -m evaluations.datasets.mtrag_federated --config REF --n N --out PATH` builds the partition and emits the config that searches it. ### Fixed diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 182d5977..d4541045 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -41,6 +41,7 @@ async def evaluate_dataset( vacuum_interval: int = 100, multimodal_only: bool = False, judge_model: ModelConfig | None = None, + retrieval_limit: int | None = None, target: Target = "rag-capability", capability_model: ModelConfig | None = None, case_ids: set[str] | None = None, @@ -72,6 +73,7 @@ async def evaluate_dataset( db_path=db_path, multimodal_only=multimodal_only, document_filter=document_filter, + retrieval_limit=retrieval_limit, ) if not skip_qa: @@ -169,6 +171,14 @@ def run( None, "--limit", help="Limit number of test cases for both retrieval and QA." ), name: str | None = typer.Option(None, "--name", help="Override evaluation name."), + retrieval_limit: int | None = typer.Option( + None, + "--retrieval-limit", + help=( + "Candidates each database fetches, overriding the dataset's. " + "Sets how deep hybrid search looks before its results are scored." + ), + ), vacuum_interval: int = typer.Option( 100, "--vacuum-interval", help="Vacuum every N documents during DB population." ), @@ -235,6 +245,7 @@ def run( vacuum_interval=vacuum_interval, multimodal_only=multimodal_only, judge_model=judge_model_config, + retrieval_limit=retrieval_limit, target=target_value, capability_model=capability_model_config, case_ids=_load_case_ids(filter_ids), diff --git a/evaluations/evaluations/retrieval.py b/evaluations/evaluations/retrieval.py index 9768c587..7a3e36c4 100644 --- a/evaluations/evaluations/retrieval.py +++ b/evaluations/evaluations/retrieval.py @@ -24,6 +24,7 @@ async def run_retrieval_benchmark( db_path: Path | None = None, multimodal_only: bool = False, document_filter: str | None = None, + retrieval_limit: int | 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.") @@ -72,6 +73,7 @@ async def run_retrieval_benchmark( evaluators=list(spec.retrieval_evaluators), ) + fetch = retrieval_limit or spec.retrieval_limit db = ( None if spec.uses_configured_databases(config, db_path) @@ -82,7 +84,7 @@ async def run_retrieval_benchmark( async def retrieval_target(question: str) -> list[str]: chunks = await rag.search( query=question, - limit=spec.retrieval_limit, + limit=fetch, include_images=False, filter=document_filter, ) diff --git a/evaluations/tests/test_mtrag_federated.py b/evaluations/tests/test_mtrag_federated.py index fd93357d..2368ada5 100644 --- a/evaluations/tests/test_mtrag_federated.py +++ b/evaluations/tests/test_mtrag_federated.py @@ -17,6 +17,22 @@ from evaluations.datasets.mtrag_federated import ( ) +def _smoke_config(): + """A config placing two databases, so the run resolves a federated client.""" + from haiku.rag.config.models import AppConfig + + return AppConfig.model_validate( + { + "lancedb": { + "databases": { + "clapnq_0": "/tmp/a.lancedb", + "clapnq_1": "/tmp/b.lancedb", + } + } + } + ) + + def record(passage_id: str, title: str) -> dict[str, str]: return {"_id": passage_id, "title": title, "text": f"text of {passage_id}"} @@ -193,3 +209,42 @@ class TestPoolComposition: def test_the_default_budget_clears_the_gold_floor(self) -> None: assert DEFAULT_BUDGET > GOLD_TITLE_FLOOR + + +class TestRetrievalLimitOverride: + async def test_override_replaces_the_spec_value(self, monkeypatch) -> None: + """Fetch depth is a run knob: hybrid search degenerates below roughly 50 + candidates, so every regime would otherwise need its own dataset.""" + seen: list[int | None] = [] + + async def fake_search(self, query, limit=None, **kwargs): # noqa: ANN001 + seen.append(limit) + return [] + + from haiku.rag.client import HaikuRAG + + monkeypatch.setattr(HaikuRAG, "search", fake_search) + from evaluations.retrieval import run_retrieval_benchmark + + await run_retrieval_benchmark( + MTRAG_FEDERATED_SPEC, + _smoke_config(), + limit=1, + retrieval_limit=77, + ) + assert seen and set(seen) == {77} + + async def test_spec_value_is_the_default(self, monkeypatch) -> None: + seen: list[int | None] = [] + + async def fake_search(self, query, limit=None, **kwargs): # noqa: ANN001 + seen.append(limit) + return [] + + from haiku.rag.client import HaikuRAG + + monkeypatch.setattr(HaikuRAG, "search", fake_search) + from evaluations.retrieval import run_retrieval_benchmark + + await run_retrieval_benchmark(MTRAG_FEDERATED_SPEC, _smoke_config(), limit=1) + assert seen and set(seen) == {MTRAG_FEDERATED_SPEC.retrieval_limit}