diff --git a/CHANGELOG.md b/CHANGELOG.md index 53c47954..29337ca3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,7 @@ ### Changed - A successful DELETE job auto-prunes dead jobs with the same `(source_id, uri)`. New `JobRepo.prune_dead(source_id, uri)`. +- Reranker built once per `HaikuRAG` client instead of per `search()` call. ## [0.50.0] - 2026-05-27 diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 13b7643e..6af61629 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -7,6 +7,7 @@ import tempfile from collections.abc import AsyncGenerator from datetime import datetime from enum import Enum +from functools import cached_property from pathlib import Path from typing import TYPE_CHECKING, overload from urllib.parse import urlparse @@ -31,6 +32,7 @@ if TYPE_CHECKING: from PIL import Image as PILImage from haiku.rag.ingester.sources.base import Source + from haiku.rag.reranking.base import RerankerBase from haiku.rag.sandbox import AnalysisResult from haiku.rag.store.models.citation import Citation @@ -87,6 +89,15 @@ class HaikuRAG: """Whether the client is in read-only mode.""" return self.store.is_read_only + @cached_property + def reranker(self) -> "RerankerBase | None": + """The configured reranker, built once and reused across searches. + + None when reranking is disabled. Local rerankers load model weights on + construction, so building per search would reload them on every query. + """ + return get_reranker(config=self._config) + async def __aenter__(self): """Async context manager entry — initializes store and repositories.""" self.store = Store( diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 8cb2ab50..e441f51a 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -1,7 +1,6 @@ import base64 from typing import TYPE_CHECKING -from haiku.rag.reranking import get_reranker from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX @@ -42,7 +41,7 @@ async def search( if search_type is None: search_type = "hybrid" - reranker = get_reranker(config=client._config) + reranker = client.reranker if reranker is None: chunk_results = await client.chunk_repository.search( diff --git a/tests/test_search.py b/tests/test_search.py index 6ec7b827..a8e109af 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -395,6 +395,37 @@ async def test_search_with_bytes_query_uses_multimodal_embedder( assert received_kwargs["query"] == "" +@pytest.mark.asyncio +async def test_reranker_built_once_across_searches(temp_db_path, monkeypatch): + """The reranker is constructed once per client and reused across searches, + rather than rebuilt (reloading model weights) on every query.""" + from haiku.rag.store.models.chunk import Chunk + + build_count = 0 + + class StubReranker: + async def rerank(self, query, chunks, top_n): + return [(chunk, 1.0) for chunk in chunks][:top_n] + + def fake_get_reranker(config): + nonlocal build_count + build_count += 1 + return StubReranker() + + monkeypatch.setattr("haiku.rag.client.get_reranker", fake_get_reranker) + + async def fake_chunk_search(query, limit, search_type, filter): + return [(Chunk(content="x", metadata={}), 0.5)] + + async with HaikuRAG(temp_db_path, create=True) as rag: + rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign] + await rag.search("first", include_images=False) + await rag.search("second", include_images=False) + await rag.search("third", include_images=False) + + assert build_count == 1 + + @pytest.mark.asyncio async def test_search_with_pil_image_works_like_bytes(temp_db_path, monkeypatch): from PIL import Image as PILImageModule