diff --git a/CHANGELOG.md b/CHANGELOG.md index 06494949..cba408f5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Changed + +- Custom rerankers override `RerankerBase._rerank` instead of `rerank`; the base `rerank` handles the empty-input short-circuit. + ## [0.63.1] - 2026-06-29 ### Changed diff --git a/haiku_rag_slim/haiku/rag/reranking/base.py b/haiku_rag_slim/haiku/rag/reranking/base.py index bdbe76f6..6c31d986 100644 --- a/haiku_rag_slim/haiku/rag/reranking/base.py +++ b/haiku_rag_slim/haiku/rag/reranking/base.py @@ -7,7 +7,14 @@ class RerankerBase: async def rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 + ) -> list[tuple[Chunk, float]]: + if not chunks: + return [] + return await self._rerank(query, chunks, top_n) + + async def _rerank( + self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: raise NotImplementedError( - "Reranker is an abstract class. Please implement the rerank method in a subclass." + "Reranker is an abstract class. Please implement the _rerank method in a subclass." ) diff --git a/haiku_rag_slim/haiku/rag/reranking/cohere.py b/haiku_rag_slim/haiku/rag/reranking/cohere.py index 9d38a631..25571bd7 100644 --- a/haiku_rag_slim/haiku/rag/reranking/cohere.py +++ b/haiku_rag_slim/haiku/rag/reranking/cohere.py @@ -14,12 +14,9 @@ class CohereReranker(RerankerBase): # pragma: no cover # Cohere SDK reads CO_API_KEY from environment by default self._client = cohere.AsyncClientV2() - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - documents = [chunk.content for chunk in chunks] model_name = self._model or "rerank-v3.5" diff --git a/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py b/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py index bbeff5b6..a12d45ee 100644 --- a/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py +++ b/haiku_rag_slim/haiku/rag/reranking/cross_encoder.py @@ -26,12 +26,9 @@ class CrossEncoderReranker(RerankerBase): self._model = model self._reranker = CrossEncoder(model) - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - documents = [chunk.content for chunk in chunks] rankings = await asyncio.to_thread( lambda: self._reranker.rank(query, documents, top_k=top_n) diff --git a/haiku_rag_slim/haiku/rag/reranking/jina.py b/haiku_rag_slim/haiku/rag/reranking/jina.py index d0e156fe..0144edb3 100644 --- a/haiku_rag_slim/haiku/rag/reranking/jina.py +++ b/haiku_rag_slim/haiku/rag/reranking/jina.py @@ -15,12 +15,9 @@ class JinaReranker(RerankerBase): if not self._api_key: raise ValueError("JINA_API_KEY environment variable required") - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - documents = [chunk.content for chunk in chunks] async with httpx.AsyncClient() as client: diff --git a/haiku_rag_slim/haiku/rag/reranking/jina_local.py b/haiku_rag_slim/haiku/rag/reranking/jina_local.py index 9b8949f7..16bc463c 100644 --- a/haiku_rag_slim/haiku/rag/reranking/jina_local.py +++ b/haiku_rag_slim/haiku/rag/reranking/jina_local.py @@ -26,12 +26,9 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover self._reranker = AutoModel.from_pretrained(model, trust_remote_code=True) self._reranker.eval() - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - documents = [chunk.content for chunk in chunks] results = await asyncio.to_thread( diff --git a/haiku_rag_slim/haiku/rag/reranking/mxbai.py b/haiku_rag_slim/haiku/rag/reranking/mxbai.py index 01ef50b9..c53713a8 100644 --- a/haiku_rag_slim/haiku/rag/reranking/mxbai.py +++ b/haiku_rag_slim/haiku/rag/reranking/mxbai.py @@ -25,12 +25,9 @@ class MxBAIReranker(RerankerBase): ) self._client = MxbaiRerankV2(model_name, disable_transformers_warnings=True) - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - documents = [chunk.content for chunk in chunks] results = await asyncio.to_thread( diff --git a/haiku_rag_slim/haiku/rag/reranking/vllm.py b/haiku_rag_slim/haiku/rag/reranking/vllm.py index ab370110..96271219 100644 --- a/haiku_rag_slim/haiku/rag/reranking/vllm.py +++ b/haiku_rag_slim/haiku/rag/reranking/vllm.py @@ -9,12 +9,9 @@ class VLLMReranker(RerankerBase): # pragma: no cover self._model = model self._base_url = base_url - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - # Prepare documents for reranking documents = [chunk.content for chunk in chunks] diff --git a/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py b/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py index 2880b86d..633fa769 100644 --- a/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py +++ b/haiku_rag_slim/haiku/rag/reranking/zeroentropy.py @@ -17,22 +17,9 @@ class ZeroEntropyReranker(RerankerBase): # pragma: no cover # Zero Entropy SDK reads ZEROENTROPY_API_KEY from environment by default self._client = AsyncZeroEntropy() - async def rerank( + async def _rerank( self, query: str, chunks: list[Chunk], top_n: int = 10 ) -> list[tuple[Chunk, float]]: - """Rerank the given chunks based on relevance to the query. - - Args: - query: The query to rank against - chunks: The chunks to rerank - top_n: The number of top results to return - - Returns: - A list of (chunk, score) tuples, sorted by relevance - """ - if not chunks: - return [] - # Prepare documents for Zero Entropy API documents = [chunk.content for chunk in chunks] @@ -44,13 +31,9 @@ class ZeroEntropyReranker(RerankerBase): # pragma: no cover documents=documents, ) - # Extract results and map back to chunks # Zero Entropy returns results sorted by relevance with scores reranked_results = [] - - # Get top_n results - for i, result in enumerate(response.results[:top_n]): - # Zero Entropy returns index and score for each document + for result in response.results[:top_n]: chunk_index = result.index score = result.relevance_score diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 0fc3c8f0..e67d3f3c 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -1,4 +1,5 @@ from pathlib import Path +from unittest.mock import MagicMock import pytest @@ -7,6 +8,14 @@ from haiku.rag.reranking import get_reranker from haiku.rag.reranking.base import RerankerBase from haiku.rag.store.models.chunk import Chunk +# Providers whose constructor loads a model in-process. Factory-routing tests +# patch the loader so they assert dispatch without paying the model load. +HEAVY_LOADERS = { + "mxbai": "MxbaiRerankV2", + "jina-local": "AutoModel", + "cross-encoder": "CrossEncoder", +} + @pytest.fixture(scope="module") def vcr_cassette_dir(): @@ -36,11 +45,16 @@ async def test_reranker_base(): expected_model = Config.reranking.model.name if Config.reranking.model else None assert reranker._model == expected_model + # Empty input short-circuits in the base class without dispatching to _rerank. + assert await reranker.rerank("query", []) == [] + + # The actual rerank step is abstract. with pytest.raises(NotImplementedError): - await reranker.rerank("query", []) + await reranker.rerank("query", chunks) @pytest.mark.asyncio +@pytest.mark.integration async def test_mxbai_reranker(): try: from haiku.rag.config import Config @@ -62,18 +76,6 @@ async def test_mxbai_reranker(): pytest.skip("MxBAI package not installed") -@pytest.mark.asyncio -async def test_mxbai_reranker_empty_chunks(): - try: - from haiku.rag.reranking.mxbai import MxBAIReranker - - reranker = MxBAIReranker() - result = await reranker.rerank("query", [], top_n=2) - assert result == [] - except ImportError: - pytest.skip("MxBAI package not installed") - - @pytest.mark.asyncio @pytest.mark.vcr() async def test_cohere_reranker(): @@ -251,6 +253,10 @@ class TestGetReranker: mod = pytest.importorskip(class_module) expected_class = getattr(mod, class_name) + loader_attr = HEAVY_LOADERS.get(provider) + if loader_attr: + monkeypatch.setattr(mod, loader_attr, MagicMock()) + for key, value in env_vars.items(): monkeypatch.setenv(key, value) @@ -277,17 +283,6 @@ def test_jina_reranker_missing_api_key(monkeypatch): JinaReranker("jina-reranker-v3") -@pytest.mark.asyncio -async def test_jina_reranker_empty_chunks(monkeypatch): - monkeypatch.setenv("JINA_API_KEY", "test-api-key") - - from haiku.rag.reranking.jina import JinaReranker - - reranker = JinaReranker("jina-reranker-v3") - result = await reranker.rerank("query", [], top_n=2) - assert result == [] - - @pytest.mark.asyncio @pytest.mark.vcr() async def test_jina_reranker(monkeypatch): @@ -332,6 +327,7 @@ async def test_jina_local_reranker(): @pytest.mark.asyncio +@pytest.mark.integration async def test_cross_encoder_reranker(): try: from haiku.rag.reranking.cross_encoder import CrossEncoderReranker @@ -347,15 +343,3 @@ async def test_cross_encoder_reranker(): assert "0" in top_ids or "2" in top_ids except ImportError: pytest.skip("sentence-transformers not installed") - - -@pytest.mark.asyncio -async def test_cross_encoder_reranker_empty_chunks(): - try: - from haiku.rag.reranking.cross_encoder import CrossEncoderReranker - - reranker = CrossEncoderReranker("cross-encoder/ms-marco-MiniLM-L-6-v2") - result = await reranker.rerank("query", [], top_n=2) - assert result == [] - except ImportError: - pytest.skip("sentence-transformers not installed")