Merge pull request #251 from ggozad/feat/jina-reranker

Jina Reranker v3 (local & API)

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+- **Jina Reranker v3**: Added support for Jina reranking with API mode (`provider: jina`) and local inference (`provider: jina-local`, requires `[jina]` extra)
+- **Model Downloads**: `download-models` now pre-downloads HuggingFace models for `sentence-transformers`, `mxbai`, and `jina-local`
+- **Reranker Factory**: Removed unreliable `id(config)`-based caching from `get_reranker()`; factory now always instantiates fresh
+
 ## [0.26.7] - 2026-01-20
 
 ### Added
@@ -17,6 +21,8 @@
 
 ### Changed
 
+- **MCP Error Handling**: MCP tools now let exceptions propagate naturally; FastMCP converts them to proper MCP error responses
+- **Chunk Contextualization**: Consolidated duplicate `contextualize` logic into `Chunk.contextualize_content()` method
 - **Type Checker**: Replaced pyright with [ty](https://github.com/astral-sh/ty), Astral's extremely fast Python type checker
   - Added explicit `Agent[Deps, Output]` type annotations to all pydantic-ai agents for better type inference
   - Removed ~24 unnecessary `# type: ignore` comments that ty correctly infers

docs/configuration/providers.md

@@ -382,3 +382,43 @@ reranking:
 ```
 
 **Note:** vLLM reranking uses the `/v1/rerank` API endpoint. You need to run a vLLM server separately with a reranking model loaded.
+
+### Jina AI
+
+Jina provides high-quality reranking with two deployment options: API mode and local inference.
+
+#### API Mode
+
+Use the Jina Reranker API for cloud-based reranking:
+
+```yaml
+reranking:
+  model:
+    provider: jina
+    name: jina-reranker-v3
+```
+
+Set your API key via environment variable:
+
+```bash
+export JINA_API_KEY=your-api-key
+```
+
+#### Local Mode
+
+For local inference, install the jina extra:
+
+```bash
+uv pip install haiku.rag-slim[jina]
+```
+
+Then configure:
+
+```yaml
+reranking:
+  model:
+    provider: jina-local
+    name: jinaai/jina-reranker-v3
+```
+
+**Note:** The Jina Reranker v3 local model is licensed under CC BY-NC 4.0, which restricts commercial use. For commercial applications, use the API mode instead.

haiku_rag_slim/haiku/rag/client.py

@@ -1656,9 +1656,11 @@ class HaikuRAG:
         """Download required models, yielding progress events.
 
         Yields DownloadProgress events for:
-        - Docling models (status="docling_start", "docling_done")
-        - HuggingFace tokenizer (status="tokenizer_start", "tokenizer_done")
-        - Ollama models (status="pulling", "downloading", "done", or other Ollama statuses)
+        - Docling models
+        - HuggingFace tokenizer
+        - Sentence-transformers embedder (if configured)
+        - HuggingFace reranker models (mxbai, jina-local)
+        - Ollama models
         """
         # Docling models
         try:
@@ -1678,6 +1680,53 @@ class HaikuRAG:
         await asyncio.to_thread(AutoTokenizer.from_pretrained, tokenizer_name)
         yield DownloadProgress(model=tokenizer_name, status="done")
 
+        # Sentence-transformers embedder
+        if (
+            self._config.embeddings.model.provider == "sentence-transformers"
+        ):  # pragma: no cover
+            try:
+                from sentence_transformers import (  # type: ignore[import-not-found]
+                    SentenceTransformer,
+                )
+
+                model_name = self._config.embeddings.model.name
+                yield DownloadProgress(model=model_name, status="start")
+                await asyncio.to_thread(SentenceTransformer, model_name)
+                yield DownloadProgress(model=model_name, status="done")
+            except ImportError:
+                pass
+
+        # HuggingFace reranker models
+        if self._config.reranking.model:  # pragma: no cover
+            provider = self._config.reranking.model.provider
+            model_name = self._config.reranking.model.name
+
+            if provider == "mxbai":
+                try:
+                    from mxbai_rerank import MxbaiRerankV2
+
+                    yield DownloadProgress(model=model_name, status="start")
+                    await asyncio.to_thread(
+                        MxbaiRerankV2, model_name, disable_transformers_warnings=True
+                    )
+                    yield DownloadProgress(model=model_name, status="done")
+                except ImportError:
+                    pass
+
+            elif provider == "jina-local":
+                try:
+                    from transformers import AutoModel
+
+                    yield DownloadProgress(model=model_name, status="start")
+                    await asyncio.to_thread(
+                        AutoModel.from_pretrained,
+                        model_name,
+                        trust_remote_code=True,
+                    )
+                    yield DownloadProgress(model=model_name, status="done")
+                except ImportError:
+                    pass
+
         # Collect Ollama models from config
         required_models: set[str] = set()
         if self._config.embeddings.model.provider == "ollama":

haiku_rag_slim/haiku/rag/reranking/__init__.py

@@ -3,8 +3,6 @@ import os
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.reranking.base import RerankerBase
 
-_reranker_cache: dict[int, RerankerBase | None] = {}
-
 
 def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
     """
@@ -17,50 +15,56 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
     Returns:
         A reranker instance if configured, None otherwise.
     """
-    # Use config id as cache key to support multiple configs
-    config_id = id(config)
-    if config_id in _reranker_cache:
-        return _reranker_cache[config_id]
-
-    reranker: RerankerBase | None = None
-
     if config.reranking.model and config.reranking.model.provider == "mxbai":
         try:
             from haiku.rag.reranking.mxbai import MxBAIReranker
 
             os.environ["TOKENIZERS_PARALLELISM"] = "true"
-            reranker = MxBAIReranker()
+            return MxBAIReranker()
         except ImportError:  # pragma: no cover
-            reranker = None
+            return None
 
-    elif config.reranking.model and config.reranking.model.provider == "cohere":
+    if config.reranking.model and config.reranking.model.provider == "cohere":
         try:
             from haiku.rag.reranking.cohere import CohereReranker
 
-            reranker = CohereReranker()
+            return CohereReranker()
         except ImportError:  # pragma: no cover
-            reranker = None
+            return None
 
-    elif config.reranking.model and config.reranking.model.provider == "vllm":
+    if config.reranking.model and config.reranking.model.provider == "vllm":
         try:
             from haiku.rag.reranking.vllm import VLLMReranker
 
             base_url = config.reranking.model.base_url
             if not base_url:
                 raise ValueError("vLLM reranker requires base_url in reranking.model")
-            reranker = VLLMReranker(config.reranking.model.name, base_url)
+            return VLLMReranker(config.reranking.model.name, base_url)
         except ImportError:  # pragma: no cover
-            reranker = None
+            return None
 
-    elif config.reranking.model and config.reranking.model.provider == "zeroentropy":
+    if config.reranking.model and config.reranking.model.provider == "zeroentropy":
         try:
             from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
 
-            # Use configured model or default to zerank-1
             model = config.reranking.model.name or "zerank-1"
-            reranker = ZeroEntropyReranker(model)
+            return ZeroEntropyReranker(model)
         except ImportError:  # pragma: no cover
-            reranker = None
+            return None
 
-    _reranker_cache[config_id] = reranker
-    return reranker
+    if config.reranking.model and config.reranking.model.provider == "jina":
+        from haiku.rag.reranking.jina import JinaReranker
+
+        model = config.reranking.model.name or "jina-reranker-v3"
+        return JinaReranker(model)
+
+    if config.reranking.model and config.reranking.model.provider == "jina-local":
+        try:
+            from haiku.rag.reranking.jina_local import JinaLocalReranker
+
+            model = config.reranking.model.name or "jinaai/jina-reranker-v3"
+            return JinaLocalReranker(model)
+        except ImportError:  # pragma: no cover
+            return None
+
+    return None

haiku_rag_slim/haiku/rag/reranking/jina.py

@@ -0,0 +1,50 @@
+import os
+
+import httpx
+
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+
+class JinaReranker(RerankerBase):
+    """Jina AI reranker using the Jina Reranker API."""
+
+    def __init__(self, model: str = "jina-reranker-v3"):
+        self._model = model
+        self._api_key = os.environ.get("JINA_API_KEY")
+        if not self._api_key:
+            raise ValueError("JINA_API_KEY environment variable required")
+
+    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:
+            response = await client.post(
+                "https://api.jina.ai/v1/rerank",
+                json={
+                    "model": self._model,
+                    "query": query,
+                    "documents": documents,
+                    "top_n": top_n,
+                },
+                headers={
+                    "Authorization": f"Bearer {self._api_key}",
+                    "Content-Type": "application/json",
+                },
+            )
+            response.raise_for_status()
+
+            result = response.json()
+
+            scored_chunks = []
+            for item in result.get("results", []):
+                index = item["index"]
+                score = item["relevance_score"]
+                scored_chunks.append((chunks[index], score))
+
+            return scored_chunks

haiku_rag_slim/haiku/rag/reranking/jina_local.py

@@ -0,0 +1,37 @@
+try:
+    from transformers import (
+        AutoModel,  # pyright: ignore[reportMissingImports]
+    )
+except ImportError as e:
+    raise ImportError(
+        "transformers is not installed. Please install it with `pip install transformers torch` "
+        "or use the jina optional dependency."
+    ) from e
+
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+
+class JinaLocalReranker(RerankerBase):  # pragma: no cover
+    """Jina reranker using local model inference via transformers.
+
+    Note: The Jina Reranker v3 model is licensed under CC BY-NC 4.0,
+    which restricts commercial use.
+    """
+
+    def __init__(self, model: str = "jinaai/jina-reranker-v3"):
+        self._model = model
+        self._reranker = AutoModel.from_pretrained(model, trust_remote_code=True)
+        self._reranker.eval()
+
+    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 = self._reranker.rerank(query, documents, top_n=top_n)
+
+        return [(chunks[r["index"]], float(r["relevance_score"])) for r in results]

haiku_rag_slim/pyproject.toml

@@ -44,6 +44,7 @@ voyageai = ["voyageai>=0.3.7"]
 mxbai = ["mxbai-rerank>=0.1.6"]
 cohere = ["cohere>=5.20.1"]
 zeroentropy = ["zeroentropy>=0.1.0a7"]
+jina = ["transformers>=4.40.0", "torch>=2.0.0"]
 # TUI (chat and inspect commands)
 tui = ["textual>=7.3.0", "textual-image>=0.8.5"]
 # Model providers (delegated to pydantic-ai-slim)

tests/cassettes/test_reranker/test_jina_reranker.yaml

@@ -0,0 +1,76 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - '*/*'
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '1266'
+      content-type:
+      - application/json
+      host:
+      - api.jina.ai
+    method: POST
+    parsed_body:
+      documents:
+      - To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the Pulitzer
+        Prize, and has become a classic of modern American literature.
+      - The novel Moby-Dick was written by Herman Melville and first published in 1851. It is considered a masterpiece of
+        American literature and deals with complex themes of obsession, revenge, and the conflict between good and evil.
+      - Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville,
+        Alabama. She received the Pulitzer Prize for Fiction in 1961.
+      - Jane Austen was an English novelist known primarily for her six major novels, which interpret, critique and comment
+        upon the British landed gentry at the end of the 18th century.
+      - The Harry Potter series, which consists of seven fantasy novels written by British author J.K. Rowling, is among the
+        most popular and critically acclaimed books of the modern era.
+      - The Great Gatsby, a novel written by American author F. Scott Fitzgerald, was published in 1925. The story is set
+        in the Jazz Age and follows the life of millionaire Jay Gatsby and his pursuit of Daisy Buchanan.
+      model: jina-reranker-v3
+      query: Who wrote 'To Kill a Mockingbird'?
+      top_n: 2
+    uri: https://api.jina.ai/v1/rerank
+  response:
+    headers:
+      alt-svc:
+      - h3=":443"; ma=86400
+      cache-control:
+      - private
+      connection:
+      - keep-alive
+      content-length:
+      - '570'
+      content-type:
+      - application/json
+      expires:
+      - Wed, 21 Jan 2026 08:21:12 GMT
+      nel:
+      - '{"report_to":"cf-nel","success_fraction":0.0,"max_age":604800}'
+      report-to:
+      - '{"group":"cf-nel","max_age":604800,"endpoints":[{"url":"https://a.nel.cloudflare.com/report/v4?s=TX0p0eBpMlJ8P0KA1iF7CLglvZSaumFNwZXrU%2BYgFDRyrC0xrRB3NyaekMeKtZQKWf8GpGD%2Bb5WCh4vN67uaCzgah3QbrVh3PvU%2FSEMF6BZcfFMY"}]}'
+      transfer-encoding:
+      - chunked
+      vary:
+      - Accept-Encoding
+    parsed_body:
+      model: jina-reranker-v3
+      object: list
+      results:
+      - document:
+          text: To Kill a Mockingbird is a novel by Harper Lee published in 1960. It was immediately successful, winning the
+            Pulitzer Prize, and has become a classic of modern American literature.
+        index: 0
+        relevance_score: 0.42858967
+      - document:
+          text: Harper Lee, an American novelist widely known for her novel To Kill a Mockingbird, was born in 1926 in Monroeville,
+            Alabama. She received the Pulitzer Prize for Fiction in 1961.
+        index: 2
+        relevance_score: 0.07394931
+      usage:
+        total_tokens: 490
+    status:
+      code: 200
+      message: OK
+version: 1

tests/test_reranker.py

@@ -3,7 +3,7 @@ from pathlib import Path
 import pytest
 
 from haiku.rag.config.models import AppConfig, ModelConfig, RerankingConfig
-from haiku.rag.reranking import _reranker_cache, get_reranker
+from haiku.rag.reranking import get_reranker
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
@@ -13,14 +13,6 @@ def vcr_cassette_dir():
     return str(Path(__file__).parent / "cassettes" / "test_reranker")
 
 
-@pytest.fixture(autouse=True)
-def clear_reranker_cache():
-    """Clear the reranker cache before each test."""
-    _reranker_cache.clear()
-    yield
-    _reranker_cache.clear()
-
-
 chunks = [
     Chunk(content=content, document_id=str(i))
     for i, content in enumerate(
@@ -202,24 +194,6 @@ class TestGetReranker:
         except ImportError:
             pytest.skip("Zero Entropy package not installed")
 
-    def test_caching_returns_same_instance(self):
-        config = AppConfig(reranking=RerankingConfig(model=None))
-        result1 = get_reranker(config)
-        result2 = get_reranker(config)
-        assert result1 is result2
-
-    def test_different_configs_get_separate_cache_entries(self):
-        config1 = AppConfig(reranking=RerankingConfig(model=None))
-        config2 = AppConfig(reranking=RerankingConfig(model=None))
-
-        result1 = get_reranker(config1)
-        result2 = get_reranker(config2)
-
-        # Both return None, but they should be cached separately
-        assert result1 is None
-        assert result2 is None
-        assert len(_reranker_cache) == 2
-
     def test_unknown_provider_returns_none(self):
         config = AppConfig(
             reranking=RerankingConfig(
@@ -228,3 +202,97 @@ class TestGetReranker:
         )
         result = get_reranker(config)
         assert result is None
+
+    def test_jina_provider(self, monkeypatch):
+        monkeypatch.setenv("JINA_API_KEY", "test-api-key")
+
+        from haiku.rag.reranking.jina import JinaReranker
+
+        config = AppConfig(
+            reranking=RerankingConfig(
+                model=ModelConfig(provider="jina", name="jina-reranker-v3")
+            )
+        )
+        result = get_reranker(config)
+        assert isinstance(result, JinaReranker)
+        assert result._model == "jina-reranker-v3"
+
+    def test_jina_local_provider(self):
+        try:
+            from haiku.rag.reranking.jina_local import JinaLocalReranker
+
+            config = AppConfig(
+                reranking=RerankingConfig(
+                    model=ModelConfig(
+                        provider="jina-local", name="jinaai/jina-reranker-v3"
+                    )
+                )
+            )
+            result = get_reranker(config)
+            assert isinstance(result, JinaLocalReranker)
+            assert result._model == "jinaai/jina-reranker-v3"
+        except ImportError:
+            pytest.skip("Jina local dependencies not installed")
+
+
+def test_jina_reranker_missing_api_key(monkeypatch):
+    monkeypatch.delenv("JINA_API_KEY", raising=False)
+
+    from haiku.rag.reranking.jina import JinaReranker
+
+    with pytest.raises(ValueError, match="JINA_API_KEY environment variable required"):
+        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):
+    import os
+
+    # Only set dummy key if real key not present (for VCR playback)
+    if not os.environ.get("JINA_API_KEY"):
+        monkeypatch.setenv("JINA_API_KEY", "test-api-key")
+
+    from haiku.rag.reranking.jina import JinaReranker
+
+    reranker = JinaReranker("jina-reranker-v3")
+
+    reranked = await reranker.rerank(
+        "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
+    )
+    assert len(reranked) == 2
+    assert all(isinstance(score, float) for chunk, score in reranked)
+    # Check that the top results are relevant to Harper Lee / To Kill a Mockingbird
+    top_ids = [chunk.document_id for chunk, score in reranked]
+    assert "0" in top_ids or "2" in top_ids  # These chunks mention the book/author
+
+
+@pytest.mark.asyncio
+@pytest.mark.integration
+async def test_jina_local_reranker():
+    try:
+        from haiku.rag.reranking.jina_local import JinaLocalReranker
+
+        reranker = JinaLocalReranker("jinaai/jina-reranker-v3")
+
+        reranked = await reranker.rerank(
+            "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
+        )
+        assert len(reranked) == 2
+        assert all(isinstance(score, float) for chunk, score in reranked)
+        # Check that the top results are relevant to Harper Lee / To Kill a Mockingbird
+        top_ids = [chunk.document_id for chunk, score in reranked]
+        assert "0" in top_ids or "2" in top_ids  # These chunks mention the book/author
+    except ImportError:
+        pytest.skip("Jina local dependencies not installed")

uv.lock

@@ -1397,6 +1397,10 @@ google = [
 groq = [
     { name = "pydantic-ai-slim", extra = ["groq"] },
 ]
+jina = [
+    { name = "torch" },
+    { name = "transformers" },
+]
 mistral = [
     { name = "pydantic-ai-slim", extra = ["mistral"] },
 ]
@@ -1440,12 +1444,14 @@ requires-dist = [
     { name = "rich", specifier = ">=14.2.0" },
     { name = "textual", marker = "extra == 'tui'", specifier = ">=7.3.0" },
     { name = "textual-image", marker = "extra == 'tui'", specifier = ">=0.8.5" },
+    { name = "torch", marker = "extra == 'jina'", specifier = ">=2.0.0" },
+    { name = "transformers", marker = "extra == 'jina'", specifier = ">=4.40.0" },
     { name = "typer", specifier = ">=0.19.2,<0.20.0" },
     { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.7" },
     { name = "watchfiles", specifier = ">=1.1.1" },
     { name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a7" },
 ]
-provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
+provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jina", "tui", "anthropic", "groq", "google", "mistral", "bedrock", "vertexai"]
 
 [[package]]
 name = "hf-xet"