Find a document in whichever database holds it

`get_document_by_id`, `get_document_by_uri` and `get_chunk_by_id` read
through repositories a client covering a set does not have, so a lookup by
identifier raised AttributeError and `resolve_document` with it. An
identifier says nothing about which database holds it, so every database is
asked at once and the first that has it, in configured order, answers.

On the evaluation side, `--db` overrides the configured set as the CLI
documents, and population refuses a set rather than ingesting into a
database the run would not read. A case filter matching nothing raises
instead of reporting 0.0000 as though it were a score.

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

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.

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,

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]:

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,
+        )

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
 

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"]:

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):