haiku.rag 0.4.2__tar.gz → 0.5.0__tar.gz

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.4.2
+Version: 0.5.0
 Summary: Retrieval Augmented Generation (RAG) with SQLite
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -18,9 +18,9 @@ Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Typing :: Typed
 Requires-Python: >=3.10
+Requires-Dist: docling>=2.15.0
 Requires-Dist: fastmcp>=2.8.1
 Requires-Dist: httpx>=0.28.1
-Requires-Dist: markitdown[audio-transcription,docx,pdf,pptx,xlsx]>=0.1.2
 Requires-Dist: mxbai-rerank>=0.1.6
 Requires-Dist: ollama>=0.5.1
 Requires-Dist: pydantic>=2.11.7
@@ -55,7 +55,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite.
 - **Reranking**: Default search result reranking with MixedBread AI or Cohere
 - **Question answering**: Built-in QA agents on your documents
 - **File monitoring**: Auto-index files when run as server
-- **40+ file formats**: PDF, DOCX, HTML, Markdown, audio, URLs
+- **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs
 - **MCP server**: Expose as tools for AI assistants
 - **CLI & Python API**: Use from command line or Python
 

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/README.md

@@ -13,7 +13,7 @@ Retrieval-Augmented Generation (RAG) library on SQLite.
 - **Reranking**: Default search result reranking with MixedBread AI or Cohere
 - **Question answering**: Built-in QA agents on your documents
 - **File monitoring**: Auto-index files when run as server
-- **40+ file formats**: PDF, DOCX, HTML, Markdown, audio, URLs
+- **40+ file formats**: PDF, DOCX, HTML, Markdown, code files, URLs
 - **MCP server**: Expose as tools for AI assistants
 - **CLI & Python API**: Use from command line or Python
 

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/docs/configuration.md

@@ -150,7 +150,4 @@ DEFAULT_DATA_DIR="/path/to/data"
 ```bash
 # Chunk size for document processing
 CHUNK_SIZE=256
-
-# Chunk overlap for better context
-CHUNK_OVERLAP=32
 ```

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/docs/index.md

@@ -10,7 +10,7 @@
 - **Reranking**: Optional result reranking with MixedBread AI or Cohere
 - **Question Answering**: Built-in QA agents using Ollama, OpenAI, or Anthropic.
 - **File monitoring**: Automatically index files when run as a server
-- **Extended file format support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, audio and more. Or add a URL!
+- **Extended file format support**: Parse 40+ file formats including PDF, DOCX, HTML, Markdown, code files and more. Or add a URL!
 - **MCP server**: Exposes functionality as MCP tools
 - **CLI commands**: Access all functionality from your terminal
 - **Python client**: Call `haiku.rag` from your own python applications

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/docs/server.md

@@ -35,7 +35,8 @@ The server can parse 40+ file formats including:
 - Microsoft Office (DOCX, XLSX, PPTX)
 - HTML and Markdown
 - Plain text files
-- Audio files
+- Code files (Python, JavaScript, etc.)
+- Images (processed via OCR)
 - And more...
 
 URLs are also supported for web content.

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.4.2"
+version = "0.5.0"
 description = "Retrieval Augmented Generation (RAG) with SQLite"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
@@ -22,9 +22,9 @@ classifiers = [
 ]
 
 dependencies = [
+    "docling>=2.15.0",
     "fastmcp>=2.8.1",
     "httpx>=0.28.1",
-    "markitdown[audio-transcription,docx,pdf,pptx,xlsx]>=0.1.2",
     "mxbai-rerank>=0.1.6",
     "ollama>=0.5.1",
     "pydantic>=2.11.7",

haiku_rag-0.5.0/src/haiku/rag/chunker.py

@@ -0,0 +1,60 @@
+from io import BytesIO
+from typing import ClassVar
+
+import tiktoken
+from docling.chunking import HybridChunker  # type: ignore
+from docling.document_converter import DocumentConverter
+from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
+from docling_core.types.io import DocumentStream
+
+from haiku.rag.config import Config
+
+
+class Chunker:
+    """A class that chunks text into smaller pieces for embedding and retrieval.
+
+    Uses docling's structure-aware chunking to create semantically meaningful chunks
+    that respect document boundaries.
+
+    Args:
+        chunk_size: The maximum size of a chunk in tokens.
+    """
+
+    encoder: ClassVar[tiktoken.Encoding] = tiktoken.encoding_for_model("gpt-4o")
+
+    def __init__(
+        self,
+        chunk_size: int = Config.CHUNK_SIZE,
+    ):
+        self.chunk_size = chunk_size
+        tokenizer = OpenAITokenizer(
+            tokenizer=tiktoken.encoding_for_model("gpt-4o"), max_tokens=chunk_size
+        )
+
+        self.chunker = HybridChunker(tokenizer=tokenizer)  # type: ignore
+
+    async def chunk(self, text: str) -> list[str]:
+        """Split the text into chunks using docling's structure-aware chunking.
+
+        Args:
+            text: The text to be split into chunks.
+
+        Returns:
+            A list of text chunks with semantic boundaries.
+        """
+        if not text:
+            return []
+
+        # Convert to docling document
+        bytes_io = BytesIO(text.encode("utf-8"))
+        doc_stream = DocumentStream(name="text.md", stream=bytes_io)
+        converter = DocumentConverter()
+        result = converter.convert(doc_stream)
+        doc = result.document
+
+        # Chunk using docling's hybrid chunker
+        chunks = list(self.chunker.chunk(doc))
+        return [self.chunker.contextualize(chunk) for chunk in chunks]
+
+
+chunker = Chunker()

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/src/haiku/rag/cli.py

@@ -1,4 +1,5 @@
 import asyncio
+from importlib.metadata import version
 from pathlib import Path
 
 import typer
@@ -26,8 +27,23 @@ async def check_version():
         console.print("[yellow]Please update.[/yellow]")
 
 
+def version_callback(value: bool):
+    if value:
+        v = version("haiku.rag")
+        console.print(f"haiku.rag version {v}")
+        raise typer.Exit()
+
+
 @cli.callback()
-def main():
+def main(
+    _version: bool = typer.Option(
+        False,
+        "-v",
+        "--version",
+        callback=version_callback,
+        help="Show version and exit",
+    ),
+):
     """haiku.rag CLI - SQLite-based RAG system"""
     # Run version check before any command
     event_loop.run_until_complete(check_version())

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/src/haiku/rag/config.py

@@ -27,7 +27,6 @@ class AppConfig(BaseModel):
     QA_MODEL: str = "qwen3"
 
     CHUNK_SIZE: int = 256
-    CHUNK_OVERLAP: int = 32
 
     OLLAMA_BASE_URL: str = "http://localhost:11434"
 

haiku_rag-0.5.0/src/haiku/rag/reader.py

@@ -0,0 +1,110 @@
+from pathlib import Path
+from typing import ClassVar
+
+from docling.document_converter import DocumentConverter
+
+
+class FileReader:
+    # Extensions supported by docling
+    docling_extensions: ClassVar[list[str]] = [
+        ".asciidoc",
+        ".bmp",
+        ".csv",
+        ".docx",
+        ".html",
+        ".xhtml",
+        ".jpeg",
+        ".jpg",
+        ".md",
+        ".pdf",
+        ".png",
+        ".pptx",
+        ".tiff",
+        ".xlsx",
+        ".xml",
+        ".webp",
+    ]
+
+    # Plain text extensions that we'll read directly
+    text_extensions: ClassVar[list[str]] = [
+        ".astro",
+        ".c",
+        ".cpp",
+        ".css",
+        ".go",
+        ".h",
+        ".hpp",
+        ".java",
+        ".js",
+        ".json",
+        ".kt",
+        ".mdx",
+        ".mjs",
+        ".php",
+        ".py",
+        ".rb",
+        ".rs",
+        ".svelte",
+        ".swift",
+        ".ts",
+        ".tsx",
+        ".txt",
+        ".vue",
+        ".yaml",
+        ".yml",
+    ]
+
+    # Code file extensions with their markdown language identifiers for syntax highlighting
+    code_markdown_identifier: ClassVar[dict[str, str]] = {
+        ".astro": "astro",
+        ".c": "c",
+        ".cpp": "cpp",
+        ".css": "css",
+        ".go": "go",
+        ".h": "c",
+        ".hpp": "cpp",
+        ".java": "java",
+        ".js": "javascript",
+        ".json": "json",
+        ".kt": "kotlin",
+        ".mjs": "javascript",
+        ".php": "php",
+        ".py": "python",
+        ".rb": "ruby",
+        ".rs": "rust",
+        ".svelte": "svelte",
+        ".swift": "swift",
+        ".ts": "typescript",
+        ".tsx": "tsx",
+        ".vue": "vue",
+        ".yaml": "yaml",
+        ".yml": "yaml",
+    }
+
+    extensions: ClassVar[list[str]] = docling_extensions + text_extensions
+
+    @staticmethod
+    def parse_file(path: Path) -> str:
+        try:
+            file_extension = path.suffix.lower()
+
+            if file_extension in FileReader.docling_extensions:
+                # Use docling for complex document formats
+                converter = DocumentConverter()
+                result = converter.convert(path)
+                return result.document.export_to_markdown()
+            elif file_extension in FileReader.text_extensions:
+                # Read plain text files directly
+                content = path.read_text(encoding="utf-8")
+
+                # Wrap code files (but not plain txt) in markdown code blocks for better presentation
+                if file_extension in FileReader.code_markdown_identifier:
+                    language = FileReader.code_markdown_identifier[file_extension]
+                    return f"```{language}\n{content}\n```"
+
+                return content
+            else:
+                # Fallback: try to read as text
+                return path.read_text(encoding="utf-8")
+        except Exception:
+            raise ValueError(f"Failed to parse file: {path}")

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/src/haiku/rag/store/engine.py

@@ -37,6 +37,11 @@ class Store:
         db = sqlite3.connect(self.db_path)
         db.enable_load_extension(True)
         sqlite_vec.load(db)
+
+        # Enable WAL mode for better concurrency (skip for in-memory databases)
+        if self.db_path != ":memory:":
+            db.execute("PRAGMA journal_mode=WAL")
+
         self._connection = db
         existing_tables = [
             row[0]

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/src/haiku/rag/store/repositories/settings.py

@@ -63,7 +63,6 @@ class SettingsRepository:
             "EMBEDDINGS_MODEL",
             "EMBEDDINGS_VECTOR_DIM",
             "CHUNK_SIZE",
-            "CHUNK_OVERLAP",
         ]
 
         errors = []

haiku_rag-0.5.0/tests/test_chunker.py

@@ -0,0 +1,34 @@
+import pytest
+from datasets import Dataset
+
+from haiku.rag.chunker import Chunker
+
+
+@pytest.mark.asyncio
+async def test_chunker(qa_corpus: Dataset):
+    chunker = Chunker()
+    doc = qa_corpus[0]["document_extracted"]
+    chunks = await Chunker().chunk(doc)
+
+    # Ensure that the text is split into multiple chunks
+    assert len(chunks) > 1
+
+    # Ensure that chunks are reasonably sized (allowing more flexibility for structure-aware chunking)
+    total_tokens = 0
+    for chunk in chunks:
+        encoded_tokens = Chunker.encoder.encode(chunk, disallowed_special=())
+        token_count = len(encoded_tokens)
+        total_tokens += token_count
+
+        # Each chunk should be reasonably sized (allowing more flexibility than the old strict limits)
+        assert (
+            token_count <= chunker.chunk_size * 1.2
+        )  # Allow some flexibility for semantic boundaries
+        assert token_count > 5  # Ensure chunks aren't too small
+
+    # Ensure that all chunks together contain roughly the same content as original
+    original_tokens = len(Chunker.encoder.encode(doc, disallowed_special=()))
+
+    # Due to structure-aware chunking, we might have some variation in token count
+    # but it should be reasonable
+    assert abs(total_tokens - original_tokens) <= original_tokens * 0.1

haiku_rag-0.5.0/tests/test_reader.py

@@ -0,0 +1,22 @@
+import tempfile
+from pathlib import Path
+
+from haiku.rag.reader import FileReader
+
+
+def test_code_file_wrapped_in_code_block():
+    """Test that code files are wrapped in markdown code blocks."""
+    python_code = '''def hello_world():
+    print("Hello, World!")
+    return "success"'''
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py") as f:
+        f.write(python_code)
+        f.flush()
+        temp_path = Path(f.name)
+
+        result = FileReader.parse_file(temp_path)
+
+        assert result.startswith("```python\n")
+        assert result.endswith("\n```")
+        assert "def hello_world():" in result

{haiku_rag-0.4.2 → haiku_rag-0.5.0}/tests/test_search.py

@@ -36,7 +36,7 @@ async def test_search_qa_corpus(qa_corpus: Dataset):
         created_document = await doc_repo.create(document)
         documents.append((created_document, doc_data))
 
-    for i in range(num_documents):  # Test with first few documents
+    for i in range(5):  # Test with first few documents
         target_document, doc_data = documents[i]
         question = doc_data["question"]
 
@@ -50,6 +50,10 @@ async def test_search_qa_corpus(qa_corpus: Dataset):
         target_document_ids = {chunk.document_id for chunk, _ in fts_results}
         assert target_document.id in target_document_ids
 
+    for i in range(num_documents):  # Test with first few documents
+        target_document, doc_data = documents[i]
+        question = doc_data["question"]
+
         # Test hybrid search
         hybrid_results = await chunk_repo.search_chunks_hybrid(question, limit=5)
         target_document_ids = {chunk.document_id for chunk, _ in hybrid_results}