Separate fetching from ranking in search

`search` fetched, reranked and truncated in one pass, with the reranker's
over-fetch and the reranking itself interleaved in the same branch. Searching
several databases needs to fuse their candidates before anything is ranked, so
the phases have to be separable.

`_fetch` returns one database's candidates, over-fetching only when a reranker
will re-order them. `_rank` orders and cuts them, leaving an image query's vector
ranking alone since there is no text for a reranker to score against. The
over-fetch multiplier is named rather than a literal 10 at the point of use.

Both check the query type before reading `client.reranker`, which is a
cached_property that builds the reranker on first access and loads model weights
for a local one. An image query never used it and must not start.

No behaviour change: the same suite passes, and the search outputs digest
identically to before.
This commit is contained in:
Yiorgis Gozadinos 2026-08-19 12:08:07 +03:00
parent 735489d723
commit 569947b28d
No known key found for this signature in database
2 changed files with 129 additions and 34 deletions

View file

@ -38,40 +38,8 @@ async def search(
if limit is None:
limit = client._config.search.limit
if isinstance(query, str):
if search_type is None:
search_type = "hybrid"
reranker = client.reranker
if reranker is None:
chunk_results = await client.chunk_repository.search(
query, limit, search_type, filter
)
else:
search_limit = limit * 10
raw_results = await client.chunk_repository.search(
query, search_limit, search_type, filter
)
chunks = [chunk for chunk, _ in raw_results]
if client._config.reranking.multimodal:
await _attach_picture_data(client, chunks)
chunk_results = await reranker.rerank(query, chunks, top_n=limit)
else:
embedder = client.embedder
if not embedder.supports_images:
raise ValueError(
"Image queries require a multimodal embedder. Set "
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere "
"model."
)
query_vector = await embedder.embed_image(query)
chunk_results = await client.chunk_repository.search(
query="",
limit=limit,
filter=filter,
query_vector=query_vector,
)
candidates = await _fetch(client, query, limit, search_type, filter)
chunk_results = await _rank(client, query, candidates, limit)
results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results]
results = _dedup_picture_chunks(results)
@ -82,6 +50,73 @@ async def search(
return results
# Candidates per requested result when a reranker will re-order them.
_RERANK_OVERFETCH = 10
async def _fetch(
client: "HaikuRAG",
query: "str | bytes | PILImage.Image",
limit: int,
search_type: SearchType | None,
filter: str | None,
) -> list[tuple[Chunk, float]]:
"""Candidates from one database, ranked by that database.
Over-fetches when a reranker will re-order them. Separate from `_rank` so a
caller searching several databases can fuse their candidates before anything
is ranked or enriched.
"""
if isinstance(query, str):
if search_type is None:
search_type = "hybrid"
fetch_limit = limit * _RERANK_OVERFETCH if client.reranker else limit
return await client.chunk_repository.search(
query, fetch_limit, search_type, filter
)
embedder = client.embedder
if not embedder.supports_images:
raise ValueError(
"Image queries require a multimodal embedder. Set "
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere "
"model."
)
query_vector = await embedder.embed_image(query)
return await client.chunk_repository.search(
query="",
limit=limit,
filter=filter,
query_vector=query_vector,
)
async def _rank(
client: "HaikuRAG",
query: "str | bytes | PILImage.Image",
candidates: list[tuple[Chunk, float]],
limit: int,
) -> 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.
"""
if not isinstance(query, str):
return candidates[:limit]
reranker = client.reranker
if reranker is None:
return candidates[:limit]
chunks = [chunk for chunk, _ in candidates]
if client._config.reranking.multimodal:
await _attach_picture_data(client, chunks)
return await reranker.rerank(query, chunks, top_n=limit)
async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None:
"""Attach picture bytes to synthetic picture chunks in-place, so a
multimodal reranker can score the pixels instead of just the chunk's

View file

@ -0,0 +1,60 @@
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.client.search import _fetch, _rank
from haiku.rag.store.models import Chunk
@pytest.fixture
def exploding_reranker(monkeypatch):
"""A reranker that cannot be built, so any access fails loudly."""
def boom(self):
raise AssertionError("reranker built for a query that cannot use it")
monkeypatch.setattr(HaikuRAG, "reranker", property(boom), raising=True)
@pytest.mark.asyncio
async def test_image_ranking_never_builds_the_reranker(
temp_db_path, exploding_reranker
):
"""Local rerankers load model weights on construction, so an image query,
which has no text to score against, must not touch one."""
async with HaikuRAG(temp_db_path, create=True) as rag:
candidates = [
(Chunk(id="a", document_id="d", content="one"), 0.9),
(Chunk(id="b", document_id="d", content="two"), 0.8),
]
ranked = await _rank(rag, b"image-bytes", candidates, limit=1)
assert [c.id for c, _ in ranked] == ["a"]
@pytest.mark.asyncio
async def test_image_fetch_never_builds_the_reranker(
temp_db_path, exploding_reranker, monkeypatch
):
async with HaikuRAG(temp_db_path, create=True) as rag:
seen = {}
async def fake_search(
query, limit, search_type="hybrid", filter=None, query_vector=None
):
seen["limit"] = limit
return []
async def fake_embed_image(self, image):
return [0.1] * 8
monkeypatch.setattr(rag.chunk_repository, "search", fake_search)
monkeypatch.setattr(type(rag.embedder), "embed_image", fake_embed_image)
monkeypatch.setattr(
type(rag.embedder), "supports_images", property(lambda self: True)
)
await _fetch(rag, b"image-bytes", 5, None, None)
# No over-fetch: nothing will re-rank these.
assert seen["limit"] == 5