From 409f60e1bf0c46e907da8c29b6e1cdbcb390ecf0 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 80c95fb0..b46b52ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ and rerun `create-index`. ### 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}