Search one database the ordinary way, however it was selected

`sources=["alpha"]` on a client covering a set went through fusion, which
scores position: a result the database ranked at 0.6549 was reported as
1/(60+rank). Embedding also moved ahead of the repository, so a filter matching
no document embedded the query anyway.

A selection of one now runs the single-database search. Fusion reconciles
rankings from separate indexes, and one ranking has nothing to reconcile.

A reranker that returns chunks it built rather than the ones it was given loses
which database each came from, since ownership is by identity. That is named
now instead of surfacing as a KeyError, and stated on `RerankerBase._rerank`.
This commit is contained in:
Yiorgis Gozadinos 2026-08-27 14:40:57 +03:00
parent 87367759f9
commit 9ed66e3a06
No known key found for this signature in database
3 changed files with 111 additions and 0 deletions

View file

@ -91,6 +91,13 @@ async def search_sources(
if not names:
return []
selected = await client.clients_for(names)
if len(selected) == 1:
# One database is an ordinary search, whatever the client covers: fusion
# would replace its hybrid scores with ranks, and embedding up front
# would embed for a filter the repository can see matches nothing.
return await selected[0].search(
query, limit, search_type, filter, include_images
)
resolved = _resolved_search_type(query, search_type)
if resolved != "fts":
client._require_one_embedder(selected)
@ -178,7 +185,14 @@ async def _fuse(
)
)
reranked = await reranker.rerank(query, chunks, top_n=limit)
# Identity, since chunk ids repeat between copies of a database.
owner_of = {id(chunk): client for client, chunk, _ in owned}
if any(id(chunk) not in owner_of for chunk, _ in reranked):
raise ValueError(
f"{type(reranker).__name__} returned chunks that are not the "
"ones it was given, so the database each came from is lost; "
"a reranker must return objects from the list passed to it"
)
return [(owner_of[id(chunk)], chunk, score) for chunk, score in reranked]
scored: list[tuple[float, HaikuRAG, Chunk]] = []

View file

@ -14,6 +14,12 @@ class RerankerBase:
async def _rerank(
self, query: str, chunks: list[Chunk], top_n: int = 10
) -> list[tuple[Chunk, float]]:
"""Score and order `chunks`, returning the top `top_n`.
Return objects taken from `chunks`, not copies: searching several
databases maps a scored chunk back to the one holding it by identity,
because chunk ids repeat between copies of a database.
"""
raise NotImplementedError(
"Reranker is an abstract class. Please implement the _rerank method in a subclass."
)

View file

@ -333,6 +333,97 @@ class TestOneReranker:
assert closes == [1], f"closed {len(closes)} times"
class TestNarrowingToOneDatabase:
"""A selection of one is an ordinary search. Fusion exists to reconcile
rankings from separate indexes, and there is nothing to reconcile."""
@pytest.mark.asyncio
async def test_narrowing_keeps_the_database_s_own_scores(self, tmp_path):
"""RRF scores position, so fusing one ranking would report 1/(60+rank)
where the database reported a hybrid score."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats", "alpha on dogs"])
await _seed(config, "beta", ["beta document about cats"])
async with HaikuRAG(config=config) as covering:
narrowed = await covering.search(
"cats", search_type="fts", sources=["alpha"]
)
async with HaikuRAG(config=config, sources=["alpha"]) as one:
native = await one.search("cats", search_type="fts")
assert [r.chunk_id for r in narrowed] == [r.chunk_id for r in native]
assert [r.score for r in narrowed] == [r.score for r in native]
assert all(r.source == "alpha" for r in narrowed)
@pytest.mark.asyncio
async def test_narrowing_does_not_embed_for_a_filter_matching_nothing(
self, tmp_path, monkeypatch
):
"""One database embeds inside the repository, which returns early when
the filter matches no document. Fusing would have embedded first."""
from haiku.rag.embeddings import EmbedderWrapper
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
def explode(self, query):
raise AssertionError("embedded a query no document could match")
monkeypatch.setattr(EmbedderWrapper, "embed_query", explode)
async with HaikuRAG(config=config) as covering:
results = await covering.search(
"cats", filter="uri = 'test://nothing'", sources=["alpha"]
)
assert results == []
class TestFusingWhatARerankerReturns:
@pytest.mark.asyncio
async def test_a_reranker_returning_copies_is_named(self, tmp_path):
"""Candidates are mapped back to their database by identity, because
chunk ids repeat between copies of one. A reranker that rebuilds its
chunks loses that, and saying so beats a KeyError."""
from haiku.rag.client.search import _fuse
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
class Rebuilds:
async def rerank(self, query, chunks, top_n=10):
return [(chunk.model_copy(), 1.0) for chunk in chunks[:top_n]]
async with HaikuRAG(config=config) as rag:
clients = await rag.clients_for(["alpha", "beta"])
rag.__dict__["reranker"] = Rebuilds()
per_source = [
await c.chunk_repository.search("cats", 5, "fts") for c in clients
]
with pytest.raises(ValueError, match="objects from the list"):
await _fuse(rag, clients, "cats", per_source, 5)
class TestComparingEmbedders:
@pytest.mark.asyncio
async def test_a_database_recording_no_embedder_is_not_compared(self, tmp_path):
"""A database whose settings never recorded one cannot disagree with a
database that did, so there is nothing to reject."""
config = _config(tmp_path, ["alpha", "beta"])
await _seed(config, "alpha", ["alpha document about cats"])
await _seed(config, "beta", ["beta document about cats"])
async with HaikuRAG(config=config) as rag:
alpha, beta = await rag.clients_for(["alpha", "beta"])
beta.store.stored_embedding = None
rag._require_one_embedder([alpha, beta])
class TestOneNamedDatabase:
@pytest.mark.asyncio
async def test_a_single_named_database_keeps_its_name(self, tmp_path):