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:
parent
735489d723
commit
569947b28d
2 changed files with 129 additions and 34 deletions
|
|
@ -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
|
||||
|
|
|
|||
60
tests/test_search_phases.py
Normal file
60
tests/test_search_phases.py
Normal 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
|
||||
Loading…
Reference in a new issue