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

Order cross-database fusion by cosine similarity to the query

Browse source 
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.
This commit is contained in:
Yiorgis Gozadinos 2026-09-01 13:45:29 +03:00
parent 82fe91bb76
commit 7be71ebbac
No known key found for this signature in database
6 changed files with 230 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
CHANGELOG.md
View file

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

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

68
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,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]

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

144
tests/multi_db/test_search.py
View file

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

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