From 7be71ebbac4f91b31e00adfc859cc2b6fb248ee0 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 1 Sep 2026 13:45:29 +0300 Subject: [PATCH] Order cross-database fusion by cosine similarity to the query Retrieval scores are each database's own rank arithmetic; the databases in a selection share an embedder, so similarity in that one space is the signal comparable across databases by construction. Measured product to product against score ordering: +8.3 to +16.6pp recall@5 across five cells on two corpora, flat in collection count and corpus shape where score ordering dips with both, closing roughly 60% of the gap to a reranker; order-sensitivity residual 0.00pp in every cell. Exact ties collapse from 51-81% of candidates to under 1%. Full-text-only searches keep retrieval-score order, having no query vector. The vector column already travels with every search result, so the similarity costs no additional transfer; per-chunk embeddings are materialized only for the federated path that reads them. --- CHANGELOG.md | 7 +- docs/configuration/storage.md | 14 +- haiku_rag_slim/haiku/rag/client/search.py | 68 ++++++--- .../haiku/rag/store/repositories/chunk.py | 13 +- tests/multi_db/test_search.py | 144 +++++++++++++++++- tests/test_chunk.py | 20 +++ 6 files changed, 230 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 02a020f8..8158b094 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,9 +9,10 @@ ### Fixed -- Cross-database fusion without a reranker orders the union by retrieval - score, with within-database rank breaking ties, instead of round-robin by - database declaration order. Fused results carry the retrieval score. +- 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 diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index dc349596..fdd06e42 100644 --- a/docs/configuration/storage.md +++ b/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 by retrieval score when reranking is disabled, with within-database rank breaking score ties. `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. - -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. +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. If a selected database is unavailable, the operation fails with `SourceUnavailableError`, which names that database. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index e1f8adcf..acdf9ce7 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/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,21 +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, the union is ordered by the - raw retrieval score. Scores from separate indexes are not calibrated, but a - hybrid score is each database's own rank agreement (lancedb fuses vector and - FTS with RRF inside the database), which carries across databases; rank - interleaving instead guarantees every database slots regardless of content, - which on domain-split collections allocates no better than chance. Equal - scores 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 candidate's own, so downstream re-sorts (context - expansion) preserve this order. + care where a candidate came from. Without one, the union is ordered by + 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) @@ -173,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 - # precedes `reranker`, which builds the reranker on first access. + # The reranker interface takes a text query, so an image query skips it, + # and the check precedes `reranker`, which builds the reranker on first + # access. if isinstance(query, str): reranker = federator.reranker if reranker is not None: @@ -203,10 +212,36 @@ async def _fuse( for client, candidates in zip(clients, per_source, strict=True): for rank, (chunk, score) in enumerate(candidates): scored.append((1.0 / (_RRF_K + rank + 1), score, client, chunk)) + + 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. _RRF_K = 60 @@ -276,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 - candidates keep the vector ranking. Its type is checked before - `client.reranker`, which builds the reranker on first access and loads model - weights for a local one. + The reranker interface takes a text query, so an image query keeps the + vector ranking. Its type is checked before `client.reranker`, which builds + the reranker on first access and loads model weights for a local one. """ if not isinstance(query, str): return candidates[:limit] diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index 3cac1864..f78d6adf 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/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)) diff --git a/tests/multi_db/test_search.py b/tests/multi_db/test_search.py index b55329b6..00ace6a5 100644 --- a/tests/multi_db/test_search.py +++ b/tests/multi_db/test_search.py @@ -474,9 +474,10 @@ class TestNarrowingToOneDatabase: class TestFusionWithoutAReranker: - """Without a reranker, the union is ordered by retrieval score; score ties - resolve by within-database rank, and only a tie on both falls to configured - order. These pin what that produces.""" + """Without a reranker, the union is ordered by cosine similarity to the + 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]]: @@ -490,7 +491,7 @@ class TestFusionWithoutAReranker: 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"]) @@ -499,9 +500,142 @@ class TestFusionWithoutAReranker: 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_cosine_orders_the_union(self, tmp_path): + """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 diff --git a/tests/test_chunk.py b/tests/test_chunk.py index be15b96d..374816cb 100644 --- a/tests/test_chunk.py +++ b/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