From b8cf8f5198a93407a38e51b2eab8b5b773a3d5ea Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 20 Jan 2026 11:49:17 +0200 Subject: [PATCH 1/5] Remove reranker cache --- CHANGELOG.md | 4 +++ .../haiku/rag/reranking/__init__.py | 35 +++++++------------ tests/test_reranker.py | 28 +-------------- 3 files changed, 17 insertions(+), 50 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e4391d35..3b2f862b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # Changelog ## [Unreleased] +- **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 +19,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 diff --git a/haiku_rag_slim/haiku/rag/reranking/__init__.py b/haiku_rag_slim/haiku/rag/reranking/__init__.py index b5e52c50..8c475221 100644 --- a/haiku_rag_slim/haiku/rag/reranking/__init__.py +++ b/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,41 @@ 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 + return None diff --git a/tests/test_reranker.py b/tests/test_reranker.py index efc8f81c..a4c34e0c 100644 --- a/tests/test_reranker.py +++ b/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( From 7ff014684ab931c9c2f48156c777c3b511746387 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 21 Jan 2026 10:28:02 +0200 Subject: [PATCH 2/5] Support for jina reranker, both local and API --- docs/configuration/providers.md | 40 ++++++++++ .../haiku/rag/reranking/__init__.py | 15 ++++ haiku_rag_slim/haiku/rag/reranking/jina.py | 50 ++++++++++++ .../haiku/rag/reranking/jina_local.py | 46 +++++++++++ haiku_rag_slim/pyproject.toml | 1 + .../test_reranker/test_jina_reranker.yaml | 76 +++++++++++++++++++ tests/test_reranker.py | 73 ++++++++++++++++++ uv.lock | 8 +- 8 files changed, 308 insertions(+), 1 deletion(-) create mode 100644 haiku_rag_slim/haiku/rag/reranking/jina.py create mode 100644 haiku_rag_slim/haiku/rag/reranking/jina_local.py create mode 100644 tests/cassettes/test_reranker/test_jina_reranker.yaml diff --git a/docs/configuration/providers.md b/docs/configuration/providers.md index 12bb9603..7f7bbded 100644 --- a/docs/configuration/providers.md +++ b/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. diff --git a/haiku_rag_slim/haiku/rag/reranking/__init__.py b/haiku_rag_slim/haiku/rag/reranking/__init__.py index 8c475221..f3fbb2c3 100644 --- a/haiku_rag_slim/haiku/rag/reranking/__init__.py +++ b/haiku_rag_slim/haiku/rag/reranking/__init__.py @@ -52,4 +52,19 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None: except ImportError: # pragma: no cover return None + 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 diff --git a/haiku_rag_slim/haiku/rag/reranking/jina.py b/haiku_rag_slim/haiku/rag/reranking/jina.py new file mode 100644 index 00000000..d0e156fe --- /dev/null +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/reranking/jina_local.py b/haiku_rag_slim/haiku/rag/reranking/jina_local.py new file mode 100644 index 00000000..2f8a228f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/reranking/jina_local.py @@ -0,0 +1,46 @@ +try: + from transformers import ( + AutoModelForSequenceClassification, # 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): + """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 = AutoModelForSequenceClassification.from_pretrained( + model, trust_remote_code=True + ) + + 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] + sentence_pairs = [[query, doc] for doc in documents] + + scores = self._reranker.compute_score(sentence_pairs) + + # Handle both single score and list of scores + if isinstance(scores, (int, float)): + scores = [scores] + + scored_chunks = list(zip(chunks, scores, strict=False)) + scored_chunks.sort(key=lambda x: x[1], reverse=True) + + return [(chunk, float(score)) for chunk, score in scored_chunks[:top_n]] diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 533eea17..61a2d37b 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/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) diff --git a/tests/cassettes/test_reranker/test_jina_reranker.yaml b/tests/cassettes/test_reranker/test_jina_reranker.yaml new file mode 100644 index 00000000..94a27344 --- /dev/null +++ b/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 diff --git a/tests/test_reranker.py b/tests/test_reranker.py index a4c34e0c..d79cfcc4 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -202,3 +202,76 @@ 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") + + +@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 +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") diff --git a/uv.lock b/uv.lock index 412dbe7c..ed24b923 100644 --- a/uv.lock +++ b/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" From 36a675e1bb802b61c0fe9b6854ad6b755bbade5d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 21 Jan 2026 10:40:22 +0200 Subject: [PATCH 3/5] Add jina, mxbai and sentence-transformers to download-models when used --- CHANGELOG.md | 2 ++ haiku_rag_slim/haiku/rag/client.py | 53 ++++++++++++++++++++++++++++-- tests/test_reranker.py | 1 + 3 files changed, 53 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b2f862b..53441028 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ # 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 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 9f1460b1..056446fe 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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,51 @@ 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": + 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: + 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 AutoModelForSequenceClassification + + yield DownloadProgress(model=model_name, status="start") + await asyncio.to_thread( + AutoModelForSequenceClassification.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": diff --git a/tests/test_reranker.py b/tests/test_reranker.py index d79cfcc4..74f19b6f 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -259,6 +259,7 @@ async def test_jina_reranker(monkeypatch): @pytest.mark.asyncio +@pytest.mark.integration async def test_jina_local_reranker(): try: from haiku.rag.reranking.jina_local import JinaLocalReranker From ce95dc47a5d05221421efbdf072015ec486c4246 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 21 Jan 2026 11:09:35 +0200 Subject: [PATCH 4/5] Do not run coverage on download-models & code that requires integration tests --- haiku_rag_slim/haiku/rag/client.py | 6 ++++-- .../haiku/rag/reranking/jina_local.py | 2 +- tests/test_reranker.py | 20 +++++++++++++++++++ 3 files changed, 25 insertions(+), 3 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 056446fe..b2507493 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -1681,7 +1681,9 @@ class HaikuRAG: yield DownloadProgress(model=tokenizer_name, status="done") # Sentence-transformers embedder - if self._config.embeddings.model.provider == "sentence-transformers": + if ( + self._config.embeddings.model.provider == "sentence-transformers" + ): # pragma: no cover try: from sentence_transformers import ( # type: ignore[import-not-found] SentenceTransformer, @@ -1695,7 +1697,7 @@ class HaikuRAG: pass # HuggingFace reranker models - if self._config.reranking.model: + if self._config.reranking.model: # pragma: no cover provider = self._config.reranking.model.provider model_name = self._config.reranking.model.name diff --git a/haiku_rag_slim/haiku/rag/reranking/jina_local.py b/haiku_rag_slim/haiku/rag/reranking/jina_local.py index 2f8a228f..bbe97898 100644 --- a/haiku_rag_slim/haiku/rag/reranking/jina_local.py +++ b/haiku_rag_slim/haiku/rag/reranking/jina_local.py @@ -12,7 +12,7 @@ from haiku.rag.reranking.base import RerankerBase from haiku.rag.store.models.chunk import Chunk -class JinaLocalReranker(RerankerBase): +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, diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 74f19b6f..8d1d3f9f 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -235,6 +235,26 @@ class TestGetReranker: 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): From a1abfd9666e01bc8cec0f9602b666cda8ccf39d2 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 21 Jan 2026 11:26:23 +0200 Subject: [PATCH 5/5] Jina uses AutoModel not AutoModelForSequenceClassification --- haiku_rag_slim/haiku/rag/client.py | 4 ++-- .../haiku/rag/reranking/jina_local.py | 19 +++++-------------- 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index b2507493..58831c93 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -1715,11 +1715,11 @@ class HaikuRAG: elif provider == "jina-local": try: - from transformers import AutoModelForSequenceClassification + from transformers import AutoModel yield DownloadProgress(model=model_name, status="start") await asyncio.to_thread( - AutoModelForSequenceClassification.from_pretrained, + AutoModel.from_pretrained, model_name, trust_remote_code=True, ) diff --git a/haiku_rag_slim/haiku/rag/reranking/jina_local.py b/haiku_rag_slim/haiku/rag/reranking/jina_local.py index bbe97898..b76e6e64 100644 --- a/haiku_rag_slim/haiku/rag/reranking/jina_local.py +++ b/haiku_rag_slim/haiku/rag/reranking/jina_local.py @@ -1,6 +1,6 @@ try: from transformers import ( - AutoModelForSequenceClassification, # pyright: ignore[reportMissingImports] + AutoModel, # pyright: ignore[reportMissingImports] ) except ImportError as e: raise ImportError( @@ -21,9 +21,8 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover def __init__(self, model: str = "jinaai/jina-reranker-v3"): self._model = model - self._reranker = AutoModelForSequenceClassification.from_pretrained( - model, trust_remote_code=True - ) + 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 @@ -32,15 +31,7 @@ class JinaLocalReranker(RerankerBase): # pragma: no cover return [] documents = [chunk.content for chunk in chunks] - sentence_pairs = [[query, doc] for doc in documents] - scores = self._reranker.compute_score(sentence_pairs) + results = self._reranker.rerank(query, documents, top_n=top_n) - # Handle both single score and list of scores - if isinstance(scores, (int, float)): - scores = [scores] - - scored_chunks = list(zip(chunks, scores, strict=False)) - scored_chunks.sort(key=lambda x: x[1], reverse=True) - - return [(chunk, float(score)) for chunk, score in scored_chunks[:top_n]] + return [(chunks[r["index"]], float(r["relevance_score"])) for r in results]