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}