569947b28d
`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.
60 lines
1.9 KiB
Python
60 lines
1.9 KiB
Python
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
|