diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index e2f169da..c47d9b98 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -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 diff --git a/tests/test_search_phases.py b/tests/test_search_phases.py new file mode 100644 index 00000000..05cc5a22 --- /dev/null +++ b/tests/test_search_phases.py @@ -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