Ask across databases and name the source of each citation

`ask(sources=[…])` scopes a question to some of the configured databases, carried
on the capability state so its search tool searches those. `Citation.source` names
the database a cited chunk came from, resolved from the search results the model
saw, which already carry it.

Context expansion routes each result through the database it came from: a
federating client has no repositories of its own.

The cite fallback, which looks up an id absent from this run's results, searches
only the selected databases. A chunk id says nothing about which database holds
it, so placing one means asking, and asking outside the selection would let a
question scoped to some databases cite another.

The loosely-specced client mocks in the capability tests now say they stand in for
a single-database client. A bare AsyncMock answers any attribute with a truthy
Mock, so `_federated` sent the fallback down the multi-database branch, and
`_source` reached a validated field.

CHANGELOG.md

@@ -39,6 +39,7 @@
 - The `haiku.rag` package declares the `jina` extra, so `provider: jina-local` is supported by declaration rather than through `cross-encoder`'s transitive `transformers` and `torch`. Raises the full package's torch floor to 2.0.
 - `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout.
 - `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured.
+- `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another.
 
 ### Changed
 

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -35,7 +35,7 @@ from haiku.rag.capabilities._tools import (
 from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config.models import AppConfig
-from haiku.rag.store.models.chunk import SearchResult
+from haiku.rag.store.models.chunk import Chunk, SearchResult
 from haiku.rag.store.models.citation import Citation, resolve_citations
 from haiku.rag.tools.search import build_image_content_from_results
 
@@ -91,6 +91,7 @@ class EvidenceState(BaseModel):
     citations: list[str] = Field(default_factory=list)
     evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord)
     document_filter: str | None = None
+    sources: list[str] | None = None
     searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
 
     def begin_invocation(self) -> None:
@@ -101,8 +102,8 @@ class EvidenceState(BaseModel):
         leaves a later citation unable to resolve against the expanded result the
         model saw, recording no provenance for it.
 
-        `document_filter` scopes the conversation, and `evidence` carries question
-        identity, so neither is working evidence.
+        `document_filter` and `sources` scope the conversation, and `evidence`
+        carries question identity, so none of them is working evidence.
         """
         self.citations.clear()
         self.searches.clear()
@@ -136,6 +137,21 @@ def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -
     return False
 
 
+async def _first_holding(
+    clients: "list[HaikuRAG]", chunk_id: str
+) -> "tuple[HaikuRAG, Chunk] | None":
+    """The first client holding this chunk, and the chunk.
+
+    A chunk id says nothing about which database holds it, so the only way to
+    place one is to ask. Returns None when none of them has it.
+    """
+    for client in clients:
+        chunk = await client.get_chunk_by_id(chunk_id)
+        if chunk is not None:
+            return client, chunk
+    return None
+
+
 @dataclass
 class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
     db_path: Path
@@ -468,6 +484,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
                 query,
                 limit=limit,
                 document_filter=self.state.document_filter,
+                sources=self.state.sources,
             )
         state = self.state
         # A model can search the same query twice with different limits, and the
@@ -504,20 +521,35 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
         if missing:
             async with self.rag_lock:
                 rag = await self._ensure_rag()
+                # A chunk id says nothing about which database holds it, so a
+                # federating client looks through the ones it covers.
+                if rag._federated:
+                    lookups = await rag.clients_for(
+                        getattr(self.state, "sources", None) or list(rag._federated)
+                    )
+                else:
+                    lookups = [rag]
                 synthetic: list[SearchResult] = []
-                documents: dict[str, Any] = {}
+                documents: dict[tuple[str | None, str], Any] = {}
                 for chunk_id in missing:
-                    chunk = await rag.get_chunk_by_id(chunk_id)
-                    if chunk is None or not chunk.document_id:
+                    found = await _first_holding(lookups, chunk_id)
+                    if found is None:
                         continue
-                    document = documents.get(chunk.document_id)
-                    if chunk.document_id not in documents:
-                        document = await rag.get_document_by_id(chunk.document_id)
-                        documents[chunk.document_id] = document
+                    owner, chunk = found
+                    if not chunk.document_id:
+                        continue
+                    key = (owner._source, chunk.document_id)
+                    if key not in documents:
+                        documents[key] = await owner.get_document_by_id(
+                            chunk.document_id
+                        )
+                    document = documents[key]
                     chunk.document_uri = document.uri if document else None
                     chunk.document_title = document.title if document else None
                     chunk.document_meta = document.metadata if document else {}
-                    synthetic.append(SearchResult.from_chunk(chunk, score=1.0))
+                    result = SearchResult.from_chunk(chunk, score=1.0)
+                    result.source = owner._source
+                    synthetic.append(result)
                 citations.extend(resolve_citations(missing, synthetic))
 
         if not citations:

haiku_rag_slim/haiku/rag/capabilities/_tools.py

@@ -18,9 +18,12 @@ async def search_corpus(
     query: str,
     limit: int | None = None,
     document_filter: str | None = None,
+    sources: list[str] | None = None,
 ) -> tuple[str, list[SearchResult]]:
     """Search and context-expand results for a capability tool."""
-    results = await rag.search(query, limit=limit, filter=document_filter)
+    results = await rag.search(
+        query, limit=limit, filter=document_filter, sources=sources
+    )
     results = await rag.expand_context(results)
     formatted = "\n\n---\n\n".join(
         result.format_for_agent(rank=index + 1, total=len(results))

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -635,10 +635,11 @@ class HaikuRAG:
         question: str,
         filter: str | None = None,
         images: Sequence[bytes] | None = None,
+        sources: list[str] | None = None,
     ) -> "tuple[str, list[Citation]]":
         from haiku.rag.client.agents import ask
 
-        return await ask(self, question, filter, images)
+        return await ask(self, question, filter, images, sources)
 
     async def analyze(
         self,

haiku_rag_slim/haiku/rag/client/agents.py

@@ -40,6 +40,7 @@ async def ask(
     question: str,
     filter: str | None = None,
     images: Sequence[bytes] | None = None,
+    sources: list[str] | None = None,
 ) -> "tuple[str, list[Citation]]":
     """Ask a question against the knowledge base via the RAG capability.
 
@@ -61,13 +62,17 @@ async def ask(
     from haiku.rag.utils import get_model
 
     capability = create_capability(
-        db_path=client.store.db_path,
+        db_path=None if client._federated else client.store.db_path,
         config=client._config,
         rag=client,
         defer_loading=False,
     )
     deps = _AgentDeps(
-        state={"rag": RAGState(document_filter=filter).model_dump(mode="json")}
+        state={
+            "rag": RAGState(document_filter=filter, sources=sources).model_dump(
+                mode="json"
+            )
+        }
     )
     user_prompt = _build_user_prompt(question, images, client._config.qa.model)
     model = get_model(client._config.qa.model, client._config)

haiku_rag_slim/haiku/rag/client/search.py

@@ -378,6 +378,28 @@ async def expand_context(
     chunks were created without docling metadata (e.g., custom chunks passed
     to import_document).
     """
+    # A federating client has no repositories of its own, so each result expands
+    # through the database it came from.
+    if client._federated:
+        by_source: dict[str, list[SearchResult]] = {}
+        unsourced: list[SearchResult] = []
+        for result in search_results:
+            if result.source:
+                by_source.setdefault(result.source, []).append(result)
+            else:
+                unsourced.append(result)
+        owners = await client.clients_for(list(by_source))
+        expanded_groups = await asyncio.gather(
+            *(
+                expand_context(owner, by_source[owner._source])
+                for owner in owners
+                if owner._source
+            )
+        )
+        merged = unsourced + [r for group in expanded_groups for r in group]
+        merged.sort(key=lambda r: r.score, reverse=True)
+        return merged
+
     from haiku.rag.context import expand_with_items, window_for
 
     max_chars = client._config.search.max_context_chars

haiku_rag_slim/haiku/rag/store/models/citation.py

@@ -22,6 +22,10 @@ class Citation(BaseModel):
     ``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
     into the cited result (always includes ``chunk_id``).
 
+    ``source`` names the configured database the cited chunk came from, and is
+    None when only one is configured. It is the name from ``lancedb.databases``,
+    never a path or URI.
+
     ``doc_item_refs`` are the ``self_ref`` values of every item in the cited
     content — the exact items the model saw. Visual grounding resolves bounding
     boxes from them so the rendered pages match the citation precisely.
@@ -37,6 +41,7 @@ class Citation(BaseModel):
 
     index: int | None = None
     document_id: str
+    source: str | None = None
     chunk_id: str
     chunk_ids: list[str] = Field(default_factory=list)
     chunk_meta: dict = Field(default_factory=dict)
@@ -69,6 +74,7 @@ def resolve_citations(
         citations.append(
             Citation(
                 document_id=r.document_id or "",
+                source=r.source,
                 chunk_id=chunk_id,
                 chunk_ids=r.chunk_ids or [chunk_id],
                 chunk_meta=r.chunk_meta,

tests/capabilities/test_capabilities.py

@@ -316,6 +316,10 @@ async def test_capability_isolated_per_run_and_round_trips_state(temp_db_path):
 async def test_run_error_closes_resources_and_propagates(temp_db_path):
     capability = create_rag(db_path=temp_db_path, config=AppConfig())
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     capability.rag = client
     error = RuntimeError("model failed")
 
@@ -385,6 +389,10 @@ async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db
     capability = create_rag(db_path=temp_db_path, config=AppConfig())
     capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     client.get_chunk_by_id.side_effect = [
         Chunk(id="chunk-1", document_id="doc-1", content="first"),
         Chunk(id="chunk-2", document_id="doc-1", content="second"),
@@ -411,6 +419,10 @@ async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path):
     capability = create_rag(db_path=temp_db_path, config=AppConfig())
     capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     client.get_chunk_by_id.side_effect = [
         Chunk(id="chunk-1", document_id="doc-1", content="first"),
         None,
@@ -455,6 +467,10 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path):
         },
     )
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     client.get_chunk_by_id.return_value = None
     capability.rag = client
 
@@ -1214,6 +1230,10 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path):
         return ModelResponse(parts=next(calls))
 
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     client.get_chunk_by_id.return_value = Chunk(
         id="chunk-1", document_id="doc-1", content="evidence"
     )
@@ -1315,6 +1335,10 @@ async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path):
     """Compaction rehydrates through the owner, which already holds the connection."""
     capability = create_rag(db_path=temp_db_path, config=AppConfig())
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     client.document_item_repository.get_picture_bytes.return_value = b"picture-bytes"
     capability.rag = client
 
@@ -1354,6 +1378,10 @@ async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_p
     capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
     capability.epoch = 5
     client = AsyncMock()
+    # Stands in for a single-database client: a bare AsyncMock's auto
+    # attributes are truthy Mocks, and `_source` reaches a validated field.
+    client._federated = {}
+    client._source = None
     client.get_chunk_by_id.return_value = Chunk(
         id="chunk-1", document_id="doc-1", content="evidence"
     )

tests/test_multi_db_ask.py

@@ -0,0 +1,201 @@
+import pytest
+from pydantic_ai import ModelRetry
+
+from haiku.rag.capabilities.rag import RAGState, create_capability
+from haiku.rag.client import HaikuRAG
+from haiku.rag.store.models import SearchResult
+from haiku.rag.store.models.citation import resolve_citations
+from tests.test_multi_db import _config, _seed
+
+
+class TestExpansionRouting:
+    @pytest.mark.asyncio
+    async def test_expansion_routes_each_result_to_its_database(self, tmp_path):
+        """A federating client has no repositories of its own, so expansion has
+        to go through the database each result came from."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            results = await rag.search("cats", search_type="fts", limit=10)
+            expanded = await rag.expand_context(results)
+
+        assert {r.source for r in expanded} == {"alpha", "beta"}
+        for r in expanded:
+            assert r.source is not None
+            assert r.source in r.content
+
+
+class TestCitationSource:
+    def test_a_citation_carries_the_result_source(self):
+        result = SearchResult(
+            content="body",
+            score=0.9,
+            source="alpha",
+            chunk_id="c1",
+            document_id="d1",
+            document_uri="test://alpha/one",
+        )
+
+        [citation] = resolve_citations(["c1"], [result])
+
+        assert citation.source == "alpha"
+
+    def test_a_single_database_citation_has_no_source(self):
+        result = SearchResult(
+            content="body",
+            score=0.9,
+            chunk_id="c1",
+            document_id="d1",
+            document_uri="test://one",
+        )
+
+        [citation] = resolve_citations(["c1"], [result])
+
+        assert citation.source is None
+
+
+class TestAskAcrossDatabases:
+    @pytest.mark.asyncio
+    async def test_the_capability_searches_the_selected_databases(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            capability = create_capability(config=config, rag=rag, defer_loading=False)
+            capability.state = RAGState(sources=["alpha"])
+
+            formatted = await capability._search("cats", limit=10)
+
+        assert isinstance(formatted, str)
+        assert "alpha" in formatted
+        assert "beta document" not in formatted
+
+    @pytest.mark.asyncio
+    async def test_searching_all_databases_reaches_both(self, tmp_path):
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            capability = create_capability(config=config, rag=rag, defer_loading=False)
+            capability.state = RAGState()
+
+            formatted = await capability._search("cats", limit=10)
+
+        assert isinstance(formatted, str)
+        assert "alpha document" in formatted
+        assert "beta document" in formatted
+
+
+class TestCiteFallback:
+    @pytest.mark.asyncio
+    async def test_an_id_from_a_selected_database_resolves_with_its_source(
+        self, tmp_path
+    ):
+        """The fallback exists for a real id this run's searches did not return.
+        Across databases it looks through the selected ones and records which
+        held it."""
+        from tests.capabilities.test_capabilities import Deps, make_context
+
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(
+            config, "alpha", ["alpha document about cats", "alpha on aardvarks"]
+        )
+        await _seed(config, "beta", ["beta document about cats"])
+
+        async with HaikuRAG(config=config) as rag:
+            alpha = (await rag.clients_for(["alpha"]))[0]
+            chunks = await alpha.chunk_repository.list_all()
+            [aardvark] = [c for c in chunks if "aardvark" in c.content]
+            assert aardvark.id is not None
+
+            capability = create_capability(config=config, rag=rag, defer_loading=False)
+            deps = Deps(
+                state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
+            )
+            run = await capability.for_run(make_context(deps))
+            # The search returns the cats chunk, never the aardvark one.
+            await run._search("cats", limit=10)
+
+            await run._cite([aardvark.id])
+
+        assert run.state is not None
+        [citation] = list(run.state.citation_index.values())
+        assert citation.chunk_id == aardvark.id
+        assert citation.source == "alpha"
+
+    @pytest.mark.asyncio
+    async def test_an_id_outside_the_selected_databases_does_not_resolve(
+        self, tmp_path
+    ):
+        """A question scoped to one database must not produce a citation from
+        another: the fallback looks only where the question looked."""
+        from tests.capabilities.test_capabilities import Deps, make_context
+
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about dogs"])
+
+        async with HaikuRAG(config=config) as rag:
+            beta = (await rag.clients_for(["beta"]))[0]
+            [outside] = await beta.chunk_repository.list_all(limit=1)
+            assert outside.id is not None
+
+            capability = create_capability(config=config, rag=rag, defer_loading=False)
+            deps = Deps(
+                state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
+            )
+            run = await capability.for_run(make_context(deps))
+            await run._search("cats", limit=10)
+
+            with pytest.raises(ModelRetry):
+                await run._cite([outside.id])
+
+
+class TestFederatedEdges:
+    @pytest.mark.asyncio
+    async def test_expansion_passes_through_results_without_a_source(self, tmp_path):
+        """A caller can hand `expand_context` results it built itself. Those name
+        no database, so there is nowhere to expand them from."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        handmade = SearchResult(content="handmade", score=0.4, doc_item_refs=[])
+
+        async with HaikuRAG(config=config) as rag:
+            found = await rag.search("cats", search_type="fts", limit=10)
+            expanded = await rag.expand_context([*found, handmade])
+
+        assert "handmade" in [r.content for r in expanded]
+        scores = [r.score for r in expanded]
+        assert scores == sorted(scores, reverse=True), "merged in score order"
+
+    @pytest.mark.asyncio
+    async def test_a_chunk_without_a_document_is_not_cited(self, tmp_path):
+        """`Chunk.document_id` is optional, and a citation without a document has
+        nothing to point at."""
+        from unittest.mock import AsyncMock, patch
+
+        from haiku.rag.capabilities.rag import RAGCapability
+        from haiku.rag.store.models import Chunk
+        from tests.capabilities.test_capabilities import Deps, make_context
+
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+
+        orphan = AsyncMock()
+        orphan._federated = {}
+        orphan._source = None
+        orphan.get_chunk_by_id.return_value = Chunk(
+            id="orphan", document_id=None, content="no document"
+        )
+
+        capability = create_capability(config=config, defer_loading=False)
+        run = await capability.for_run(make_context(Deps()))
+        with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=orphan)):
+            with pytest.raises(ModelRetry):
+                await run._cite(["orphan"])