Merge pull request #126 from ggozad/feat/zero-entropy

Support for zero entropy reranker

README.md

@@ -15,7 +15,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
 - **Multiple QA providers**: Any provider/model supported by Pydantic AI
 - **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI
 - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
-- **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM
+- **Reranking**: Default search result reranking with MixedBread AI, Cohere, Zero Entropy, or vLLM
 - **Question answering**: Built-in QA agents on your documents
 - **File monitoring**: Auto-index files when run as server
 - **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs

docs/configuration.md

@@ -82,7 +82,7 @@ embeddings:
   vector_dim: 4096
 
 reranking:
-  provider: ""  # Empty to disable, or mxbai, cohere, vllm
+  provider: ""  # Empty to disable, or mxbai, cohere, zeroentropy, vllm
   model: ""
 
 qa:
@@ -412,7 +412,13 @@ reranking:
 
 ### Cohere
 
-Cohere reranking is included in the default installation:
+Install with cohere extras:
+
+```bash
+uv pip install haiku.rag[cohere]
+```
+
+Then configure:
 
 ```yaml
 reranking:
@@ -426,6 +432,28 @@ Set your API key via environment variable:
 export CO_API_KEY=your-api-key
 ```
 
+### Zero Entropy
+
+Install with zeroentropy extras:
+
+```bash
+uv pip install haiku.rag[zeroentropy]
+```
+
+Then configure:
+
+```yaml
+reranking:
+  provider: zeroentropy
+  model: zerank-1  # Currently the only available model
+```
+
+Set your API key via environment variable:
+
+```bash
+export ZEROENTROPY_API_KEY=your-api-key
+```
+
 ### vLLM
 
 For high-performance local reranking using dedicated reranking models:

docs/installation.md

@@ -10,25 +10,40 @@ This includes support for:
 - **Ollama** (default embedding provider using `mxbai-embed-large`)
 - **OpenAI** (GPT models for QA and embeddings)
 - **Anthropic** (Claude models for QA)
-- **Cohere** (reranking models)
 - **vLLM** (high-performance local inference for embeddings, QA, and reranking)
 
 ## Provider-Specific Installation
 
-For additional embedding providers, install with extras:
+For additional providers, install with extras:
 
-### VoyageAI
+### Embedding Providers
+
+#### VoyageAI
 
 ```bash
 uv pip install haiku.rag[voyageai]
 ```
 
-### MixedBread AI Reranking
+### Reranking Providers
+
+#### MixedBread AI
 
 ```bash
 uv pip install haiku.rag[mxbai]
 ```
 
+#### Cohere
+
+```bash
+uv pip install haiku.rag[cohere]
+```
+
+#### Zero Entropy
+
+```bash
+uv pip install haiku.rag[zeroentropy]
+```
+
 ### vLLM Setup
 
 vLLM requires no additional installation - it works with the base haiku.rag package. However, you need to run vLLM servers separately:

pyproject.toml

