Support for zero entropy reranker
This commit is contained in:
parent
f9b6894974
commit
2f00686605
8 changed files with 179 additions and 8 deletions
|
|
@ -15,7 +15,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
|
||||||
- **Multiple QA providers**: Any provider/model supported by Pydantic AI
|
- **Multiple QA providers**: Any provider/model supported by Pydantic AI
|
||||||
- **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI
|
- **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI
|
||||||
- **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
|
- **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
|
- **Question answering**: Built-in QA agents on your documents
|
||||||
- **File monitoring**: Auto-index files when run as server
|
- **File monitoring**: Auto-index files when run as server
|
||||||
- **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs
|
- **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs
|
||||||
|
|
|
||||||
|
|
@ -82,7 +82,7 @@ embeddings:
|
||||||
vector_dim: 4096
|
vector_dim: 4096
|
||||||
|
|
||||||
reranking:
|
reranking:
|
||||||
provider: "" # Empty to disable, or mxbai, cohere, vllm
|
provider: "" # Empty to disable, or mxbai, cohere, zeroentropy, vllm
|
||||||
model: ""
|
model: ""
|
||||||
|
|
||||||
qa:
|
qa:
|
||||||
|
|
@ -412,7 +412,13 @@ reranking:
|
||||||
|
|
||||||
### Cohere
|
### Cohere
|
||||||
|
|
||||||
Cohere reranking is included in the default installation:
|
Install with cohere extras:
|
||||||
|
|
||||||
|
```bash
|
||||||
|
uv pip install haiku.rag[cohere]
|
||||||
|
```
|
||||||
|
|
||||||
|
Then configure:
|
||||||
|
|
||||||
```yaml
|
```yaml
|
||||||
reranking:
|
reranking:
|
||||||
|
|
@ -426,6 +432,28 @@ Set your API key via environment variable:
|
||||||
export CO_API_KEY=your-api-key
|
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
|
### vLLM
|
||||||
|
|
||||||
For high-performance local reranking using dedicated reranking models:
|
For high-performance local reranking using dedicated reranking models:
|
||||||
|
|
|
||||||
|
|
@ -10,25 +10,40 @@ This includes support for:
|
||||||
- **Ollama** (default embedding provider using `mxbai-embed-large`)
|
- **Ollama** (default embedding provider using `mxbai-embed-large`)
|
||||||
- **OpenAI** (GPT models for QA and embeddings)
|
- **OpenAI** (GPT models for QA and embeddings)
|
||||||
- **Anthropic** (Claude models for QA)
|
- **Anthropic** (Claude models for QA)
|
||||||
- **Cohere** (reranking models)
|
|
||||||
- **vLLM** (high-performance local inference for embeddings, QA, and reranking)
|
- **vLLM** (high-performance local inference for embeddings, QA, and reranking)
|
||||||
|
|
||||||
## Provider-Specific Installation
|
## Provider-Specific Installation
|
||||||
|
|
||||||
For additional embedding providers, install with extras:
|
For additional providers, install with extras:
|
||||||
|
|
||||||
### VoyageAI
|
### Embedding Providers
|
||||||
|
|
||||||
|
#### VoyageAI
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv pip install haiku.rag[voyageai]
|
uv pip install haiku.rag[voyageai]
|
||||||
```
|
```
|
||||||
|
|
||||||
### MixedBread AI Reranking
|
### Reranking Providers
|
||||||
|
|
||||||
|
#### MixedBread AI
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv pip install haiku.rag[mxbai]
|
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 Setup
|
||||||
|
|
||||||
vLLM requires no additional installation - it works with the base haiku.rag package. However, you need to run vLLM servers separately:
|
vLLM requires no additional installation - it works with the base haiku.rag package. However, you need to run vLLM servers separately:
|
||||||
|
|
|
||||||
|
|
@ -42,6 +42,8 @@ dependencies = [
|
||||||
voyageai = ["voyageai>=0.3.5"]
|
voyageai = ["voyageai>=0.3.5"]
|
||||||
mxbai = ["mxbai-rerank>=0.1.6"]
|
mxbai = ["mxbai-rerank>=0.1.6"]
|
||||||
a2a = ["fasta2a>=0.1.0"]
|
a2a = ["fasta2a>=0.1.0"]
|
||||||
|
cohere = ["cohere>=5.0.0"]
|
||||||
|
zeroentropy = ["zeroentropy>=0.1.0a6"]
|
||||||
|
|
||||||
[project.scripts]
|
[project.scripts]
|
||||||
haiku-rag = "haiku.rag.cli:cli"
|
haiku-rag = "haiku.rag.cli:cli"
|
||||||
|
|
|
||||||
|
|
@ -41,5 +41,23 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
|
||||||
except ImportError:
|
except ImportError:
|
||||||
reranker = None
|
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
|
_reranker_cache[config_id] = reranker
|
||||||
return reranker
|
return reranker
|
||||||
|
|
|
||||||
59
src/haiku/rag/reranking/zeroentropy.py
Normal file
59
src/haiku/rag/reranking/zeroentropy.py
Normal 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
|
||||||
|
|
@ -9,6 +9,7 @@ from haiku.rag.store.models.chunk import Chunk
|
||||||
|
|
||||||
COHERE_AVAILABLE = bool(os.getenv("CO_API_KEY"))
|
COHERE_AVAILABLE = bool(os.getenv("CO_API_KEY"))
|
||||||
VLLM_RERANK_AVAILABLE = bool(Config.providers.vllm.rerank_base_url)
|
VLLM_RERANK_AVAILABLE = bool(Config.providers.vllm.rerank_base_url)
|
||||||
|
ZEROENTROPY_AVAILABLE = bool(os.getenv("ZEROENTROPY_API_KEY"))
|
||||||
|
|
||||||
chunks = [
|
chunks = [
|
||||||
Chunk(content=content, document_id=str(i))
|
Chunk(content=content, document_id=str(i))
|
||||||
|
|
@ -89,3 +90,26 @@ async def test_vllm_reranker():
|
||||||
except Exception:
|
except Exception:
|
||||||
# Skip test if vLLM rerank server is not available
|
# Skip test if vLLM rerank server is not available
|
||||||
pytest.skip("vLLM rerank server 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
27
uv.lock
|
|
@ -1173,12 +1173,18 @@ dependencies = [
|
||||||
a2a = [
|
a2a = [
|
||||||
{ name = "fasta2a" },
|
{ name = "fasta2a" },
|
||||||
]
|
]
|
||||||
|
cohere = [
|
||||||
|
{ name = "cohere" },
|
||||||
|
]
|
||||||
mxbai = [
|
mxbai = [
|
||||||
{ name = "mxbai-rerank" },
|
{ name = "mxbai-rerank" },
|
||||||
]
|
]
|
||||||
voyageai = [
|
voyageai = [
|
||||||
{ name = "voyageai" },
|
{ name = "voyageai" },
|
||||||
]
|
]
|
||||||
|
zeroentropy = [
|
||||||
|
{ name = "zeroentropy" },
|
||||||
|
]
|
||||||
|
|
||||||
[package.dev-dependencies]
|
[package.dev-dependencies]
|
||||||
dev = [
|
dev = [
|
||||||
|
|
@ -1197,6 +1203,7 @@ dev = [
|
||||||
|
|
||||||
[package.metadata]
|
[package.metadata]
|
||||||
requires-dist = [
|
requires-dist = [
|
||||||
|
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.0.0" },
|
||||||
{ name = "docling", specifier = ">=2.58.0" },
|
{ name = "docling", specifier = ">=2.58.0" },
|
||||||
{ name = "fasta2a", marker = "extra == 'a2a'", specifier = ">=0.1.0" },
|
{ name = "fasta2a", marker = "extra == 'a2a'", specifier = ">=0.1.0" },
|
||||||
{ name = "fastmcp", specifier = ">=2.13.0.2" },
|
{ name = "fastmcp", specifier = ">=2.13.0.2" },
|
||||||
|
|
@ -1214,8 +1221,9 @@ requires-dist = [
|
||||||
{ name = "typer", specifier = ">=0.19.2,<0.20.0" },
|
{ name = "typer", specifier = ">=0.19.2,<0.20.0" },
|
||||||
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.5" },
|
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.5" },
|
||||||
{ name = "watchfiles", specifier = ">=1.1.1" },
|
{ 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]
|
[package.metadata.requires-dev]
|
||||||
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" },
|
{ 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]]
|
[[package]]
|
||||||
name = "zipp"
|
name = "zipp"
|
||||||
version = "3.23.0"
|
version = "3.23.0"
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue