haiku.rag 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

haiku/rag/app.py

@@ -32,9 +32,9 @@ class HaikuRAGApp:
                 f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"
             )
 
-    async def add_document_from_source(self, file_path: Path):
+    async def add_document_from_source(self, source: str):
         async with HaikuRAG(db_path=self.db_path) as self.client:
-            doc = await self.client.create_document_from_source(file_path)
+            doc = await self.client.create_document_from_source(source)
             self._rich_print_document(doc, truncate=True)
             self.console.print(
                 f"[b]Document with id [cyan]{doc.id}[/cyan] added successfully.[/b]"

haiku/rag/cli.py

@@ -81,7 +81,7 @@ def add_document_text(
 
 @cli.command("add-src", help="Add a document from a file path or URL")
 def add_document_src(
-    file_path: Path = typer.Argument(
+    source: str = typer.Argument(
         help="The file path or URL of the document to add",
     ),
     db: Path = typer.Option(
@@ -91,7 +91,7 @@ def add_document_src(
     ),
 ):
     app = HaikuRAGApp(db_path=db)
-    asyncio.run(app.add_document_from_source(file_path=file_path))
+    asyncio.run(app.add_document_from_source(source=source))
 
 
 @cli.command("get", help="Get and display a document by its ID")

haiku/rag/client.py

@@ -319,7 +319,7 @@ class HaikuRAG:
         return await self.document_repository.list_all(limit=limit, offset=offset)
 
     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]]:
         """Search for relevant chunks using hybrid search (vector similarity + full-text search) with reranking.
 
@@ -331,8 +331,10 @@ class HaikuRAG:
         Returns:
             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)
 
         # Get more initial results (3X) for reranking
@@ -340,7 +342,6 @@ class HaikuRAG:
             query, limit * 3, k
         )
         # Apply reranking
-        reranker = get_reranker()
         chunks = [chunk for chunk, _ in search_results]
         reranked_results = await reranker.rerank(query, chunks, top_n=limit)
 

haiku/rag/config.py

@@ -19,9 +19,8 @@ class AppConfig(BaseModel):
     EMBEDDINGS_MODEL: str = "mxbai-embed-large"
     EMBEDDINGS_VECTOR_DIM: int = 1024
 
-    RERANK: bool = True
-    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"

haiku/rag/reranking/__init__.py

@@ -1,37 +1,40 @@
 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:
+def get_reranker() -> RerankerBase | None:
     """
     Factory function to get the appropriate reranker based on the configuration.
+    Returns None if if reranking is disabled.
     """
     global _reranker
     if _reranker is not None:
         return _reranker
+
     if Config.RERANK_PROVIDER == "mxbai":
-        from haiku.rag.reranking.mxbai import MxBAIReranker
+        try:
+            from haiku.rag.reranking.mxbai import MxBAIReranker
 
-        _reranker = MxBAIReranker()
-        return _reranker
+            _reranker = MxBAIReranker()
+            return _reranker
+        except ImportError:
+            return None
 
     if Config.RERANK_PROVIDER == "cohere":
         try:
             from haiku.rag.reranking.cohere import CohereReranker
+
+            _reranker = CohereReranker()
+            return _reranker
         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 None
+
+    if Config.RERANK_PROVIDER == "ollama":
+        from haiku.rag.reranking.ollama import OllamaReranker
+
+        _reranker = OllamaReranker()
         return _reranker
 
-    raise ValueError(f"Unsupported reranker provider: {Config.RERANK_PROVIDER}")
+    return None

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)))]

{haiku_rag-0.5.1.dist-info → haiku_rag-0.5.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.5.1
+Version: 0.5.2
 Summary: Retrieval Augmented Generation (RAG) with SQLite
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -21,8 +21,7 @@ Requires-Python: >=3.11
 Requires-Dist: docling>=2.15.0
 Requires-Dist: fastmcp>=2.8.1
 Requires-Dist: httpx>=0.28.1
-Requires-Dist: mxbai-rerank>=0.1.6
-Requires-Dist: ollama>=0.5.1
+Requires-Dist: ollama>=0.5.3
 Requires-Dist: pydantic>=2.11.7
 Requires-Dist: python-dotenv>=1.1.0
 Requires-Dist: rich>=14.0.0
@@ -34,6 +33,8 @@ Provides-Extra: anthropic
 Requires-Dist: anthropic>=0.56.0; extra == 'anthropic'
 Provides-Extra: cohere
 Requires-Dist: cohere>=5.16.1; extra == 'cohere'
+Provides-Extra: mxbai
+Requires-Dist: mxbai-rerank>=0.1.6; extra == 'mxbai'
 Provides-Extra: openai
 Requires-Dist: openai>=1.0.0; extra == 'openai'
 Provides-Extra: voyageai

{haiku_rag-0.5.1.dist-info → haiku_rag-0.5.2.dist-info}/RECORD

@@ -1,9 +1,9 @@
 haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-haiku/rag/app.py,sha256=FpLVyP1-zAq_XPmU8CPVLkuIAeuhBOGvMqhYS8RbN40,7649
+haiku/rag/app.py,sha256=kuvULOIdgwqzJMaKtb1znStc1YAqB1-RkZ0fwdg6TBk,7642
 haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
-haiku/rag/cli.py,sha256=rk4uUwN_FdMC-rai9_R2sgXXMI3TIWKRtdWWHg_WoWM,5865
-haiku/rag/client.py,sha256=pFcrPkQo1h1zJ76jts-72goP_kGVtnJNfLuoT8qpsb8,15795
-haiku/rag/config.py,sha256=8mlQ8gYFxxq1q9gi9tjY9StjqhfhiHkO1FvS4b0et0E,1633
+haiku/rag/cli.py,sha256=5CcWcBQ47KCZ1wl7DpLzMXtgJZ1nz5Hci8AYp72oXEI,5855
+haiku/rag/client.py,sha256=K51l_orUc3BeKGTHX4xC7YClY9M4Eijpac5Hg1_q6LE,15815
+haiku/rag/config.py,sha256=jiy5vg-YbYa7yY-834Dd9omFtfMBXQBYXmHRJXMPjrs,1581
 haiku/rag/logging.py,sha256=zTTGpGq5tPdcd7RpCbd9EGw1IZlQDbYkrCg9t9pqRc4,580
 haiku/rag/mcp.py,sha256=tMN6fNX7ZtAER1R6DL1GkC9HZozTC4HzuQs199p7icI,4551
 haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
@@ -20,10 +20,11 @@ haiku/rag/qa/base.py,sha256=4ZTM_l5FAZ9cA0f8NeqRJiUAmjatwCTmSoclFw0gTFQ,1349
 haiku/rag/qa/ollama.py,sha256=EGUi4urSx9nrnsr5j-qHVDVOnvRTbSMKUbMvXEMIcxM,2381
 haiku/rag/qa/openai.py,sha256=dF32sGgVt8mZi5oVxByaeECs9NqLjvDiZnnpJBsrHm8,3968
 haiku/rag/qa/prompts.py,sha256=8uYMxHzbzI9vo2FPkCSSNTh_RNL96WkBbUWPCMBlLpo,1315
-haiku/rag/reranking/__init__.py,sha256=DsPCdU94wRzDCYl6hz2DySOMWwOvNxKviqKAUfyykK8,1118
+haiku/rag/reranking/__init__.py,sha256=fwC3pauteJwh9Ulm2270QvwAdwr4NMr4RUEuolC-wKU,1063
 haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
 haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
 haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
+haiku/rag/reranking/ollama.py,sha256=tCrLlNNDBCZu7J3to1gvBq-sOvN1flYEA7E3H3Jq0mU,2790
 haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
 haiku/rag/store/engine.py,sha256=cOMBToLilI1Di1qQrFzGLqtRMsuvtiX0Q5RNIEzQy9w,6232
 haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
@@ -36,8 +37,8 @@ haiku/rag/store/repositories/document.py,sha256=ki8LiDukwU1469Yw51i0rQFvBzUQeYkF
 haiku/rag/store/repositories/settings.py,sha256=qZLXvLsErnCWL0nBQQNfRnatHzCKhtUDLvUK9k-W_fU,2463
 haiku/rag/store/upgrades/__init__.py,sha256=kKS1YWT_P-CYKhKtokOLTIFNKf9jlfjFFr8lyIMeogM,100
 haiku/rag/store/upgrades/v0_3_4.py,sha256=GLogKZdZ40NX1vBHKdOJju7fFzNUCHoEnjSZg17Hm2U,663
-haiku_rag-0.5.1.dist-info/METADATA,sha256=X4r-1CBCTef3_T9HWPgCHi5XumqOSF4tlHfUpxO533E,4198
-haiku_rag-0.5.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.5.1.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.5.1.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.5.1.dist-info/RECORD,,
+haiku_rag-0.5.2.dist-info/METADATA,sha256=b91HARmgPKSy_4LIhna9EoacKb9I_f-cRRTgHqaG-S8,4238
+haiku_rag-0.5.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.5.2.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.5.2.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.5.2.dist-info/RECORD,,