haiku.rag 0.5.1__tar.gz → 0.5.4__tar.gz

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.5.1
+Version: 0.5.4
 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
@@ -75,6 +76,9 @@ haiku-rag search "query"
 # Ask questions
 haiku-rag ask "Who is the author of haiku.rag?"
 
+# Ask questions with citations
+haiku-rag ask "Who is the author of haiku.rag?" --cite
+
 # Rebuild database (re-chunk and re-embed all documents)
 haiku-rag rebuild
 
@@ -100,6 +104,10 @@ async with HaikuRAG("database.db") as client:
     # Ask questions
     answer = await client.ask("Who is the author of haiku.rag?")
     print(answer)
+
+    # Ask questions with citations
+    answer = await client.ask("Who is the author of haiku.rag?", cite=True)
+    print(answer)
 ```
 
 ## MCP Server

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/README.md

@@ -33,6 +33,9 @@ haiku-rag search "query"
 # Ask questions
 haiku-rag ask "Who is the author of haiku.rag?"
 
+# Ask questions with citations
+haiku-rag ask "Who is the author of haiku.rag?" --cite
+
 # Rebuild database (re-chunk and re-embed all documents)
 haiku-rag rebuild
 
@@ -58,6 +61,10 @@ async with HaikuRAG("database.db") as client:
     # Ask questions
     answer = await client.ask("Who is the author of haiku.rag?")
     print(answer)
+
+    # Ask questions with citations
+    answer = await client.ask("Who is the author of haiku.rag?", cite=True)
+    print(answer)
 ```
 
 ## MCP Server

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/docs/cli.md

@@ -64,7 +64,12 @@ Ask questions about your documents:
 haiku-rag ask "Who is the author of haiku.rag?"
 ```
 
-The QA agent will search your documents for relevant information and provide a comprehensive answer.
+Ask questions with citations showing source documents:
+```bash
+haiku-rag ask "Who is the author of haiku.rag?" --cite
+```
+
+The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used.
 
 ## Configuration
 

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/docs/configuration.md

@@ -105,15 +105,37 @@ ANTHROPIC_API_KEY="your-api-key"
 
 ## 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** 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:
 
 ```bash
-RERANK=false
+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:
+
+```bash
+uv pip install haiku.rag[mxbai]
 ```
 
-### MixedBread AI (Default)
+Then configure:
 
 ```bash
 RERANK_PROVIDER="mxbai"
@@ -150,4 +172,10 @@ DEFAULT_DATA_DIR="/path/to/data"
 ```bash
 # Chunk size for document processing
 CHUNK_SIZE=256
+
+# Number of adjacent chunks to include before/after retrieved chunks for context
+# 0 = no expansion (default), 1 = include 1 chunk before and after, etc.
+# When expanded chunks overlap or are adjacent, they are automatically merged
+# into single chunks with continuous content to eliminate duplication
+CONTEXT_CHUNK_RADIUS=0
 ```

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/docs/python.md

@@ -130,6 +130,26 @@ for chunk, relevance_score in results:
     print(f"Document metadata: {chunk.document_meta}")
 ```
 
+### Expanding Search Context
+
+Expand search results with adjacent chunks for more complete context:
+
+```python
+# Get initial search results
+search_results = await client.search("machine learning", limit=3)
+
+# Expand with adjacent chunks based on CONTEXT_CHUNK_RADIUS setting
+expanded_results = await client.expand_context(search_results)
+
+# The expanded results contain chunks with combined content from adjacent chunks
+for chunk, score in expanded_results:
+    print(f"Expanded content: {chunk.content}")  # Now includes before/after chunks
+```
+
+**Smart Merging**: When expanded chunks overlap or are adjacent within the same document, they are automatically merged into single chunks with continuous content. This eliminates duplication and provides coherent text blocks. The merged chunk uses the highest relevance score from the original chunks.
+
+This is automatically used by the QA system when `CONTEXT_CHUNK_RADIUS > 0` to provide better answers with more complete context.
+
 ## Question Answering
 
 Ask questions about your documents:
@@ -139,6 +159,13 @@ answer = await client.ask("Who is the author of haiku.rag?")
 print(answer)
 ```
 
-The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer.
+Ask questions with citations showing source documents:
+
+```python
+answer = await client.ask("Who is the author of haiku.rag?", cite=True)
+print(answer)
+```
+
+The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources.
 
 The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.5.1"
+version = "0.5.4"
 description = "Retrieval Augmented Generation (RAG) with SQLite"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
@@ -25,8 +25,7 @@ dependencies = [
     "docling>=2.15.0",
     "fastmcp>=2.8.1",
     "httpx>=0.28.1",
-    "mxbai-rerank>=0.1.6",
-    "ollama>=0.5.1",
+    "ollama>=0.5.3",
     "pydantic>=2.11.7",
     "python-dotenv>=1.1.0",
     "rich>=14.0.0",
@@ -41,6 +40,7 @@ voyageai = ["voyageai>=0.3.2"]
 openai = ["openai>=1.0.0"]
 anthropic = ["anthropic>=0.56.0"]
 cohere = ["cohere>=5.16.1"]
+mxbai = ["mxbai-rerank>=0.1.6"]
 
 [project.scripts]
 haiku-rag = "haiku.rag.cli:cli"

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/src/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]"
@@ -62,10 +62,10 @@ class HaikuRAGApp:
             for chunk, score in results:
                 self._rich_print_search_result(chunk, score)
 
-    async def ask(self, question: str):
+    async def ask(self, question: str, cite: bool = False):
         async with HaikuRAG(db_path=self.db_path) as self.client:
             try:
-                answer = await self.client.ask(question)
+                answer = await self.client.ask(question, cite=cite)
                 self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                 self.console.print()
                 self.console.print("[bold green]Answer:[/bold green]")

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/src/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")
@@ -160,9 +160,14 @@ def ask(
         "--db",
         help="Path to the SQLite database file",
     ),
+    cite: bool = typer.Option(
+        False,
+        "--cite",
+        help="Include citations in the response",
+    ),
 ):
     app = HaikuRAGApp(db_path=db)
-    asyncio.run(app.ask(question=question))
+    asyncio.run(app.ask(question=question, cite=cite))
 
 
 @cli.command("settings", help="Display current configuration settings")

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/src/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,25 +342,151 @@ 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)
 
         # Return reranked results with scores from reranker
         return reranked_results
 
-    async def ask(self, question: str) -> str:
+    async def expand_context(
+        self, search_results: list[tuple[Chunk, float]]
+    ) -> list[tuple[Chunk, float]]:
+        """Expand search results with adjacent chunks, merging overlapping chunks.
+
+        Args:
+            search_results: List of (chunk, score) tuples from search.
+
+        Returns:
+            List of (chunk, score) tuples with expanded and merged context chunks.
+        """
+        if Config.CONTEXT_CHUNK_RADIUS == 0:
+            return search_results
+
+        # Group chunks by document_id to handle merging within documents
+        document_groups = {}
+        for chunk, score in search_results:
+            doc_id = chunk.document_id
+            if doc_id not in document_groups:
+                document_groups[doc_id] = []
+            document_groups[doc_id].append((chunk, score))
+
+        results = []
+
+        for doc_id, doc_chunks in document_groups.items():
+            # Get all expanded ranges for this document
+            expanded_ranges = []
+            for chunk, score in doc_chunks:
+                adjacent_chunks = await self.chunk_repository.get_adjacent_chunks(
+                    chunk, Config.CONTEXT_CHUNK_RADIUS
+                )
+
+                all_chunks = adjacent_chunks + [chunk]
+
+                # Get the range of orders for this expanded chunk
+                orders = [c.metadata.get("order", 0) for c in all_chunks]
+                min_order = min(orders)
+                max_order = max(orders)
+
+                expanded_ranges.append(
+                    {
+                        "original_chunk": chunk,
+                        "score": score,
+                        "min_order": min_order,
+                        "max_order": max_order,
+                        "all_chunks": sorted(
+                            all_chunks, key=lambda c: c.metadata.get("order", 0)
+                        ),
+                    }
+                )
+
+            # Merge overlapping/adjacent ranges
+            merged_ranges = self._merge_overlapping_ranges(expanded_ranges)
+
+            # Create merged chunks
+            for merged_range in merged_ranges:
+                combined_content_parts = [c.content for c in merged_range["all_chunks"]]
+
+                # Use the first original chunk for metadata
+                original_chunk = merged_range["original_chunks"][0]
+
+                merged_chunk = Chunk(
+                    id=original_chunk.id,
+                    document_id=original_chunk.document_id,
+                    content="".join(combined_content_parts),
+                    metadata=original_chunk.metadata,
+                    document_uri=original_chunk.document_uri,
+                    document_meta=original_chunk.document_meta,
+                )
+
+                # Use the highest score from merged chunks
+                best_score = max(merged_range["scores"])
+                results.append((merged_chunk, best_score))
+
+        return results
+
+    def _merge_overlapping_ranges(self, expanded_ranges):
+        """Merge overlapping or adjacent expanded ranges."""
+        if not expanded_ranges:
+            return []
+
+        # Sort by min_order
+        sorted_ranges = sorted(expanded_ranges, key=lambda x: x["min_order"])
+        merged = []
+
+        current = {
+            "min_order": sorted_ranges[0]["min_order"],
+            "max_order": sorted_ranges[0]["max_order"],
+            "original_chunks": [sorted_ranges[0]["original_chunk"]],
+            "scores": [sorted_ranges[0]["score"]],
+            "all_chunks": sorted_ranges[0]["all_chunks"],
+        }
+
+        for range_info in sorted_ranges[1:]:
+            # Check if ranges overlap or are adjacent (max_order + 1 >= min_order)
+            if current["max_order"] >= range_info["min_order"] - 1:
+                # Merge ranges
+                current["max_order"] = max(
+                    current["max_order"], range_info["max_order"]
+                )
+                current["original_chunks"].append(range_info["original_chunk"])
+                current["scores"].append(range_info["score"])
+
+                # Merge all_chunks and deduplicate by order
+                all_chunks_dict = {}
+                for chunk in current["all_chunks"] + range_info["all_chunks"]:
+                    order = chunk.metadata.get("order", 0)
+                    all_chunks_dict[order] = chunk
+                current["all_chunks"] = [
+                    all_chunks_dict[order] for order in sorted(all_chunks_dict.keys())
+                ]
+            else:
+                # No overlap, add current to merged and start new
+                merged.append(current)
+                current = {
+                    "min_order": range_info["min_order"],
+                    "max_order": range_info["max_order"],
+                    "original_chunks": [range_info["original_chunk"]],
+                    "scores": [range_info["score"]],
+                    "all_chunks": range_info["all_chunks"],
+                }
+
+        # Add the last range
+        merged.append(current)
+        return merged
+
+    async def ask(self, question: str, cite: bool = False) -> str:
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
+            cite: Whether to include citations in the response.
 
         Returns:
             The generated answer as a string.
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self)
+        qa_agent = get_qa_agent(self, use_citations=cite)
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[int, None]:

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/src/haiku/rag/config.py

@@ -19,14 +19,14 @@ 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"
 
     CHUNK_SIZE: int = 256
+    CONTEXT_CHUNK_RADIUS: int = 0
 
     OLLAMA_BASE_URL: str = "http://localhost:11434"
 

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/src/haiku/rag/qa/__init__.py

@@ -4,12 +4,16 @@ from haiku.rag.qa.base import QuestionAnswerAgentBase
 from haiku.rag.qa.ollama import QuestionAnswerOllamaAgent
 
 
-def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase:
+def get_qa_agent(
+    client: HaikuRAG, model: str = "", use_citations: bool = False
+) -> QuestionAnswerAgentBase:
     """
     Factory function to get the appropriate QA agent based on the configuration.
     """
     if Config.QA_PROVIDER == "ollama":
-        return QuestionAnswerOllamaAgent(client, model or Config.QA_MODEL)
+        return QuestionAnswerOllamaAgent(
+            client, model or Config.QA_MODEL, use_citations
+        )
 
     if Config.QA_PROVIDER == "openai":
         try:
@@ -20,7 +24,9 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase:
                 "Please install haiku.rag with the 'openai' extra:"
                 "uv pip install haiku.rag[openai]"
             )
-        return QuestionAnswerOpenAIAgent(client, model or Config.QA_MODEL)
+        return QuestionAnswerOpenAIAgent(
+            client, model or Config.QA_MODEL, use_citations
+        )
 
     if Config.QA_PROVIDER == "anthropic":
         try:
@@ -31,6 +37,8 @@ def get_qa_agent(client: HaikuRAG, model: str = "") -> QuestionAnswerAgentBase:
                 "Please install haiku.rag with the 'anthropic' extra:"
                 "uv pip install haiku.rag[anthropic]"
             )
-        return QuestionAnswerAnthropicAgent(client, model or Config.QA_MODEL)
+        return QuestionAnswerAnthropicAgent(
+            client, model or Config.QA_MODEL, use_citations
+        )
 
     raise ValueError(f"Unsupported QA provider: {Config.QA_PROVIDER}")

{haiku_rag-0.5.1 → haiku_rag-0.5.4}/src/haiku/rag/qa/anthropic.py

@@ -1,19 +1,29 @@
 from collections.abc import Sequence
 
 try:
-    from anthropic import AsyncAnthropic
-    from anthropic.types import MessageParam, TextBlock, ToolParam, ToolUseBlock
+    from anthropic import AsyncAnthropic  # type: ignore
+    from anthropic.types import (  # type: ignore
+        MessageParam,
+        TextBlock,
+        ToolParam,
+        ToolUseBlock,
+    )
 
     from haiku.rag.client import HaikuRAG
     from haiku.rag.qa.base import QuestionAnswerAgentBase
 
     class QuestionAnswerAnthropicAgent(QuestionAnswerAgentBase):
-        def __init__(self, client: HaikuRAG, model: str = "claude-3-5-haiku-20241022"):
-            super().__init__(client, model or self._model)
+        def __init__(
+            self,
+            client: HaikuRAG,
+            model: str = "claude-3-5-haiku-20241022",
+            use_citations: bool = False,
+        ):
+            super().__init__(client, model or self._model, use_citations)
             self.tools: Sequence[ToolParam] = [
                 ToolParam(
                     name="search_documents",
-                    description="Search the knowledge base for relevant documents",
+                    description="Search the knowledge base for relevant documents. Returns a JSON array with content, score, and document_uri for each result.",
                     input_schema={
                         "type": "object",
                         "properties": {
@@ -69,18 +79,10 @@ try:
                                     else 3
                                 )
 
-                                search_results = await self._client.search(
+                                context = await self._search_and_expand(
                                     query, limit=limit
                                 )
 
-                                context_chunks = []
-                                for chunk, score in search_results:
-                                    context_chunks.append(
-                                        f"Content: {chunk.content}\nScore: {score:.4f}"
-                                    )
-
-                                context = "\n\n".join(context_chunks)
-
                                 tool_results.append(
                                     {
                                         "type": "tool_result",

haiku_rag-0.5.4/src/haiku/rag/qa/base.py

@@ -0,0 +1,89 @@
+import json
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS
+
+
+class QuestionAnswerAgentBase:
+    _model: str = ""
+    _system_prompt: str = SYSTEM_PROMPT
+
+    def __init__(self, client: HaikuRAG, model: str = "", use_citations: bool = False):
+        self._model = model
+        self._client = client
+        self._system_prompt = (
+            SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT
+        )
+
+    async def answer(self, question: str) -> str:
+        raise NotImplementedError(
+            "QABase is an abstract class. Please implement the answer method in a subclass."
+        )
+
+    async def _search_and_expand(self, query: str, limit: int = 3) -> str:
+        """Search for documents and expand context, then format as JSON"""
+        search_results = await self._client.search(query, limit=limit)
+        expanded_results = await self._client.expand_context(search_results)
+        return self._format_search_results(expanded_results)
+
+    def _format_search_results(self, search_results) -> str:
+        """Format search results as JSON list of {content, score, document_uri}"""
+        formatted_results = []
+        for chunk, score in search_results:
+            formatted_results.append(
+                {
+                    "content": chunk.content,
+                    "score": score,
+                    "document_uri": chunk.document_uri,
+                }
+            )
+        return json.dumps(formatted_results, indent=2)
+
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "search_documents",
+                "description": "Search the knowledge base for relevant documents. Returns a JSON array of search results.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "The search query to find relevant documents",
+                        },
+                        "limit": {
+                            "type": "integer",
+                            "description": "Maximum number of results to return",
+                            "default": 3,
+                        },
+                    },
+                    "required": ["query"],
+                },
+                "returns": {
+                    "type": "string",
+                    "description": "JSON array of search results",
+                    "schema": {
+                        "type": "array",
+                        "items": {
+                            "type": "object",
+                            "properties": {
+                                "content": {
+                                    "type": "string",
+                                    "description": "The document text content",
+                                },
+                                "score": {
+                                    "type": "number",
+                                    "description": "Relevance score (higher is more relevant)",
+                                },
+                                "document_uri": {
+                                    "type": "string",
+                                    "description": "Source URI/path of the document",
+                                },
+                            },
+                        },
+                    },
+                },
+            },
+        }
+    ]