diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 6c78205b..ac056341 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -50,6 +50,13 @@ async def evaluate_dataset( console.print(f"Document filter: {document_filter}", style="dim") if not skip_db: + if spec.covers_a_set(config, db_path): + raise ValueError( + "lancedb.databases names several databases and population writes " + "to one, so it would ingest into a database the run does not " + "read. Pass --skip-db to evaluate the configured set, or --db " + "PATH to populate and evaluate one database." + ) console.print(f"Using dataset: {spec.key}", style="bold magenta") await populate_db( spec, config, db_path=db_path, vacuum_interval=vacuum_interval diff --git a/evaluations/evaluations/config.py b/evaluations/evaluations/config.py index 8e94f504..30c8956a 100644 --- a/evaluations/evaluations/config.py +++ b/evaluations/evaluations/config.py @@ -83,13 +83,14 @@ class DatasetSpec: compaction: bool = False experiment_metadata: dict[str, Any] | None = None - def covers_a_set(self, config) -> bool: + def covers_a_set(self, config, override_path: Path | None = None) -> bool: """Whether `lancedb.databases` names the databases to evaluate over. - A path names one database and wins over the configured set, so a run over - a set has to pass none — the client resolves it. + A path names one database and wins over the configured set, both when it + comes from `--db` and when the client resolves it, so a run over a set is + one where the configuration names several and nobody named a path. """ - return bool(config.lancedb.databases) + return bool(config.lancedb.databases) and override_path is None def db_path(self, override_path: Path | None = None) -> Path: """Get the database path. diff --git a/evaluations/evaluations/qa.py b/evaluations/evaluations/qa.py index e420d8b8..75816355 100644 --- a/evaluations/evaluations/qa.py +++ b/evaluations/evaluations/qa.py @@ -186,11 +186,20 @@ def _refusal_metrics(report_cases) -> tuple[float, float, int, int] | None: def _filter_qa_corpus(corpus, case_ids: set[str] | None): """Keep only rows whose ``id`` is in ``case_ids`` (failure-subset reruns). - Returns the corpus unchanged when ``case_ids`` is None. + Returns the corpus unchanged when ``case_ids`` is None. Matching nothing + raises: a dataset keying its rows by another name leaves every case filtered + out, and a run of no cases otherwise reports 0.0000 as though it were a score. """ if case_ids is None: return corpus - return corpus.filter(lambda row: row.get("id") in case_ids) + filtered = corpus.filter(lambda row: row.get("id") in case_ids) + if len(filtered) == 0: + raise ValueError( + f"--filter-ids matched none of the {len(corpus)} cases. " + "Check that the ids belong to this dataset and that its rows are " + "keyed by `id`." + ) + return filtered class _QARun(NamedTuple): @@ -243,7 +252,7 @@ def _prepare_qa_run( return _QARun( cases=cases, - db=None if spec.covers_a_set(config) else spec.db_path(db_path), + db=None if spec.covers_a_set(config, db_path) else spec.db_path(db_path), judge_config=judge_config, eval_name=eval_name, experiment_metadata=experiment_metadata, diff --git a/evaluations/evaluations/retrieval.py b/evaluations/evaluations/retrieval.py index 48d24690..ad6d84c3 100644 --- a/evaluations/evaluations/retrieval.py +++ b/evaluations/evaluations/retrieval.py @@ -72,7 +72,7 @@ async def run_retrieval_benchmark( evaluators=list(spec.retrieval_evaluators), ) - db = None if spec.covers_a_set(config) else spec.db_path(db_path) + db = None if spec.covers_a_set(config, db_path) else spec.db_path(db_path) async with HaikuRAG(db, config=config, read_only=True) as rag: async def retrieval_target(question: str) -> list[str]: diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 6b2a82e6..3121e883 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -1288,3 +1288,54 @@ class TestEvaluateDatasetCaseIds: ) mock_qa.assert_called_once() assert mock_qa.call_args[1]["case_ids"] == {"finqa_dev_16", "finqa_dev_66"} + + +def test_a_case_filter_matching_nothing_raises(): + """A run of no cases reports 0.0000, which reads like a score rather than a + mistake — so an id set that matches nothing fails instead.""" + import pytest + from datasets import Dataset + + from evaluations.qa import _filter_qa_corpus + + corpus = Dataset.from_list([{"query_id": "a"}, {"query_id": "b"}]) + + with pytest.raises(ValueError, match="matched none"): + _filter_qa_corpus(corpus, {"a"}) + + +def test_a_case_filter_that_matches_keeps_those_rows(): + from datasets import Dataset + + from evaluations.qa import _filter_qa_corpus + + corpus = Dataset.from_list([{"id": "a"}, {"id": "b"}, {"id": "c"}]) + + assert _filter_qa_corpus(corpus, {"a", "c"})["id"] == ["a", "c"] + + +async def test_population_refuses_a_configured_set(): + """Population writes to one database, so a set would be ingested into a + database the run never reads.""" + import pytest + from haiku.rag.config.models import AppConfig, LanceDBConfig + + from evaluations.benchmark import evaluate_dataset + from evaluations.datasets import DATASETS + + spec = next(iter(DATASETS.values())) + config = AppConfig( + lancedb=LanceDBConfig(databases={"a": "/a.lancedb", "b": "/b.lancedb"}) + ) + + with pytest.raises(ValueError, match="--skip-db"): + await evaluate_dataset( + spec, + config, + skip_db=False, + skip_retrieval=True, + skip_qa=True, + limit=None, + name=None, + db_path=None, + ) diff --git a/evaluations/tests/test_config.py b/evaluations/tests/test_config.py index a53c1b6a..6cafc5c4 100644 --- a/evaluations/tests/test_config.py +++ b/evaluations/tests/test_config.py @@ -163,6 +163,22 @@ class TestCoversASet: assert spec.covers_a_set(config) is True + def test_a_named_path_overrides_the_set(self): + """`--db` is documented as an override, so it names the one database to + evaluate even when the configuration names several.""" + from pathlib import Path as _Path + + from haiku.rag.config.models import AppConfig, LanceDBConfig + + from evaluations.datasets import DATASETS + + spec = next(iter(DATASETS.values())) + config = AppConfig( + lancedb=LanceDBConfig(databases={"a": "/a.lancedb", "b": "/b.lancedb"}) + ) + + assert spec.covers_a_set(config, _Path("/chosen.lancedb")) is False + def test_one_database_is_not_a_set(self): from haiku.rag.config.models import AppConfig diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 85c3f90f..44641ec4 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -4,13 +4,13 @@ import json import logging import mimetypes import tempfile -from collections.abc import AsyncGenerator, Sequence +from collections.abc import AsyncGenerator, Callable, Coroutine, Sequence from enum import Enum from functools import cached_property from itertools import zip_longest from pathlib import Path from time import monotonic -from typing import TYPE_CHECKING, overload +from typing import TYPE_CHECKING, Any, overload from urllib.parse import urlparse import httpx @@ -505,6 +505,10 @@ class HaikuRAG: Returns: The Document instance if found, None otherwise. """ + if self._federated: + return await self._from_any_covered( + lambda owner: owner.get_document_by_id(document_id) + ) return await self.document_repository.get_by_id(document_id) async def get_chunk_by_id(self, chunk_id: str) -> Chunk | None: @@ -516,6 +520,10 @@ class HaikuRAG: Returns: The Chunk instance if found, None otherwise. """ + if self._federated: + return await self._from_any_covered( + lambda owner: owner.get_chunk_by_id(chunk_id) + ) return await self.chunk_repository.get_by_id(chunk_id) async def get_picture_bytes( @@ -553,6 +561,10 @@ class HaikuRAG: Returns: The Document instance if found, None otherwise. """ + if self._federated: + return await self._from_any_covered( + lambda owner: owner.get_document_by_uri(uri) + ) return await self.document_repository.get_by_uri(uri) async def resolve_document(self, id_or_title: str) -> Document | None: @@ -682,6 +694,22 @@ class HaikuRAG: return sum(counts) return await self.document_repository.count(filter=filter) + async def _from_any_covered( + self, lookup: "Callable[[HaikuRAG], Coroutine[Any, Any, Any]]" + ) -> Any: + """The first result `lookup` finds in the databases this client covers. + + An id or a URI says nothing about which database holds it, so every + database is asked at once and the first that has it, in configured order, + answers. Asking in turn would cost a round trip per database for an + identifier that is missing or held by the last of them. + """ + owners = await self.clients_covering() + for found in await asyncio.gather(*(lookup(owner) for owner in owners)): + if found is not None: + return found + return None + async def clients_covering( self, sources: list[str] | None = None ) -> list["HaikuRAG"]: diff --git a/tests/test_multi_db.py b/tests/test_multi_db.py index d245214d..23b752bc 100644 --- a/tests/test_multi_db.py +++ b/tests/test_multi_db.py @@ -223,6 +223,62 @@ class TestListingAcrossDatabases: assert [d.uri for d in docs] == ["test://beta/beta one"] +class TestLookupByIdentifier: + """An id or a URI says nothing about which database holds it, and a client + covering a set has no repositories of its own.""" + + @pytest.mark.asyncio + async def test_a_document_is_found_in_whichever_database_holds_it(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + beta = (await rag.clients_for(["beta"]))[0] + [target] = await beta.document_repository.list_all(limit=1) + assert target.id is not None + + found = await rag.get_document_by_id(target.id) + by_uri = await rag.get_document_by_uri("test://alpha/alpha one") + resolved = await rag.resolve_document(target.id) + + assert found is not None and found.uri == "test://beta/beta one" + assert by_uri is not None and by_uri.uri == "test://alpha/alpha one" + assert resolved is not None and resolved.uri == "test://beta/beta one" + + @pytest.mark.asyncio + async def test_a_chunk_is_found_in_whichever_database_holds_it(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + beta = (await rag.clients_for(["beta"]))[0] + [chunk] = await beta.chunk_repository.list_all(limit=1) + assert chunk.id is not None + + found = await rag.get_chunk_by_id(chunk.id) + + assert found is not None and found.content == "beta one" + + @pytest.mark.asyncio + async def test_an_unknown_identifier_is_absent_rather_than_an_error(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha one"]) + await _seed(config, "beta", ["beta one"]) + + async with HaikuRAG(config=config) as rag: + assert ( + await rag.get_document_by_id("00000000-0000-4000-8000-000000000000") + is None + ) + assert ( + await rag.get_chunk_by_id("00000000-0000-4000-8000-000000000000") + is None + ) + assert await rag.get_document_by_uri("test://nowhere") is None + + class TestFederatedSearch: @pytest.mark.asyncio async def test_results_carry_their_source(self, tmp_path):