diff --git a/README.md b/README.md index da0f2dff..358091bc 100644 --- a/README.md +++ b/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 diff --git a/docs/configuration.md b/docs/configuration.md index 1a8acc31..d67e5bdf 100644 --- a/docs/configuration.md +++ b/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: diff --git a/docs/installation.md b/docs/installation.md index acd43c7b..36cf89d2 100644 --- a/docs/installation.md +++ b/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: diff --git a/pyproject.toml b/pyproject.toml index ff633475..ec89df98 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index 797aa0f7..f43e619e 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/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 diff --git a/src/haiku/rag/reranking/zeroentropy.py b/src/haiku/rag/reranking/zeroentropy.py new file mode 100644 index 00000000..0c66fb08 --- /dev/null +++ b/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 diff --git a/tests/test_reranker.py b/tests/test_reranker.py index cc93f584..2376dc47 100644 --- a/tests/test_reranker.py +++ b/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") diff --git a/uv.lock b/uv.lock index ac6d3943..b2c380c0 100644 --- a/uv.lock +++ b/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"