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:
parent
87367759f9
commit
9ed66e3a06
3 changed files with 111 additions and 0 deletions
|
|
@ -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]] = []
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in a new issue