From d06277408b7be5ffea47bb79fc866a7fb875f1da Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 31 Aug 2026 13:15:05 +0300 Subject: [PATCH 1/3] Break cross-database RRF rank ties by retrieval score Disjoint corpora give every database's rank-r candidate the same RRF score, and the stable sort resolved those ties to lancedb.databases declaration order, discarding the retrieval scores entirely. Ties now break on the raw retrieval score, which is uncalibrated across indexes but only ever orders candidates within one rank tier: the databases in a fusion share an embedder and ran the same search type, and it can never lift a candidate above another rank. Hybrid per-database scores are themselves rank-derived, so exact agreement still ties and keeps configured order, deterministically. The n > limit depth quota is unchanged, pending the retrieval eval. --- CHANGELOG.md | 5 +++ haiku_rag_slim/haiku/rag/client/search.py | 18 +++++++--- tests/multi_db/test_search.py | 41 +++++++++++++++++++++-- 3 files changed, 56 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1309ffc7..0233636a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,11 @@ - `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 breaks rank ties by retrieval + score instead of database declaration order. + ## [0.80.0] - 2026-08-31 ### Changed diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 9c3441a6..b241fda8 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -156,6 +156,14 @@ async def _fuse( across databases tractable: it compares query against document and does not care where a candidate came from. Without one, reciprocal rank fusion over the per-database rankings, since scores from separate indexes are not comparable. + The corpora are disjoint, so every database's rank-r candidate carries the + same RRF score; those ties break on the raw retrieval score, which is not + calibrated across indexes but only ever orders candidates within one rank + and never lifts a candidate above another rank. A hybrid score is itself + rank-derived (each database fuses its own vector and FTS rankings with + lancedb's RRFReranker), so there the tiebreak compares rank agreement, not + relevance magnitude, and databases agreeing exactly still tie, resolving + deterministically to configured order. """ owned = [ (client, chunk, score) @@ -191,12 +199,12 @@ 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): - scored.append((1.0 / (_RRF_K + rank + 1), client, chunk)) - scored.sort(key=lambda item: item[0], reverse=True) - return [(client, chunk, score) for score, client, chunk in scored[:limit]] + for rank, (chunk, score) in enumerate(candidates): + scored.append((1.0 / (_RRF_K + rank + 1), score, client, chunk)) + scored.sort(key=lambda item: (item[0], item[1]), reverse=True) + return [(client, chunk, fused) for fused, _, client, chunk in scored[:limit]] # Reciprocal rank fusion's smoothing constant, the value the literature uses. diff --git a/tests/multi_db/test_search.py b/tests/multi_db/test_search.py index 3332d9b3..de588c7e 100644 --- a/tests/multi_db/test_search.py +++ b/tests/multi_db/test_search.py @@ -527,12 +527,47 @@ class TestReciprocalRankFusion: ] @pytest.mark.asyncio - async def test_equal_scores_keep_the_configured_order(self, tmp_path): - """Every rank ties across databases, so the tiebreak decides all of it.""" - fused = await self._fuse_over(tmp_path, self._lopsided(2), 10) + async def test_ranks_tie_and_the_retrieval_score_breaks_them(self, tmp_path): + """Every rank ties across databases, so the tiebreak decides all of it, + and it must be relevance, not the order databases were configured in.""" + per_source = [self._ranked("a", 2, 0.2), self._ranked("b", 2, 0.9)] + fused = await self._fuse_over(tmp_path, per_source, 10) + + assert [source for source, _, _ in fused] == ["beta", "alpha", "beta", "alpha"] + + @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_exact_ties_keep_the_configured_order(self, tmp_path): + """Hybrid scores are rank-derived and tie exactly when databases agree, + 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_the_retrieval_score_never_overrides_the_rank(self, tmp_path): + """Raw scores are not calibrated across indexes, so a database with + inflated scores still contributes one candidate per rank.""" + fused = await self._fuse_over(tmp_path, self._lopsided(2), 10) + + assert [cid for _, cid, _ in fused] == ["a0", "b0", "a1", "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.""" From 82fe91bb7693428e5a6aa4b8b97f41fed6fc6084 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 1 Sep 2026 08:34:37 +0300 Subject: [PATCH 2/3] Order cross-database fusion by retrieval score Rank interleaving guarantees every database slots regardless of content; on domain-split collections it allocates no better than chance and costs 4.7pp recall@5 at four collections against score ordering (7.1pp at eight). Hybrid scores are each database's own vector/FTS rank agreement, which carries across databases; equal scores resolve by within-database rank, and only a tie on both falls to configured order, leaving permutation sensitivity at 0.02-0.26pp. Fused results carry the candidate's own retrieval score, so the context-expansion re-sort preserves fused order. --- CHANGELOG.md | 5 +- docs/configuration/storage.md | 2 +- haiku_rag_slim/haiku/rag/client/search.py | 24 +++++----- tests/multi_db/test_search.py | 57 +++++++++++++---------- 4 files changed, 49 insertions(+), 39 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0233636a..02a020f8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,8 +9,9 @@ ### Fixed -- Cross-database fusion without a reranker breaks rank ties by retrieval - score instead of database declaration order. +- 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. ## [0.80.0] - 2026-08-31 diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 02bfc02c..dc349596 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 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 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`. 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. diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index b241fda8..e1f8adcf 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -154,16 +154,16 @@ async def _fuse( 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 - per-database rankings, since scores from separate indexes are not comparable. - The corpora are disjoint, so every database's rank-r candidate carries the - same RRF score; those ties break on the raw retrieval score, which is not - calibrated across indexes but only ever orders candidates within one rank - and never lifts a candidate above another rank. A hybrid score is itself - rank-derived (each database fuses its own vector and FTS rankings with - lancedb's RRFReranker), so there the tiebreak compares rank agreement, not - relevance magnitude, and databases agreeing exactly still tie, resolving - deterministically to configured order. + 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. """ owned = [ (client, chunk, score) @@ -203,8 +203,8 @@ 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)) - scored.sort(key=lambda item: (item[0], item[1]), reverse=True) - return [(client, chunk, fused) for fused, _, client, chunk in scored[:limit]] + scored.sort(key=lambda item: (item[1], item[0]), reverse=True) + return [(client, chunk, score) for _, score, client, chunk in scored[:limit]] # Reciprocal rank fusion's smoothing constant, the value the literature uses. diff --git a/tests/multi_db/test_search.py b/tests/multi_db/test_search.py index de588c7e..b55329b6 100644 --- a/tests/multi_db/test_search.py +++ b/tests/multi_db/test_search.py @@ -473,9 +473,10 @@ class TestNarrowingToOneDatabase: assert results == [] -class TestReciprocalRankFusion: - """Without a reranker, scores from separate indexes are not comparable, so - fusion ranks by position. These pin what that produces.""" +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.""" @staticmethod def _ranked(source: str, count: int, top: float) -> list[tuple[Chunk, float]]: @@ -502,38 +503,46 @@ class TestReciprocalRankFusion: return [(owner.source, chunk.id, score) for owner, chunk, score in fused] @pytest.mark.asyncio - async def test_databases_interleave_by_rank(self, tmp_path): - """Each contributes its rank-1 before either contributes its rank-2.""" + 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] == [ - 1 / 61, - 1 / 61, - 1 / 62, - 1 / 62, - ] + assert [score for _, _, score in fused] == [0.9, 0.89, 0.2, 0.19] @pytest.mark.asyncio - async def test_ranks_tie_and_the_retrieval_score_breaks_them(self, tmp_path): - """Every rank ties across databases, so the tiebreak decides all of it, - and it must be relevance, not the order databases were configured in.""" - per_source = [self._ranked("a", 2, 0.2), self._ranked("b", 2, 0.9)] + async def test_score_ties_break_by_rank_within_the_database(self, tmp_path): + """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 [source for source, _, _ in fused] == ["beta", "alpha", "beta", "alpha"] + 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): @@ -561,12 +570,12 @@ class TestReciprocalRankFusion: assert [source for source, _, _ in fused] == ["alpha", "beta", "alpha", "beta"] @pytest.mark.asyncio - async def test_the_retrieval_score_never_overrides_the_rank(self, tmp_path): - """Raw scores are not calibrated across indexes, so a database with - inflated scores still contributes one candidate per rank.""" + 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", "b0", "a1", "b1"] + 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): @@ -575,8 +584,8 @@ class TestReciprocalRankFusion: assert [(source, cid) for source, cid, _ in fused] == [ ("alpha", "a0"), - ("beta", "b0"), ("alpha", "a1"), + ("alpha", "a2"), ] From 7be71ebbac4f91b31e00adfc859cc2b6fb248ee0 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 1 Sep 2026 13:45:29 +0300 Subject: [PATCH 3/3] 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