Cache reranker on the client instead of rebuilding per search
This commit is contained in:
parent
c0c83d8037
commit
37a78a4e9b
4 changed files with 44 additions and 2 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue