From 53c8289790410146a084eb594aa7714e84686114 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 10 Aug 2025 20:22:50 +0200 Subject: [PATCH 1/3] Update ollama --- pyproject.toml | 2 +- uv.lock | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c5fd668b..58f8b459 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -25,7 +25,7 @@ dependencies = [ "docling>=2.15.0", "fastmcp>=2.8.1", "httpx>=0.28.1", - "ollama>=0.5.1", + "ollama>=0.5.3", "pydantic>=2.11.7", "python-dotenv>=1.1.0", "rich>=14.0.0", diff --git a/uv.lock b/uv.lock index 64107a5d..eb47a20a 100644 --- a/uv.lock +++ b/uv.lock @@ -934,7 +934,7 @@ requires-dist = [ { name = "fastmcp", specifier = ">=2.8.1" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" }, - { name = "ollama", specifier = ">=0.5.1" }, + { name = "ollama", specifier = ">=0.5.3" }, { name = "openai", marker = "extra == 'openai'", specifier = ">=1.0.0" }, { name = "pydantic", specifier = ">=2.11.7" }, { name = "python-dotenv", specifier = ">=1.1.0" }, @@ -1829,15 +1829,15 @@ wheels = [ [[package]] name = "ollama" -version = "0.5.1" +version = "0.5.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "httpx" }, { name = "pydantic" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/8d/96/c7fe0d2d1b3053be614822a7b722c7465161b3672ce90df71515137580a0/ollama-0.5.1.tar.gz", hash = "sha256:5a799e4dc4e7af638b11e3ae588ab17623ee019e496caaf4323efbaa8feeff93", size = 41112, upload-time = "2025-05-30T21:32:48.679Z" } +sdist = { url = "https://files.pythonhosted.org/packages/91/6d/ae96027416dcc2e98c944c050c492789502d7d7c0b95a740f0bb39268632/ollama-0.5.3.tar.gz", hash = "sha256:40b6dff729df3b24e56d4042fd9d37e231cee8e528677e0d085413a1d6692394", size = 43331, upload-time = "2025-08-07T21:44:10.422Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d6/76/3f96c8cdbf3955d7a73ee94ce3e0db0755d6de1e0098a70275940d1aff2f/ollama-0.5.1-py3-none-any.whl", hash = "sha256:4c8839f35bc173c7057b1eb2cbe7f498c1a7e134eafc9192824c8aecb3617506", size = 13369, upload-time = "2025-05-30T21:32:47.429Z" }, + { url = "https://files.pythonhosted.org/packages/be/f6/2091e50b8b6c3e6901f6eab283d5efd66fb71c86ddb1b4d68766c3eeba0f/ollama-0.5.3-py3-none-any.whl", hash = "sha256:a8303b413d99a9043dbf77ebf11ced672396b59bec27e6d5db67c88f01b279d2", size = 13490, upload-time = "2025-08-07T21:44:09.353Z" }, ] [[package]] From 5103e6c2d79a0d778a179f6194862751344b63e7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 10 Aug 2025 20:50:17 +0200 Subject: [PATCH 2/3] Add Ollama reranker --- docs/configuration.md | 14 ++++- src/haiku/rag/config.py | 4 +- src/haiku/rag/reranking/__init__.py | 6 +++ src/haiku/rag/reranking/ollama.py | 84 +++++++++++++++++++++++++++++ tests/test_reranker.py | 15 +++++- 5 files changed, 118 insertions(+), 5 deletions(-) create mode 100644 src/haiku/rag/reranking/ollama.py diff --git a/docs/configuration.md b/docs/configuration.md index 5e78ee8a..3f42a369 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -107,9 +107,19 @@ ANTHROPIC_API_KEY="your-api-key" 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 **automatically enabled** if you install the appropriate reranking provider package. +Reranking is **automatically enabled** by default using Ollama, or if you install the appropriate reranking provider package. -### MixedBread AI (Default) +### 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" +``` + +### MixedBread AI For MxBAI reranking, install with mxbai extras: diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 251d4a46..e4732f48 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -19,8 +19,8 @@ class AppConfig(BaseModel): EMBEDDINGS_MODEL: str = "mxbai-embed-large" EMBEDDINGS_VECTOR_DIM: int = 1024 - RERANK_PROVIDER: str = "mxbai" - RERANK_MODEL: str = "mixedbread-ai/mxbai-rerank-base-v2" + RERANK_PROVIDER: str = "ollama" + RERANK_MODEL: str = "qwen3" 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 80df6bcb..449d35c3 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/haiku/rag/reranking/__init__.py @@ -35,4 +35,10 @@ 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 new file mode 100644 index 00000000..727c546b --- /dev/null +++ b/src/haiku/rag/reranking/ollama.py @@ -0,0 +1,84 @@ +import json + +from ollama import AsyncClient +from pydantic import BaseModel + +from haiku.rag.config import Config +from haiku.rag.reranking.base import RerankerBase +from haiku.rag.store.models.chunk import Chunk + +OLLAMA_OPTIONS = {"temperature": 0.0, "seed": 42, "num_ctx": 16384} + + +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 + self._client = AsyncClient(host=Config.OLLAMA_BASE_URL) + + 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}) + + # Create the prompt for reranking + 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.""" + + 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()} + +Please rank these documents by relevance to the query and return the top {top_n} results as JSON.""" + + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt}, + ] + + try: + response = await self._client.chat( + model=self._model, + messages=messages, + format=RerankResponse.model_json_schema(), + options=OLLAMA_OPTIONS, + ) + + content = response["message"]["content"] + + parsed_response = RerankResponse.model_validate(json.loads(content)) + return [ + (chunks[result.index], result.relevance_score) + for result in parsed_response.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 00e28b3b..5ce13156 100644 --- a/tests/test_reranker.py +++ b/tests/test_reranker.py @@ -21,7 +21,7 @@ chunks = [ @pytest.mark.asyncio async def test_reranker_base(): reranker = RerankerBase() - assert reranker._model == "mixedbread-ai/mxbai-rerank-base-v2" + assert reranker._model == "qwen3" with pytest.raises(NotImplementedError): await reranker.rerank("query", []) @@ -58,3 +58,16 @@ 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) From ffe867f7bb1a01980ecf45ffc5bba6eff5bc265b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sun, 10 Aug 2025 21:02:17 +0200 Subject: [PATCH 3/3] Allow to disable reranking by having an empty provider --- docs/configuration.md | 8 ++++++++ src/haiku/rag/reranking/__init__.py | 8 ++------ 2 files changed, 10 insertions(+), 6 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 3f42a369..c1507e49 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -109,6 +109,14 @@ Reranking improves search quality by re-ordering the initial search results usin Reranking is **automatically enabled** by default using Ollama, or if you install the appropriate reranking provider package. +### 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: diff --git a/src/haiku/rag/reranking/__init__.py b/src/haiku/rag/reranking/__init__.py index 449d35c3..e668968f 100644 --- a/src/haiku/rag/reranking/__init__.py +++ b/src/haiku/rag/reranking/__init__.py @@ -1,22 +1,18 @@ from haiku.rag.config import Config from haiku.rag.reranking.base import RerankerBase -try: - from haiku.rag.reranking.cohere import CohereReranker -except ImportError: - pass - _reranker: RerankerBase | None = None def get_reranker() -> RerankerBase | None: """ Factory function to get the appropriate reranker based on the configuration. - Returns None if the required package is not available. + Returns None if if reranking is disabled. """ global _reranker if _reranker is not None: return _reranker + if Config.RERANK_PROVIDER == "mxbai": try: from haiku.rag.reranking.mxbai import MxBAIReranker