@@ -42,6 +42,8 @@ dependencies = [
 voyageai = ["voyageai>=0.3.5"]
 mxbai = ["mxbai-rerank>=0.1.6"]
 a2a = ["fasta2a>=0.1.0"]
+cohere = ["cohere>=5.0.0"]
+zeroentropy = ["zeroentropy>=0.1.0a6"]
 
 [project.scripts]
 haiku-rag = "haiku.rag.cli:cli"

src/haiku/rag/reranking/__init__.py

@@ -41,5 +41,23 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
         except ImportError:
             reranker = None
 
+    elif config.reranking.provider == "vllm":
+        try:
+            from haiku.rag.reranking.vllm import VLLMReranker
+
+            reranker = VLLMReranker(config.reranking.model)
+        except ImportError:
+            reranker = None
+
+    elif config.reranking.provider == "zeroentropy":
+        try:
+            from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
+
+            # Use configured model or default to zerank-1
+            model = config.reranking.model or "zerank-1"
+            reranker = ZeroEntropyReranker(model)
+        except ImportError:
+            reranker = None
+
     _reranker_cache[config_id] = reranker
     return reranker

src/haiku/rag/reranking/zeroentropy.py

@@ -0,0 +1,59 @@
+from zeroentropy import ZeroEntropy
+
+from haiku.rag.reranking.base import RerankerBase
+from haiku.rag.store.models.chunk import Chunk
+
+
+class ZeroEntropyReranker(RerankerBase):
+    """Zero Entropy reranker implementation using the zerank-1 model."""
+
+    def __init__(self, model: str = "zerank-1"):
+        """Initialize the Zero Entropy reranker.
+
+        Args:
+            model: The Zero Entropy model to use (default: "zerank-1")
+        """
+        self._model = model
+        # Zero Entropy SDK reads ZEROENTROPY_API_KEY from environment by default
+        self._client = ZeroEntropy()
+
+    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]
+
+        # Call Zero Entropy reranking API
+        response = self._client.models.rerank(
+            model=self._model,
+            query=query,
+            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
+            chunk_index = result.index
+            score = result.relevance_score
+
+            if chunk_index < len(chunks):
+                reranked_results.append((chunks[chunk_index], score))
+
+        return reranked_results

tests/test_reranker.py

@@ -9,6 +9,7 @@ from haiku.rag.store.models.chunk import Chunk
 
 COHERE_AVAILABLE = bool(os.getenv("CO_API_KEY"))
 VLLM_RERANK_AVAILABLE = bool(Config.providers.vllm.rerank_base_url)
+ZEROENTROPY_AVAILABLE = bool(os.getenv("ZEROENTROPY_API_KEY"))
 
 chunks = [
     Chunk(content=content, document_id=str(i))
@@ -89,3 +90,26 @@ async def test_vllm_reranker():
     except Exception:
         # Skip test if vLLM rerank server is not available
         pytest.skip("vLLM rerank server not available")
+
+
+@pytest.mark.asyncio
+@pytest.mark.skipif(
+    not ZEROENTROPY_AVAILABLE, reason="Zero Entropy API key not available"
+)
+async def test_zeroentropy_reranker():
+    try:
+        from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
+
+        reranker = ZeroEntropyReranker("zerank-1")
+
+        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("Zero Entropy package not installed")

uv.lock

@@ -1173,12 +1173,18 @@ dependencies = [
 a2a = [
     { name = "fasta2a" },
 ]
+cohere = [
+    { name = "cohere" },
+]
 mxbai = [
     { name = "mxbai-rerank" },
 ]
 voyageai = [
     { name = "voyageai" },
 ]
+zeroentropy = [
+    { name = "zeroentropy" },
+]
 
 [package.dev-dependencies]
 dev = [
@@ -1197,6 +1203,7 @@ dev = [
 
 [package.metadata]
 requires-dist = [
+    { name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.0.0" },
     { name = "docling", specifier = ">=2.58.0" },
     { name = "fasta2a", marker = "extra == 'a2a'", specifier = ">=0.1.0" },
     { name = "fastmcp", specifier = ">=2.13.0.2" },
@@ -1214,8 +1221,9 @@ requires-dist = [
     { name = "typer", specifier = ">=0.19.2,<0.20.0" },
     { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.5" },
     { name = "watchfiles", specifier = ">=1.1.1" },
+    { name = "zeroentropy", marker = "extra == 'zeroentropy'", specifier = ">=0.1.0a6" },
 ]
-provides-extras = ["voyageai", "mxbai", "a2a"]
+provides-extras = ["voyageai", "mxbai", "a2a", "cohere", "zeroentropy"]
 
 [package.metadata.requires-dev]
 dev = [
@@ -4718,6 +4726,23 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b4/2d/2345fce04cfd4bee161bf1e7d9cdc702e3e16109021035dbb24db654a622/yarl-1.20.1-py3-none-any.whl", hash = "sha256:83b8eb083fe4683c6115795d9fc1cfaf2cbbefb19b3a1cb68f6527460f483a77", size = 46542, upload-time = "2025-06-10T00:46:07.521Z" },
 ]
 
+[[package]]
+name = "zeroentropy"
+version = "0.1.0a6"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio" },
+    { name = "distro" },
+    { name = "httpx" },
+    { name = "pydantic" },
+    { name = "sniffio" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/08/14/05c8caaa25ae64008c2f5b021cefd30276f00845d99bab4763be300da93a/zeroentropy-0.1.0a6.tar.gz", hash = "sha256:04f38e7b40f39cfdd4bb16df0ab0b18d8f33d516c1e9a39e0494a5fb7fba358d", size = 112726, upload-time = "2025-07-08T01:51:42.796Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3b/7e/594e9ec5cda6d8f4dc249c5ca0c6b31b19886eafedcea6c41a2400a2a7b9/zeroentropy-0.1.0a6-py3-none-any.whl", hash = "sha256:0caa6c4a450af80892d42848036eea0b766e3fe0bf6a097a613ab9403fdf5ad6", size = 101501, upload-time = "2025-07-08T01:51:41.527Z" },
+]
+
 [[package]]
 name = "zipp"
 version = "3.23.0"