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 41a7a263cb
commit 025e042fd0
No known key found for this signature in database
6 changed files with 170 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
CHANGELOG.md
View file

@ -4,9 +4,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

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

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

@ -113,12 +113,15 @@ async def search_sources(
search_type=resolved,
filter=filter,
query_vector=query_vector,
with_vectors=query_vector is not None,
)
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 +152,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)
@ -203,10 +206,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

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

99
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,97 @@ 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_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