danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source 
Order cross-database fusion by cosine similarity to the query
This commit is contained in:
Yiorgis Gozadinos 2026-09-01 07:36:12 -05:00 committed by GitHub
commit f7e9535d7b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 296 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
CHANGELOG.md
View file

@ -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

14
docs/configuration/storage.md
View file

@ -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.
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.

70
haiku_rag_slim/haiku/rag/client/search.py
View file

@ -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
per-database rankings, since scores from separate indexes are not comparable.
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)
@ -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
# 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:
@ -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):
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))
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
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]

13
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -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))

216
tests/multi_db/test_search.py
View file

@ -473,9 +473,11 @@ 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 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]]:
@ -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):
"""Each contributes its rank-1 before either contributes its rank-2."""
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
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_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 [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):
"""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_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_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"),
]

20
tests/test_chunk.py
View file

@ -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