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
|
### Changed
|
||||||
|
|
||||||
- A successful DELETE job auto-prunes dead jobs with the same `(source_id, uri)`. New `JobRepo.prune_dead(source_id, uri)`.
|
- 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
|
## [0.50.0] - 2026-05-27
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -7,6 +7,7 @@ import tempfile
|
||||||
from collections.abc import AsyncGenerator
|
from collections.abc import AsyncGenerator
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
|
from functools import cached_property
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from typing import TYPE_CHECKING, overload
|
from typing import TYPE_CHECKING, overload
|
||||||
from urllib.parse import urlparse
|
from urllib.parse import urlparse
|
||||||
|
|
@ -31,6 +32,7 @@ if TYPE_CHECKING:
|
||||||
from PIL import Image as PILImage
|
from PIL import Image as PILImage
|
||||||
|
|
||||||
from haiku.rag.ingester.sources.base import Source
|
from haiku.rag.ingester.sources.base import Source
|
||||||
|
from haiku.rag.reranking.base import RerankerBase
|
||||||
from haiku.rag.sandbox import AnalysisResult
|
from haiku.rag.sandbox import AnalysisResult
|
||||||
from haiku.rag.store.models.citation import Citation
|
from haiku.rag.store.models.citation import Citation
|
||||||
|
|
||||||
|
|
@ -87,6 +89,15 @@ class HaikuRAG:
|
||||||
"""Whether the client is in read-only mode."""
|
"""Whether the client is in read-only mode."""
|
||||||
return self.store.is_read_only
|
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 def __aenter__(self):
|
||||||
"""Async context manager entry — initializes store and repositories."""
|
"""Async context manager entry — initializes store and repositories."""
|
||||||
self.store = Store(
|
self.store = Store(
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,6 @@
|
||||||
import base64
|
import base64
|
||||||
from typing import TYPE_CHECKING
|
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.chunk import Chunk, SearchResult, SearchType
|
||||||
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX
|
||||||
|
|
||||||
|
|
@ -42,7 +41,7 @@ async def search(
|
||||||
if search_type is None:
|
if search_type is None:
|
||||||
search_type = "hybrid"
|
search_type = "hybrid"
|
||||||
|
|
||||||
reranker = get_reranker(config=client._config)
|
reranker = client.reranker
|
||||||
|
|
||||||
if reranker is None:
|
if reranker is None:
|
||||||
chunk_results = await client.chunk_repository.search(
|
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"] == ""
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_search_with_pil_image_works_like_bytes(temp_db_path, monkeypatch):
|
async def test_search_with_pil_image_works_like_bytes(temp_db_path, monkeypatch):
|
||||||
from PIL import Image as PILImageModule
|
from PIL import Image as PILImageModule
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue