haiku.rag 0.4.1__tar.gz → 0.4.3__tar.gz

{haiku_rag-0.4.1 → haiku_rag-0.4.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.4.1
+Version: 0.4.3
 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.1 → haiku_rag-0.4.3}/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.1 → haiku_rag-0.4.3}/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.1 → haiku_rag-0.4.3}/docs/python.md

@@ -27,6 +27,31 @@ doc = await client.create_document(
 )
 ```
 
+With custom externally generated chunks:
+```python
+from haiku.rag.store.models.chunk import Chunk
+
+# Create custom chunks with optional embeddings
+chunks = [
+    Chunk(
+        content="This is the first chunk",
+        metadata={"section": "intro"}
+    ),
+    Chunk(
+        content="This is the second chunk",
+        metadata={"section": "body"},
+        embedding=[0.1] * 1024  # Optional pre-computed embedding
+    ),
+]
+
+doc = await client.create_document(
+    content="Full document content",
+    uri="doc://custom",
+    metadata={"source": "manual"},
+    chunks=chunks  # Use provided chunks instead of auto-generating
+)
+```
+
 From file:
 ```python
 doc = await client.create_document_from_source("path/to/document.pdf")

{haiku_rag-0.4.1 → haiku_rag-0.4.3}/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.1 → haiku_rag-0.4.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "haiku.rag"
-version = "0.4.1"
+version = "0.4.3"
 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.4.1 → haiku_rag-0.4.3}/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.1 → haiku_rag-0.4.3}/src/haiku/rag/client.py

@@ -50,7 +50,11 @@ class HaikuRAG:
         return False
 
     async def create_document(
-        self, content: str, uri: str | None = None, metadata: dict | None = None
+        self,
+        content: str,
+        uri: str | None = None,
+        metadata: dict | None = None,
+        chunks: list[Chunk] | None = None,
     ) -> Document:
         """Create a new document with optional URI and metadata.
 
@@ -58,6 +62,7 @@ class HaikuRAG:
             content: The text content of the document.
             uri: Optional URI identifier for the document.
             metadata: Optional metadata dictionary.
+            chunks: Optional list of pre-created chunks to use instead of generating new ones.
 
         Returns:
             The created Document instance.
@@ -67,7 +72,7 @@ class HaikuRAG:
             uri=uri,
             metadata=metadata or {},
         )
-        return await self.document_repository.create(document)
+        return await self.document_repository.create(document, chunks)
 
     async def create_document_from_source(
         self, source: str | Path, metadata: dict = {}

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

@@ -0,0 +1,109 @@
+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.1 → haiku_rag-0.4.3}/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.1 → haiku_rag-0.4.3}/src/haiku/rag/store/models/chunk.py

@@ -7,8 +7,9 @@ class Chunk(BaseModel):
     """
 
     id: int | None = None
-    document_id: int
+    document_id: int | None = None
     content: str
     metadata: dict = {}
     document_uri: str | None = None
     document_meta: dict = {}
+    embedding: list[float] | None = None

{haiku_rag-0.4.1 → haiku_rag-0.4.3}/src/haiku/rag/store/repositories/chunk.py

@@ -18,6 +18,8 @@ class ChunkRepository(BaseRepository[Chunk]):
         """Create a chunk in the database."""
         if self.store._connection is None:
             raise ValueError("Store connection is not available")
+        if entity.document_id is None:
+            raise ValueError("Chunk must have a document_id to be created")
 
         cursor = self.store._connection.cursor()
         cursor.execute(
@@ -34,9 +36,15 @@ class ChunkRepository(BaseRepository[Chunk]):
 
         entity.id = cursor.lastrowid
 
-        # Generate and store embedding
-        embedding = await self.embedder.embed(entity.content)
-        serialized_embedding = self.store.serialize_embedding(embedding)
+        # Generate and store embedding - use existing one if provided
+        if entity.embedding is not None:
+            # Use the provided embedding
+            serialized_embedding = self.store.serialize_embedding(entity.embedding)
+        else:
+            # Generate embedding from content
+            embedding = await self.embedder.embed(entity.content)
+            serialized_embedding = self.store.serialize_embedding(embedding)
+
         cursor.execute(
             """
             INSERT INTO chunk_embeddings (chunk_id, embedding)

{haiku_rag-0.4.1 → haiku_rag-0.4.3}/src/haiku/rag/store/repositories/document.py

@@ -1,8 +1,12 @@
 import json
+from typing import TYPE_CHECKING
 
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.base import BaseRepository
 
+if TYPE_CHECKING:
+    from haiku.rag.store.models.chunk import Chunk
+
 
 class DocumentRepository(BaseRepository[Document]):
     """Repository for Document database operations."""
@@ -16,7 +20,9 @@ class DocumentRepository(BaseRepository[Document]):
             chunk_repository = ChunkRepository(store)
         self.chunk_repository = chunk_repository
 
-    async def create(self, entity: Document) -> Document:
+    async def create(
+        self, entity: Document, chunks: list["Chunk"] | None = None
+    ) -> Document:
         """Create a document with its chunks and embeddings."""
         if self.store._connection is None:
             raise ValueError("Store connection is not available")
@@ -46,10 +52,20 @@ class DocumentRepository(BaseRepository[Document]):
             assert document_id is not None, "Failed to create document in database"
             entity.id = document_id
 
-            # Create chunks and embeddings using ChunkRepository
-            await self.chunk_repository.create_chunks_for_document(
-                document_id, entity.content, commit=False
-            )
+            # Create chunks - either use provided chunks or generate from content
+            if chunks is not None:
+                # Use provided chunks, but update their document_id and set order from list position
+                for order, chunk in enumerate(chunks):
+                    chunk.document_id = document_id
+                    # Ensure order is set from list position
+                    chunk.metadata = chunk.metadata.copy() if chunk.metadata else {}
+                    chunk.metadata["order"] = order
+                    await self.chunk_repository.create(chunk, commit=False)
+            else:
+                # Create chunks and embeddings using ChunkRepository
+                await self.chunk_repository.create_chunks_for_document(
+                    document_id, entity.content, commit=False
+                )
 
             cursor.execute("COMMIT")
             return entity

{haiku_rag-0.4.1 → haiku_rag-0.4.3}/tests/test_client.py

@@ -7,6 +7,7 @@ import pytest
 from datasets import Dataset
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.store.models.chunk import Chunk
 
 
 @pytest.mark.asyncio
@@ -449,3 +450,42 @@ async def test_client_async_context_manager():
     # Context manager should have automatically closed the connection
     # We can't easily test that the connection is closed without accessing internals,
     # but the test passing means the context manager methods work correctly
+
+
+@pytest.mark.asyncio
+async def test_client_create_document_with_custom_chunks():
+    """Test creating a document with pre-created chunks."""
+    async with HaikuRAG(":memory:") as client:
+        # Create some custom chunks with and without embeddings
+        chunks = [
+            Chunk(content="This is the first chunk", metadata={"custom": "metadata1"}),
+            Chunk(
+                content="This is the second chunk",
+                metadata={"custom": "metadata2"},
+                embedding=[0.1] * 1024,
+            ),  # With embedding
+            Chunk(content="This is the third chunk", metadata={"custom": "metadata3"}),
+        ]
+
+        # Create document with custom chunks
+        document = await client.create_document(
+            content="Full document content", chunks=chunks
+        )
+
+        assert document.id is not None
+        assert document.content == "Full document content"
+
+        # Verify the chunks were created correctly
+        doc_chunks = await client.chunk_repository.get_by_document_id(document.id)
+        assert len(doc_chunks) == 3
+
+        # Check chunks have correct content, document_id, and order from list position
+        for i, chunk in enumerate(doc_chunks):
+            assert chunk.document_id == document.id
+            assert chunk.content == chunks[i].content
+            assert (
+                chunk.metadata["order"] == i
+            )  # Order should be set from list position
+            assert (
+                chunk.metadata["custom"] == f"metadata{i + 1}"
+            )  # Original metadata preserved