Merge pull request #593 from ggozad/fix/cross-database-fusion

Order cross-database fusion by cosine similarity to the query

CHANGELOG.md

@@ -7,6 +7,13 @@
 - `docs/configuration/storage.md` "Vector Indexing" carries the measured with/without IVF_PQ retrieval comparison; `docs/benchmarks.md` states the published numbers are measured without a vector index.
 - Re-indexing note corrected: `optimize()` adds new chunks to an existing vector index, a rebuild retrains centroids.
 
+### Fixed
+
+- Cross-database fusion without a reranker orders the union by cosine
+  similarity to the query, with within-database rank breaking ties, instead
+  of round-robin by database declaration order. Full-text-only searches order
+  by retrieval score. Fused results carry the ordering score.
+
 ## [0.80.0] - 2026-08-31
 
 ### Changed

docs/configuration/storage.md

@@ -227,7 +227,7 @@ results = await client.search("query")                     # every database
 results = await client.search("query", sources=["papers"])   # one of them
 ```
 
-Candidates are combined into one ranked list with the configured reranker, or with reciprocal rank fusion when reranking is disabled. `SearchResult.source`, `Citation.source`, and `Document.source` contain the configured database name. The name is retained when a client covers only one named database. Databases configured through `lancedb.uri` are unnamed, so their `source` is `None`.
+Candidates are combined into one ranked list with the configured reranker, or by cosine similarity to the query when reranking is disabled, with within-database rank breaking ties (full-text-only searches order by retrieval score). `SearchResult.source`, `Citation.source`, and `Document.source` contain the configured database name. The name is retained when a client covers only one named database. Databases configured through `lancedb.uri` are unnamed, so their `source` is `None`.
 
 The CLI labels results and citations only when the operation spans multiple databases. A command already narrowed with `--db-name` does not repeat the name on every result.
 
@@ -245,17 +245,13 @@ The chat document filter selects by document and database: the search is narrowe
 
 #### Ranking
 
-Reciprocal rank fusion compares positions rather than scores, so each database contributes top-ranked results even when another database has stronger matches. A reranker scores the combined candidate set directly, which has been measured to help aggregate retrieval and to hurt attribution between near-identical documents.
+Without a reranker, the fused list is ordered by cosine similarity between the query vector and each candidate. The databases in a selection share an embedder, so similarity in that one space is comparable across databases, where retrieval scores are each database's own arithmetic. Ties resolve by the candidate's rank within its own database, and configured order decides only when both tie. Similarities rarely tie exactly, so declaration order decides almost nothing: on MTRAG retrieval benchmarks, reversing it left recall unchanged in every cell. Full-text-only searches have no query vector and order by retrieval score instead.
 
-Aggregate retrieval is stronger with a reranker. In a 3,045-query evaluation over a corpus split across three databases, reranking produced retrieval MAP 0.9914, compared with 0.9918 for the same corpus in one database. Without a reranker, MAP was 0.6044, compared with 0.9798 in one database. Reranking cost grows with the number of databases because each contributes candidates.
+Results are not guaranteed to spread across databases: a database with nothing relevant to a query contributes nothing, and a strong database can fill every slot. On MTRAG retrieval benchmarks over two to eight collections, cosine fusion holds recall roughly flat as collections are added, where position-based fusion lost up to half its recall at eight.
 
-A reranker scores the combined candidates with no notion of which database each came from, so on near-identical text it can pick the wrong database's chunk, where fusion keeps them apart because each database contributes its own top-ranked result. In two nine-case acceptance runs over a synthetic corpus holding one station in two databases under near-identical names, attribution was weaker with reranking: citing the right database succeeded 5 of 9 and 6 of 9 times with a reranker, against 8 of 9 and 9 of 9 without.
+A configured reranker scores the combined candidate set directly, ignoring which database each candidate came from, and remains the strongest option: roughly 6 to 8 recall points above cosine fusion on the same benchmarks. Its cost grows with the number of databases because each contributes candidates.
 
-Configure a reranker where retrieval breadth matters, and measure it where answers have to attribute between documents that read alike.
+Image queries are vector-only and skip the reranker: the reranker interface takes a text query, and multimodal reranking applies to pictures on the candidate side, not to image queries. Their fused list is ordered by cosine similarity like any other vector search.
-
-Without a reranker, consider increasing `search.limit` with the number of databases. With three complete rankings and a limit of 5, a database may contribute only one or two results. A higher limit also sends more results to the caller and model.
-
-Image queries are vector-only and skip the reranker: there is no query text to score a document against, so candidates keep their vector ranking and fusion ranks by position.
 
 If a selected database is unavailable, the operation fails with `SourceUnavailableError`, which names that database.
 

haiku_rag_slim/haiku/rag/client/search.py

@@ -105,6 +105,11 @@ async def search_sources(
     fetch_limit = _fetch_limit(client, query, limit)
     query_vector = await _embed_query(selected[0], query, resolved)
     text = query if isinstance(query, str) else ""
+    # Embeddings are read only by cosine fusion: a reranker scores the union
+    # itself, and its 10x over-fetch would materialize them for nothing.
+    uses_cosine = query_vector is not None and (
+        not isinstance(query, str) or client.reranker is None
+    )
     per_source = await gather_all(
         *(
             c.chunk_repository.search(
@@ -113,12 +118,15 @@ async def search_sources(
                 search_type=resolved,
                 filter=filter,
                 query_vector=query_vector,
+                with_vectors=uses_cosine,
             )
             for c in selected
         )
     )
 
-    ranked = await _fuse(client, selected, query, per_source, limit)
+    ranked = await _fuse(
+        client, selected, query, per_source, limit, query_vector=query_vector
+    )
 
     results: list[SearchResult] = []
     for owner, chunk, score in ranked:
@@ -149,13 +157,21 @@ async def _fuse(
     query: "str | bytes | PILImage.Image",
     per_source: list[list[tuple[Chunk, float]]],
     limit: int,
+    query_vector: list[float] | None = None,
 ) -> list[tuple["HaikuRAG", Chunk, float]]:
     """One ranked list from several, keeping each candidate's owner.
 
     A configured reranker scores the union directly, which is what makes ranking
     across databases tractable: it compares query against document and does not
-    care where a candidate came from. Without one, reciprocal rank fusion over the
+    care where a candidate came from. Without one, the union is ordered by
-    per-database rankings, since scores from separate indexes are not comparable.
+    cosine similarity to the query vector: the databases in a selection share an
+    embedder, so similarity in that one space is the signal that is comparable
+    across databases by construction, where retrieval scores are each database's
+    own rank arithmetic. A search with no query vector (full-text) orders by the
+    retrieval score instead. In both, ties resolve by within-database rank — the
+    candidate nothing in its own database beat wins — and only a tie on both
+    falls to configured order. The returned score is the one the union was
+    ordered by, so downstream re-sorts (context expansion) preserve this order.
     """
     owned = [
         (client, chunk, score)
@@ -165,8 +181,9 @@ async def _fuse(
     if not owned:
         return []
 
-    # An image query has no text for a reranker to score against, and the check
+    # The reranker interface takes a text query, so an image query skips it,
-    # precedes `reranker`, which builds the reranker on first access.
+    # and the check precedes `reranker`, which builds the reranker on first
+    # access.
     if isinstance(query, str):
         reranker = federator.reranker
         if reranker is not None:
@@ -191,12 +208,38 @@ async def _fuse(
                 )
             return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked]
 
-    scored: list[tuple[float, HaikuRAG, Chunk]] = []
+    scored: list[tuple[float, float, HaikuRAG, Chunk]] = []
     for client, candidates in zip(clients, per_source, strict=True):
-        for rank, (chunk, _) in enumerate(candidates):
+        for rank, (chunk, score) in enumerate(candidates):
-            scored.append((1.0 / (_RRF_K + rank + 1), client, chunk))
+            scored.append((1.0 / (_RRF_K + rank + 1), score, client, chunk))
-    scored.sort(key=lambda item: item[0], reverse=True)
+
-    return [(client, chunk, score) for score, client, chunk in scored[:limit]]
+    embeddings = [chunk.embedding for _, _, _, chunk in scored]
+    if query_vector is not None and all(e is not None for e in embeddings):
+        similarities = _cosine_to(query_vector, embeddings)  # ty: ignore[invalid-argument-type]
+        scored = [
+            (rank_score, similarity, client, chunk)
+            for (rank_score, _, client, chunk), similarity in zip(
+                scored, similarities, strict=True
+            )
+        ]
+    scored.sort(key=lambda item: (item[1], item[0]), reverse=True)
+    return [(client, chunk, score) for _, score, client, chunk in scored[:limit]]
+
+
+def _cosine_to(query_vector: list[float], embeddings: list[list[float]]) -> list[float]:
+    """Cosine similarity of each embedding to the query vector.
+
+    A zero-norm vector has no direction, so its similarity is 0 rather than a
+    division error.
+    """
+    import numpy as np
+
+    query = np.asarray(query_vector, dtype=np.float32)
+    matrix = np.asarray(embeddings, dtype=np.float32)
+    norms = np.linalg.norm(matrix, axis=1) * np.linalg.norm(query)
+    with np.errstate(divide="ignore", invalid="ignore"):
+        similarities = np.where(norms > 0, matrix @ query / norms, 0.0)
+    return [float(s) for s in similarities]
 
 
 # Reciprocal rank fusion's smoothing constant, the value the literature uses.
@@ -268,10 +311,9 @@ async def _rank(
 ) -> list[tuple[Chunk, float]]:
     """Order candidates and cut them to `limit`.
 
-    An image query carries no text for a reranker to score against, so its
+    The reranker interface takes a text query, so an image query keeps the
-    candidates keep the vector ranking. Its type is checked before
+    vector ranking. Its type is checked before `client.reranker`, which builds
-    `client.reranker`, which builds the reranker on first access and loads model
+    the reranker on first access and loads model weights for a local one.
-    weights for a local one.
     """
     if not isinstance(query, str):
         return candidates[:limit]

haiku_rag_slim/haiku/rag/store/repositories/chunk.py

@@ -240,6 +240,7 @@ class ChunkRepository:
         search_type: SearchType = "hybrid",
         filter: str | None = None,
         query_vector: list[float] | None = None,
+        with_vectors: bool = False,
     ) -> list[tuple[Chunk, float]]:
         """Search for relevant chunks using the specified search method.
 
@@ -304,7 +305,7 @@ class ChunkRepository:
         if chunk_filter is not None:
             results = results.where(chunk_filter)
         results = results.limit(limit)
-        return await self._process_search_results(results)
+        return await self._process_search_results(results, with_vectors=with_vectors)
 
     async def get_by_document_id(
         self,
@@ -405,7 +406,7 @@ class ChunkRepository:
         return len(df)
 
     async def _process_search_results(
-        self, query_result: "AsyncQueryBase"
+        self, query_result: "AsyncQueryBase", with_vectors: bool = False
     ) -> list[tuple[Chunk, float]]:
         """Process search results into chunks with document info and scores."""
         import pandas as pd
@@ -456,6 +457,13 @@ class ChunkRepository:
             )
             documents_map = {str(row["id"]): row for row in doc_rows}
 
+        # The query projects no columns, so the vectors are already in the
+        # frame; only the federated fusion path reads them, so materializing
+        # per-chunk lists is gated on the caller asking.
+        vectors = (
+            df["vector"].tolist() if with_vectors and "vector" in df.columns else None
+        )
+
         chunks_with_scores = []
         for i, chunk_record in enumerate(pydantic_results):
             doc = documents_map.get(chunk_record.document_id)
@@ -468,6 +476,7 @@ class ChunkRepository:
                 document_uri=doc["uri"] if doc else None,
                 document_title=doc["title"] if doc else None,
                 document_meta=json.loads(doc.get("metadata", "{}") if doc else "{}"),
+                embedding=list(vectors[i]) if vectors is not None else None,
             )
             score = scores[i] if i < len(scores) else 1.0
             chunks_with_scores.append((chunk, score))

tests/multi_db/test_search.py

@@ -473,9 +473,11 @@ class TestNarrowingToOneDatabase:
         assert results == []
 
 
-class TestReciprocalRankFusion:
+class TestFusionWithoutAReranker:
-    """Without a reranker, scores from separate indexes are not comparable, so
+    """Without a reranker, the union is ordered by cosine similarity to the
-    fusion ranks by position. These pin what that produces."""
+    query. A search with no query vector (full-text) orders by retrieval score
+    instead; in both, ties resolve by within-database rank and only a tie on
+    both falls to configured order. These pin what that produces."""
 
     @staticmethod
     def _ranked(source: str, count: int, top: float) -> list[tuple[Chunk, float]]:
@@ -489,7 +491,7 @@ class TestReciprocalRankFusion:
         second, so score order and position order disagree."""
         return [self._ranked("a", count, 0.9), self._ranked("b", count, 0.2)]
 
-    async def _fuse_over(self, tmp_path, per_source, limit):
+    async def _fuse_over(self, tmp_path, per_source, limit, query_vector=None):
         from haiku.rag.client.search import _fuse
 
         config = _config(tmp_path, ["alpha", "beta"])
@@ -498,41 +500,217 @@ class TestReciprocalRankFusion:
         async with HaikuRAG(config=config) as rag:
             assert rag.reranker is None
             clients = await rag.clients_for(["alpha", "beta"])
-            fused = await _fuse(rag, clients, "cats", per_source, limit)
+            fused = await _fuse(
+                rag, clients, "cats", per_source, limit, query_vector=query_vector
+            )
         return [(owner.source, chunk.id, score) for owner, chunk, score in fused]
 
+    @staticmethod
+    def _embedded(
+        source: str, embeddings: list[list[float]]
+    ) -> list[tuple[Chunk, float]]:
+        """A ranking whose retrieval scores descend while the embeddings are
+        the caller's, so cosine order and score order can be made to disagree."""
+        return [
+            (
+                Chunk(id=f"{source}{i}", content=f"{source} {i}", embedding=e),
+                0.9 - i / 100,
+            )
+            for i, e in enumerate(embeddings)
+        ]
+
     @pytest.mark.asyncio
-    async def test_databases_interleave_by_rank(self, tmp_path):
+    async def test_cosine_orders_the_union(self, tmp_path):
-        """Each contributes its rank-1 before either contributes its rank-2."""
+        """With a query vector, similarity to the query decides, not the
+        databases' own scores or ranks."""
+        alpha = self._embedded("a", [[0.0, 1.0], [0.6, 0.8]])
+        beta = self._embedded("b", [[1.0, 0.0], [0.8, 0.6]])
+        fused = await self._fuse_over(
+            tmp_path, [alpha, beta], 10, query_vector=[1.0, 0.0]
+        )
+
+        assert [cid for _, cid, _ in fused] == ["b0", "b1", "a1", "a0"]
+        assert [round(score, 2) for _, _, score in fused] == [1.0, 0.8, 0.6, 0.0]
+
+    @pytest.mark.asyncio
+    async def test_cosine_ties_break_by_rank_then_configured_order(self, tmp_path):
+        """Identical embeddings tie on cosine; within-database rank decides,
+        and equal ranks fall to configured order."""
+        same = [1.0, 0.0]
+        alpha = self._embedded("a", [same, same])
+        beta = self._embedded("b", [same, same])
+        fused = await self._fuse_over(
+            tmp_path, [alpha, beta], 10, query_vector=[1.0, 0.0]
+        )
+
+        assert [cid for _, cid, _ in fused] == ["a0", "b0", "a1", "b1"]
+
+    @pytest.mark.asyncio
+    async def test_a_hybrid_search_takes_the_cosine_path_end_to_end(
+        self, tmp_path, monkeypatch
+    ):
+        """The result scores are cosines, not retrieval scores: a fusion that
+        silently loses the candidate embeddings reverts to score order and
+        returns lancedb's hybrid values, which this pins against."""
+        dim = get_config().embeddings.model.vector_dim
+        toward = [1.0] + [0.0] * (dim - 1)
+        away = [0.0, 1.0] + [0.0] * (dim - 2)
+
+        config = _config(tmp_path, ["alpha", "beta"])
+        for name, embedding in (("alpha", away), ("beta", toward)):
+            async with HaikuRAG(config=config, create=True, sources=[name]) as rag:
+                doc = DoclingDocument(name=name)
+                doc.add_text(label=DocItemLabel.TEXT, text=f"{name} cats")
+                await rag.import_document(
+                    doc,
+                    [Chunk(content=f"{name} cats", embedding=embedding, order=0)],
+                    uri=f"test://{name}",
+                )
+
+        async def embed_query(self, text):
+            return toward
+
+        monkeypatch.setattr(EmbedderWrapper, "embed_query", embed_query)
+
+        async with HaikuRAG(config=config) as rag:
+            results = await rag.search("cats", limit=2)
+
+        assert [r.source for r in results] == ["beta", "alpha"]
+        assert results[0].score == pytest.approx(1.0)
+        assert results[1].score == pytest.approx(0.0)
+
+    @pytest.mark.asyncio
+    async def test_embeddings_are_materialized_only_for_cosine_fusion(
+        self, tmp_path, monkeypatch, query_embedding
+    ):
+        """A reranker scores the union itself, so its 10x over-fetch must not
+        materialize per-chunk embeddings."""
+        from haiku.rag.store.repositories.chunk import ChunkRepository
+
+        config = _config(tmp_path, ["alpha", "beta"])
+        await _seed(config, "alpha", ["alpha document about cats"])
+        await _seed(config, "beta", ["beta document about cats"])
+
+        asked: list[bool] = []
+        search = ChunkRepository.search
+
+        async def spy(self, *args, **kwargs):
+            asked.append(kwargs.get("with_vectors", False))
+            return await search(self, *args, **kwargs)
+
+        monkeypatch.setattr(ChunkRepository, "search", spy)
+
+        async with HaikuRAG(config=config) as rag:
+            await rag.search("cats")
+        assert asked == [True, True]
+
+        asked.clear()
+        monkeypatch.setattr(HaikuRAG, "reranker", property(lambda self: StubReranker()))
+        async with HaikuRAG(config=config) as rag:
+            await rag.search("cats")
+        assert asked == [False, False]
+
+        # An image query skips the reranker branch, so it takes cosine fusion
+        # and needs vectors even with a reranker configured.
+        dim = get_config().embeddings.model.vector_dim
+
+        async def embed_image(self, image):
+            return [0.1] * dim
+
+        monkeypatch.setattr(EmbedderWrapper, "supports_images", True)
+        monkeypatch.setattr(EmbedderWrapper, "embed_image", embed_image)
+        asked.clear()
+        async with HaikuRAG(config=config) as rag:
+            await rag.search(b"\x89PNG\r\n\x1a\n")
+        assert asked == [True, True]
+
+    @pytest.mark.asyncio
+    async def test_a_candidate_without_an_embedding_disables_the_cosine(self, tmp_path):
+        """One unembedded candidate makes cosine incomparable across the union,
+        so the whole fusion keeps retrieval-score order."""
+        alpha = self._embedded("a", [[0.0, 1.0]])
+        beta = self._ranked("b", 1, 0.2)
+        fused = await self._fuse_over(
+            tmp_path, [alpha, beta], 10, query_vector=[1.0, 0.0]
+        )
+
+        assert [(cid, score) for _, cid, score in fused] == [("a0", 0.9), ("b0", 0.2)]
+
+    @pytest.mark.asyncio
+    async def test_the_score_orders_the_union(self, tmp_path):
+        """A stronger database takes consecutive slots; breadth is not
+        guaranteed."""
         fused = await self._fuse_over(tmp_path, self._lopsided(3), 10)
 
         assert [(source, cid) for source, cid, _ in fused] == [
             ("alpha", "a0"),
-            ("beta", "b0"),
             ("alpha", "a1"),
-            ("beta", "b1"),
             ("alpha", "a2"),
+            ("beta", "b0"),
+            ("beta", "b1"),
             ("beta", "b2"),
         ]
 
     @pytest.mark.asyncio
-    async def test_the_score_is_the_reciprocal_of_the_rank(self, tmp_path):
+    async def test_the_score_is_the_retrieval_score(self, tmp_path):
+        """The fused score is the candidate's own, so re-sorting downstream
+        (context expansion) preserves the fused order."""
         fused = await self._fuse_over(tmp_path, self._lopsided(2), 10)
 
-        assert [score for _, _, score in fused] == [
+        assert [score for _, _, score in fused] == [0.9, 0.89, 0.2, 0.19]
-            1 / 61,
+
-            1 / 61,
+    @pytest.mark.asyncio
-            1 / 62,
+    async def test_score_ties_break_by_rank_within_the_database(self, tmp_path):
-            1 / 62,
+        """Equal scores can sit at different ranks: rank depends on what the
+        rest of a database scored. The candidate nothing in its own database
+        beat wins the tie."""
+        per_source = [
+            [
+                (Chunk(id="a0", content="a 0"), 0.9),
+                (Chunk(id="a1", content="a 1"), 0.5),
+            ],
+            [
+                (Chunk(id="b0", content="b 0"), 0.5),
+                (Chunk(id="b1", content="b 1"), 0.3),
+            ],
+        ]
+        fused = await self._fuse_over(tmp_path, per_source, 10)
+
+        assert [cid for _, cid, _ in fused] == ["a0", "b0", "a1", "b1"]
+
+    @pytest.mark.asyncio
+    async def test_the_configured_order_does_not_matter(self, tmp_path):
+        """The same candidates fuse to the same list whichever database is
+        declared first."""
+        forward = await self._fuse_over(tmp_path, self._lopsided(3), 10)
+        (tmp_path / "swapped").mkdir()
+        backward = await self._fuse_over(
+            tmp_path / "swapped",
+            [self._ranked("b", 3, 0.2), self._ranked("a", 3, 0.9)],
+            10,
+        )
+
+        assert [(cid, score) for _, cid, score in forward] == [
+            (cid, score) for _, cid, score in backward
         ]
 
     @pytest.mark.asyncio
-    async def test_equal_scores_keep_the_configured_order(self, tmp_path):
+    async def test_exact_ties_keep_the_configured_order(self, tmp_path):
-        """Every rank ties across databases, so the tiebreak decides all of it."""
+        """Hybrid scores are rank-derived and tie exactly when databases agree,
-        fused = await self._fuse_over(tmp_path, self._lopsided(2), 10)
+        so a genuine tie must still resolve deterministically."""
+        per_source = [self._ranked("a", 2, 0.9), self._ranked("b", 2, 0.9)]
+        fused = await self._fuse_over(tmp_path, per_source, 10)
 
         assert [source for source, _, _ in fused] == ["alpha", "beta", "alpha", "beta"]
 
+    @pytest.mark.asyncio
+    async def test_rank_never_overrides_the_score(self, tmp_path):
+        """A database's rank-2 with a higher score precedes another's rank-0:
+        allocation is content-driven, not round-robin."""
+        fused = await self._fuse_over(tmp_path, self._lopsided(2), 10)
+
+        assert [cid for _, cid, _ in fused] == ["a0", "a1", "b0", "b1"]
+
     @pytest.mark.asyncio
     async def test_the_limit_cuts_the_fused_list(self, tmp_path):
         """Each database was asked for enough to fill the window on its own."""
@@ -540,8 +718,8 @@ class TestReciprocalRankFusion:
 
         assert [(source, cid) for source, cid, _ in fused] == [
             ("alpha", "a0"),
-            ("beta", "b0"),
             ("alpha", "a1"),
+            ("alpha", "a2"),
         ]
 
 

tests/test_chunk.py

@@ -676,6 +676,26 @@ async def test_fts_search_does_not_warn_on_an_empty_table(temp_db_path):
         assert not records
 
 
+async def test_search_populates_embeddings_only_when_asked(temp_db_path):
+    """Vectors ride the result frame either way; the per-chunk lists are
+    materialized only for the caller that reads them (federated fusion)."""
+    async with HaikuRAG(
+        db_path=temp_db_path, config=get_config(), create=True
+    ) as client:
+        await _import_one(client)
+        await client.store.vacuum(retention_seconds=0)
+
+        plain = await client.chunk_repository.search("gardens", search_type="fts")
+        with_vectors = await client.chunk_repository.search(
+            "gardens", search_type="fts", with_vectors=True
+        )
+
+    assert plain and all(chunk.embedding is None for chunk, _ in plain)
+    assert with_vectors and all(
+        chunk.embedding is not None for chunk, _ in with_vectors
+    )
+
+
 @pytest.mark.vcr()
 async def test_chunk_repository_get_by_id_and_list_all_pagination(
     qa_corpus: list[dict[str, str]], temp_db_path