Support for zero entropy reranker

This commit is contained in:
Yiorgis Gozadinos 2025-10-30 13:36:44 +02:00
parent f9b6894974
commit 2f00686605
No known key found for this signature in database
8 changed files with 179 additions and 8 deletions

View file

@ -15,7 +15,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
- **Multiple QA providers**: Any provider/model supported by Pydantic AI
- **Research graph (multiagent)**: 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

View file

@ -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:

View file

@ -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:

View file

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

View file

@ -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

View file

@ -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

View file

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

27
uv.lock
View file

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