haiku.rag/tests/test_search_phases.py
Yiorgis Gozadinos c284d86885
Resolve the search type once per search
An image query has no text to match, so it is vector-only whatever the
caller asked for, and full-text search embeds nothing, so it needs no
agreement on embedders.
2026-08-24 15:28:21 +03:00

88 lines
2.8 KiB
Python

import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.client.search import _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 rag.search(b"image-bytes", limit=5)
# No over-fetch: nothing will re-rank these.
assert seen["limit"] == 5
@pytest.mark.asyncio
async def test_an_image_query_ignores_a_full_text_search_type(
temp_db_path, monkeypatch
):
"""An image has no text to match, so it searches by vector whatever the
caller asked for."""
async with HaikuRAG(temp_db_path, create=True) as rag:
seen = {}
async def fake_search(query, limit, search_type, filter, query_vector):
seen.update({"search_type": search_type, "query_vector": query_vector})
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 rag.search(b"image-bytes", search_type="fts")
assert seen["search_type"] == "vector"
assert seen["query_vector"] == [0.1] * 8