Hoist reranker empty-input guard into base; stop factory tests loading models

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

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."
         )

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"

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)

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:

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(

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(

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]
 

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
 

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")