Make mxbai reranker optional. Only use a reranker if mxbai or cohere are installed

This commit is contained in:
Yiorgis Gozadinos 2025-08-10 19:49:14 +02:00
parent ff42224de3
commit 8d98eed2e2
No known key found for this signature in database
7 changed files with 45 additions and 34 deletions

View file

@ -105,16 +105,20 @@ ANTHROPIC_API_KEY="your-api-key"
## Reranking ## Reranking
Reranking is **enabled by default** and improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results. Reranking improves search quality by re-ordering the initial search results using specialized models. When enabled, the system retrieves more candidates (3x the requested limit) and then reranks them to return the most relevant results.
If you use the default reranked (running locally), it can slow down searching significantly. To disable reranking for faster searches: Reranking is **automatically enabled** if you install the appropriate reranking provider package.
```bash
RERANK=false
```
### MixedBread AI (Default) ### MixedBread AI (Default)
For MxBAI reranking, install with mxbai extras:
```bash
uv pip install haiku.rag[mxbai]
```
Then configure:
```bash ```bash
RERANK_PROVIDER="mxbai" RERANK_PROVIDER="mxbai"
RERANK_MODEL="mixedbread-ai/mxbai-rerank-base-v2" RERANK_MODEL="mixedbread-ai/mxbai-rerank-base-v2"

View file

@ -25,7 +25,6 @@ dependencies = [
"docling>=2.15.0", "docling>=2.15.0",
"fastmcp>=2.8.1", "fastmcp>=2.8.1",
"httpx>=0.28.1", "httpx>=0.28.1",
"mxbai-rerank>=0.1.6",
"ollama>=0.5.1", "ollama>=0.5.1",
"pydantic>=2.11.7", "pydantic>=2.11.7",
"python-dotenv>=1.1.0", "python-dotenv>=1.1.0",
@ -41,6 +40,7 @@ voyageai = ["voyageai>=0.3.2"]
openai = ["openai>=1.0.0"] openai = ["openai>=1.0.0"]
anthropic = ["anthropic>=0.56.0"] anthropic = ["anthropic>=0.56.0"]
cohere = ["cohere>=5.16.1"] cohere = ["cohere>=5.16.1"]
mxbai = ["mxbai-rerank>=0.1.6"]
[project.scripts] [project.scripts]
haiku-rag = "haiku.rag.cli:cli" haiku-rag = "haiku.rag.cli:cli"

View file

@ -319,7 +319,7 @@ class HaikuRAG:
return await self.document_repository.list_all(limit=limit, offset=offset) return await self.document_repository.list_all(limit=limit, offset=offset)
async def search( async def search(
self, query: str, limit: int = 5, k: int = 60, rerank=Config.RERANK self, query: str, limit: int = 5, k: int = 60
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]:
"""Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking. """Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking.
@ -331,8 +331,10 @@ class HaikuRAG:
Returns: Returns:
List of (chunk, score) tuples ordered by relevance. List of (chunk, score) tuples ordered by relevance.
""" """
# Get reranker if available
reranker = get_reranker()
if not rerank: if reranker is None:
return await self.chunk_repository.search_chunks_hybrid(query, limit, k) return await self.chunk_repository.search_chunks_hybrid(query, limit, k)
# Get more initial results (3X) for reranking # Get more initial results (3X) for reranking
@ -340,7 +342,6 @@ class HaikuRAG:
query, limit * 3, k query, limit * 3, k
) )
# Apply reranking # Apply reranking
reranker = get_reranker()
chunks = [chunk for chunk, _ in search_results] chunks = [chunk for chunk, _ in search_results]
reranked_results = await reranker.rerank(query, chunks, top_n=limit) reranked_results = await reranker.rerank(query, chunks, top_n=limit)

View file

@ -19,7 +19,6 @@ class AppConfig(BaseModel):
EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_MODEL: str = "mxbai-embed-large"
EMBEDDINGS_VECTOR_DIM: int = 1024 EMBEDDINGS_VECTOR_DIM: int = 1024
RERANK: bool = True
RERANK_PROVIDER: str = "mxbai" RERANK_PROVIDER: str = "mxbai"
RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2"

View file

