Expand context through the database a result came from

Expansion branched on whether the client covered a set, so the
single-database half reached for repositories through a facade that may
have none. `expand_sources` groups results by database and hands each
group the session that owns it; `expand_context` and `visualize_chunk`
take that session, so `visualize_chunk` stops narrowing to one database
and discarding the result.

Inline `_fetch`, a pass-through to the chunk repository.

haiku_rag_slim/haiku/rag/client/__init__.py

@@ -828,9 +828,13 @@ class HaikuRAG:
         self,
         search_results: list[SearchResult],
     ) -> list[SearchResult]:
-        from haiku.rag.client.search import expand_context
+        from haiku.rag.client.search import expand_context, expand_sources
 
-        return await expand_context(self, search_results)
+        if isinstance(self._session, FederatedSession):
+            return await expand_sources(self._session, search_results)
+        return await expand_context(
+            self._single_session("expand_context"), search_results
+        )
 
     async def ask(
         self,
@@ -862,9 +866,9 @@ class HaikuRAG:
     ) -> list:
         from haiku.rag.client.search import visualize_chunk
 
-        self._single_session("visualize_chunk")
+        return await visualize_chunk(
-
+            self._single_session("visualize_chunk"), chunk, refs, expand
-        return await visualize_chunk(self, chunk, refs, expand)
+        )
 
     async def rebuild_database(
         self, mode: RebuildMode = RebuildMode.FULL

haiku_rag_slim/haiku/rag/client/search.py

@@ -10,6 +10,7 @@ if TYPE_CHECKING:
     from PIL import Image as PILImage
 
     from haiku.rag.client import HaikuRAG
+    from haiku.rag.client.session import FederatedSession, SingleDatabaseSession
 
 
 async def search(
@@ -45,13 +46,12 @@ async def search(
     query_vector = (
         None if isinstance(query, str) else await _embed_query(client, query, resolved)
     )
-    candidates = await _fetch(
+    candidates = await client.chunk_repository.search(
-        client,
+        query=query if isinstance(query, str) else "",
-        query,
+        limit=_fetch_limit(client, query, limit),
-        _fetch_limit(client, query, limit),
+        search_type=resolved,
-        resolved,
+        filter=filter,
-        filter,
+        query_vector=query_vector,
-        query_vector,
     )
     chunk_results = await _rank(client, query, candidates, limit)
 
@@ -92,12 +92,20 @@ async def search_sources(
 
     # One over-fetch decision, one query vector, and one reranker, for the whole
     # set. The databases in a selection share an embedder, so the vector is the
-    # same wherever it is computed.
+    # same wherever it is computed, and deciding the over-fetch per database would
+    # have each consult its own reranker.
     fetch_limit = _fetch_limit(client, query, limit)
     query_vector = await _embed_query(selected[0], query, resolved)
+    text = query if isinstance(query, str) else ""
     per_source = await asyncio.gather(
         *(
-            _fetch(c, query, fetch_limit, resolved, filter, query_vector)
+            c.chunk_repository.search(
+                query=text,
+                limit=fetch_limit,
+                search_type=resolved,
+                filter=filter,
+                query_vector=query_vector,
+            )
             for c in selected
         )
     )
@@ -238,29 +246,6 @@ async def _embed_query(
     return await embedder.embed_image(query)
 
 
-async def _fetch(
-    client: "HaikuRAG",
-    query: "str | bytes | PILImage.Image",
-    limit: int,
-    search_type: SearchType,
-    filter: str | None,
-    query_vector: list[float] | None,
-) -> list[tuple[Chunk, float]]:
-    """Candidates from one database, ranked by that database.
-
-    `limit` is how many to fetch, already including any over-fetch the caller
-    wants. Deciding that here would have each database consult its own reranker,
-    and a local reranker loads model weights per instance.
-    """
-    return await client.chunk_repository.search(
-        query=query if isinstance(query, str) else "",
-        limit=limit,
-        search_type=search_type,
-        filter=filter,
-        query_vector=query_vector,
-    )
-
-
 async def _rank(
     client: "HaikuRAG",
     query: "str | bytes | PILImage.Image",
@@ -416,22 +401,15 @@ async def _populate_image_data(client: "HaikuRAG", results: list[SearchResult])
             r.picture_captions = captions
 
 
-async def expand_context(
+async def expand_sources(
-    client: "HaikuRAG",
+    federated: "FederatedSession",
     search_results: list[SearchResult],
 ) -> list[SearchResult]:
-    """Expand search results with surrounding content from the document.
+    """Expand results drawn from several databases, each through its own.
 
-    Uses the document_items table for section-bounded expansion.
+    A result naming no database passes through unexpanded: it cannot be placed,
-    See haiku.rag.context for the algorithm description.
+    which is the case for results a caller built by hand.
-
-    Results without doc_item_refs pass through unexpanded. This happens when
-    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.covers_multiple:
     by_source: dict[str, list[SearchResult]] = {}
     unsourced: list[SearchResult] = []
     for result in search_results:
@@ -439,17 +417,17 @@ async def expand_context(
             by_source.setdefault(result.source, []).append(result)
         else:
             unsourced.append(result)
-        owners = await client.clients_for(list(by_source))
+    names = list(by_source)
+    sessions = await federated.sessions_for(names)
     expanded_groups = await asyncio.gather(
         *(
-                expand_context(owner, by_source[owner.source])
+            expand_context(session, by_source[name])
-                for owner in owners
+            for name, session in zip(names, sessions, strict=True)
-                if owner.source
         )
     )
     merged = unsourced + [r for group in expanded_groups for r in group]
-        # Grouping by database must not become the tiebreak: fused scores tie
+    # Grouping by database must not become the tiebreak: fused scores tie often,
-        # often, so equal scores keep the order they were fused in.
+    # so equal scores keep the order they were fused in.
     arrival = {
         result.chunk_id: rank
         for rank, result in enumerate(search_results)
@@ -469,9 +447,23 @@ async def expand_context(
     merged.sort(key=lambda r: (-r.score, fused_rank(r)))
     return merged
 
+
+async def expand_context(
+    session: "SingleDatabaseSession",
+    search_results: list[SearchResult],
+) -> list[SearchResult]:
+    """Expand search results with surrounding content from the document.
+
+    Uses the document_items table for section-bounded expansion.
+    See haiku.rag.context for the algorithm description.
+
+    Results without doc_item_refs pass through unexpanded. This happens when
+    chunks were created without docling metadata (e.g., custom chunks passed
+    to import_document).
+    """
     from haiku.rag.context import expand_with_items, window_for
 
-    max_chars = client._config.search.max_context_chars
+    max_chars = session.config.search.max_context_chars
 
     # Group by document_id for efficient processing
     document_groups: dict[str | None, list[SearchResult]] = {}
@@ -487,7 +479,7 @@ async def expand_context(
         for doc_id, doc_results in document_groups.items()
         if doc_id is not None and any(r.doc_item_refs for r in doc_results)
     }
-    repo = client.document_item_repository
+    repo = session.document_item_repository
     positions_by_document = await repo.resolve_refs_grouped(
         {
             doc_id: [ref for r in doc_results for ref in r.doc_item_refs]
@@ -526,7 +518,7 @@ async def expand_context(
 
 
 async def visualize_chunk(
-    client: "HaikuRAG",
+    session: "SingleDatabaseSession",
     chunk: "Chunk | Sequence[Chunk]",
     refs: list[str] | None = None,
     expand: bool = True,
@@ -561,7 +553,7 @@ async def visualize_chunk(
         return []
     chunks = [c for c in chunks if c.document_id == document_id]
 
-    doc = await client.document_repository.get_docling_data(document_id)
+    doc = await session.document_repository.get_docling_data(document_id)
     if not doc:
         return []
 
@@ -591,7 +583,7 @@ async def visualize_chunk(
             if (meta := c.get_chunk_metadata()).doc_item_refs
         ]
         if search_results:
-            expanded = await expand_context(client, search_results)
+            expanded = await expand_context(session, search_results)
             all_refs = []
             for result in expanded:
                 all_refs.extend(r for r in result.doc_item_refs if r not in all_refs)
@@ -622,7 +614,7 @@ async def visualize_chunk(
         boxes_by_page[bbox.page_no].append((bbox, is_matched))
 
     # Load only the needed page images
-    pages_doc = await client.document_repository.get_pages_data(document_id)
+    pages_doc = await session.document_repository.get_pages_data(document_id)
     if not pages_doc:
         return []
     page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))

tests/test_multi_db_ask.py

@@ -69,6 +69,25 @@ class TestExpansionRouting:
         assert "cats also hunt" in expanded[0].content, "expansion did not run"
         assert expanded[0].source == "alpha"
 
+    @pytest.mark.asyncio
+    async def test_a_federated_result_is_expanded_by_its_own_database(self, tmp_path):
+        """Routing is not enough: each result has to come back carrying the
+        neighbours of the database it was expanded through, and only those."""
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed_expandable(
+            config, "alpha", ["cats sleep often", "alpha follows on"]
+        )
+        await _seed_expandable(config, "beta", ["cats also hunt", "beta follows on"])
+
+        async with HaikuRAG(config=config) as rag:
+            results = await rag.search("cats", search_type="fts", limit=10)
+            expanded = await rag.expand_context(results)
+
+        content = {r.source: r.content for r in expanded}
+        assert "alpha follows on" in content["alpha"]
+        assert "beta follows on" not in content["alpha"]
+        assert "beta follows on" in content["beta"]
+
     @pytest.mark.asyncio
     async def test_expansion_keeps_tied_results_in_fused_order(self, tmp_path):
         """Fused scores tie often, so grouping by database must not reorder

tests/test_search.py

@@ -549,11 +549,9 @@ def test_dedup_does_not_collapse_across_documents():
 @pytest.mark.asyncio
 async def test_expand_context_passes_through_results_without_document(temp_db_path):
     """A result with no document_id can't be expanded; it is returned as-is."""
-    from haiku.rag.client.search import expand_context
-
     async with HaikuRAG(temp_db_path, create=True) as rag:
         orphan = SearchResult(content="loose text", score=0.5, chunk_id="c1")
-        assert await expand_context(rag, [orphan]) == [orphan]
+        assert await rag.expand_context([orphan]) == [orphan]
 
 
 @pytest.mark.asyncio