diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 4650db4d..1b10b2b1 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -15,11 +15,11 @@ The recall obtained is ~0.79 for matching in the top result, raising to ~0.91 fo | Embedding Model | Document in top 1 | Document in top 3 | Reranker | |---------------------------------------|-------------------|-------------------|------------------------| | Ollama / `mxbai-embed-large` | 0.79 | 0.91 | None | -| Ollama / `mxbai-embed-large` | 0.81 | 0.91 | `mxbai-rerank-base-v2` | -| Ollama / `nomic-embed-text` | 0.74 | 0.88 | None | +| Ollama / `mxbai-embed-large` | 0.90 | 0.95 | `mxbai-rerank-base-v2` | + ## Question/Answer evaluation @@ -29,5 +29,5 @@ Again using the same dataset, we use a QA agent to answer the question. In addit |------------------------------------|-----------------------------------|-----------|------------------------| | Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.85 | None | | Ollama / `mxbai-embed-large` | Ollama / `qwen3` | 0.72 | `mxbai-rerank-base-v2` | -| Ollama / `mxbai-embed-large` | Anthropic / `Claude Sonnet 3.7` | 0.79 | None | -| OpenAI / `text-embeddings-3-small` | OpenAI / `gpt-4-turbo` | 0.62 | None | + diff --git a/docs/configuration.md b/docs/configuration.md index bfbd200a..e02ae2d2 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -109,25 +109,7 @@ See the [Pydantic AI documentation](https://ai.pydantic.dev/models/) for the com 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. -Reranking is **disabled by default** for faster searches. You can enable it by configuring a reranking provider. - -### Disabling Reranking - -To disable reranking completely for faster searches: - -```bash -RERANK_PROVIDER="" -``` - -### Ollama (Default) - -Ollama reranking uses LLMs with structured output to rank documents by relevance: - -```bash -RERANK_PROVIDER="ollama" -RERANK_MODEL="qwen3:1.7b" # or any model that supports structured output -OLLAMA_BASE_URL="http://localhost:11434" -``` +Reranking is **disabled by default** (`RERANK_PROVIDER=""`) for faster searches. You can enable it by configuring one of the providers below. ### MixedBread AI diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 31eee040..8910fab8 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -20,7 +20,7 @@ class AppConfig(BaseModel): EMBEDDINGS_VECTOR_DIM: int = 1024 RERANK_PROVIDER: str = "" - RERANK_MODEL: str = "qwen3" + RERANK_MODEL: str = "" QA_PROVIDER: str = "ollama" QA_MODEL: str = "qwen3" diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index e668968f..f63453c6 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/haiku/rag/reranking/__init__.py @@ -31,10 +31,4 @@ def get_reranker() -> RerankerBase | None: except ImportError: return None - if Config.RERANK_PROVIDER == "ollama": - from haiku.rag.reranking.ollama import OllamaReranker - - _reranker = OllamaReranker() - return _reranker - return None diff --git a/src/haiku/rag/reranking/ollama.py b/src/haiku/rag/reranking/ollama.py deleted file mode 100644 index 9acba3dd..00000000 --- a/src/haiku/rag/reranking/ollama.py +++ /dev/null @@ -1,81 +0,0 @@ -from pydantic import BaseModel -from pydantic_ai import Agent -from pydantic_ai.models.openai import OpenAIModel -from pydantic_ai.providers.ollama import OllamaProvider - -from haiku.rag.config import Config -from haiku.rag.reranking.base import RerankerBase -from haiku.rag.store.models.chunk import Chunk - - -class RerankResult(BaseModel): - """Individual rerank result with index and relevance score.""" - - index: int - relevance_score: float - - -class RerankResponse(BaseModel): - """Response from the reranking model containing ranked results.""" - - results: list[RerankResult] - - -class OllamaReranker(RerankerBase): - def __init__(self, model: str = Config.RERANK_MODEL): - self._model = model - - # Create the reranking prompt - system_prompt = """You are a document reranking assistant. Given a query and a list of document chunks, you must rank them by relevance to the query. - -Return your response as a JSON object with a "results" array. Each result should have: -- "index": the original index of the document (integer) -- "relevance_score": a score between 0.0 and 1.0 indicating relevance (float, where 1.0 is most relevant) - -Only return the top documents up to the requested limit, ordered by decreasing relevance score. -/no_think -""" - - model_obj = OpenAIModel( - model_name=model, - provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), - ) - - self._agent = Agent( - model=model_obj, - output_type=RerankResponse, - system_prompt=system_prompt, - ) - - async def rerank( - self, query: str, chunks: list[Chunk], top_n: int = 10 - ) -> list[tuple[Chunk, float]]: - if not chunks: - return [] - - documents = [] - for i, chunk in enumerate(chunks): - documents.append({"index": i, "content": chunk.content}) - - documents_text = "" - for doc in documents: - documents_text += f"Index {doc['index']}: {doc['content']}\n\n" - - user_prompt = f"""Query: {query} - -Documents to rerank: -{documents_text.strip()} - -Rank these documents by relevance to the query and return the top {top_n} results as JSON.""" - - try: - result = await self._agent.run(user_prompt) - - return [ - (chunks[result_item.index], result_item.relevance_score) - for result_item in result.output.results[:top_n] - ] - - except Exception: - # Fallback: return chunks in original order with same score - return [(chunks[i], 1.0) for i in range(min(top_n, len(chunks)))] diff --git a/tests/test_reranker.py b/tests/test_reranker.py index 3cad4f0f..5e309c0d 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -24,7 +24,7 @@ chunks = [ @pytest.mark.asyncio async def test_reranker_base(): reranker = RerankerBase() - assert reranker._model == "qwen3" + assert reranker._model == "" with pytest.raises(NotImplementedError): await reranker.rerank("query", []) @@ -35,12 +35,16 @@ async def test_mxbai_reranker(): try: from haiku.rag.reranking.mxbai import MxBAIReranker + Config.RERANK_MODEL = "mixedbread-ai/mxbai-rerank-base-v2" reranker = MxBAIReranker() + # reranker._model = "mixedbread-ai/mxbai-rerank-base-v2" reranked = await reranker.rerank( "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) + Config.RERANK_MODEL = "" + except ImportError: pytest.skip("MxBAI package not installed") @@ -62,16 +66,3 @@ async def test_cohere_reranker(): except ImportError: pytest.skip("Cohere package not installed") - - -@pytest.mark.asyncio -async def test_ollama_reranker(): - from haiku.rag.reranking.ollama import OllamaReranker - - reranker = OllamaReranker() - reranked = await reranker.rerank( - "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)