@ -9,29 +9,30 @@ except ImportError:
_reranker: RerankerBase | None = None _reranker: RerankerBase | None = None
def get_reranker() -> RerankerBase: def get_reranker() -> RerankerBase | None:
""" """
Factory function to get the appropriate reranker based on the configuration. Factory function to get the appropriate reranker based on the configuration.
Returns None if the required package is not available.
""" """
global _reranker global _reranker
if _reranker is not None: if _reranker is not None:
return _reranker return _reranker
if Config.RERANK_PROVIDER == "mxbai": if Config.RERANK_PROVIDER == "mxbai":
from haiku.rag.reranking.mxbai import MxBAIReranker try:
from haiku.rag.reranking.mxbai import MxBAIReranker
_reranker = MxBAIReranker() _reranker = MxBAIReranker()
return _reranker return _reranker
except ImportError:
return None
if Config.RERANK_PROVIDER == "cohere": if Config.RERANK_PROVIDER == "cohere":
try: try:
from haiku.rag.reranking.cohere import CohereReranker from haiku.rag.reranking.cohere import CohereReranker
except ImportError:
raise ImportError(
"Cohere reranker requires the 'cohere' package. "
"Please install haiku.rag with the 'cohere' extra:"
"uv pip install haiku.rag[cohere]"
)
_reranker = CohereReranker()
return _reranker
raise ValueError(f"Unsupported reranker provider: {Config.RERANK_PROVIDER}") _reranker = CohereReranker()
return _reranker
except ImportError:
return None
return None

View file

@ -1,7 +1,6 @@
import pytest import pytest
from haiku.rag.reranking.base import RerankerBase from haiku.rag.reranking.base import RerankerBase
from haiku.rag.reranking.mxbai import MxBAIReranker
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
chunks = [ chunks = [
@ -30,12 +29,17 @@ async def test_reranker_base():
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_mxbai_reranker(): async def test_mxbai_reranker():
reranker = MxBAIReranker() try:
reranked = await reranker.rerank( from haiku.rag.reranking.mxbai import MxBAIReranker
"Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
) reranker = MxBAIReranker()
assert [chunk.document_id for chunk, score in reranked] == [0, 2] reranked = await reranker.rerank(
assert all(isinstance(score, float) for chunk, score in reranked) "Who wrote 'To Kill a Mockingbird'?", chunks, top_n=2
)
assert [chunk.document_id for chunk, score in reranked] == [0, 2]
assert all(isinstance(score, float) for chunk, score in reranked)
except ImportError:
pytest.skip("MxBAI package not installed")
@pytest.mark.asyncio @pytest.mark.asyncio

View file

@ -886,7 +886,6 @@ dependencies = [
{ name = "docling" }, { name = "docling" },
{ name = "fastmcp" }, { name = "fastmcp" },
{ name = "httpx" }, { name = "httpx" },
{ name = "mxbai-rerank" },
{ name = "ollama" }, { name = "ollama" },
{ name = "pydantic" }, { name = "pydantic" },
{ name = "python-dotenv" }, { name = "python-dotenv" },
@ -904,6 +903,9 @@ anthropic = [
cohere = [ cohere = [
{ name = "cohere" }, { name = "cohere" },
] ]
mxbai = [
{ name = "mxbai-rerank" },
]
openai = [ openai = [
{ name = "openai" }, { name = "openai" },
] ]
@ -931,7 +933,7 @@ requires-dist = [
{ name = "docling", specifier = ">=2.15.0" }, { name = "docling", specifier = ">=2.15.0" },
{ name = "fastmcp", specifier = ">=2.8.1" }, { name = "fastmcp", specifier = ">=2.8.1" },
{ name = "httpx", specifier = ">=0.28.1" }, { name = "httpx", specifier = ">=0.28.1" },
{ name = "mxbai-rerank", specifier = ">=0.1.6" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" },
{ name = "ollama", specifier = ">=0.5.1" }, { name = "ollama", specifier = ">=0.5.1" },
{ name = "openai", marker = "extra == 'openai'", specifier = ">=1.0.0" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=1.0.0" },
{ name = "pydantic", specifier = ">=2.11.7" }, { name = "pydantic", specifier = ">=2.11.7" },
@ -943,7 +945,7 @@ requires-dist = [
{ name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" }, { name = "voyageai", marker = "extra == 'voyageai'", specifier = ">=0.3.2" },
{ name = "watchfiles", specifier = ">=1.1.0" }, { name = "watchfiles", specifier = ">=1.1.0" },
] ]
provides-extras = ["voyageai", "openai", "anthropic", "cohere"] provides-extras = ["voyageai", "openai", "anthropic", "cohere", "mxbai"]
[package.metadata.requires-dev] [package.metadata.requires-dev]
dev = [ dev = [