Merge pull request #31 from ggozad/fix/docling-chunking

Refactor reading & chunking around docling documents.

src/haiku/rag/chunker.py

@@ -1,11 +1,9 @@
-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 docling_core.types.doc.document import DoclingDocument
 
 from haiku.rag.config import Config
 
@@ -33,27 +31,20 @@ class Chunker:
 
         self.chunker = HybridChunker(tokenizer=tokenizer)  # type: ignore
 
-    async def chunk(self, text: str) -> list[str]:
+    async def chunk(self, document: DoclingDocument) -> list[str]:
-        """Split the text into chunks using docling's structure-aware chunking.
+        """Split the document into chunks using docling's structure-aware chunking.
 
         Args:
-            text: The text to be split into chunks.
+            document: The DoclingDocument to be split into chunks.
 
         Returns:
             A list of text chunks with semantic boundaries.
         """
-        if not text:
+        if document is None:
             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))
+        chunks = list(self.chunker.chunk(document))
         return [self.chunker.contextualize(chunk) for chunk in chunks]
 
 

src/haiku/rag/cli.py

@@ -1,4 +1,5 @@
 import asyncio
+import warnings
 from importlib.metadata import version
 from pathlib import Path
 
@@ -9,6 +10,9 @@ from haiku.rag.app import HaikuRAGApp
 from haiku.rag.config import Config
 from haiku.rag.utils import is_up_to_date
 
+if not Config.ENV == "development":
+    warnings.filterwarnings("ignore")
+
 cli = typer.Typer(
     context_settings={"help_option_names": ["-h", "--help"]}, no_args_is_help=True
 )

src/haiku/rag/client.py

@@ -16,6 +16,7 @@ from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.utils import text_to_docling_document
 
 
 class HaikuRAG:
@@ -49,6 +50,24 @@ class HaikuRAG:
         self.close()
         return False
 
+    async def _create_document_with_docling(
+        self,
+        docling_document,
+        uri: str | None = None,
+        metadata: dict | None = None,
+        chunks: list[Chunk] | None = None,
+    ) -> Document:
+        """Create a new document from DoclingDocument."""
+        content = docling_document.export_to_markdown()
+        document = Document(
+            content=content,
+            uri=uri,
+            metadata=metadata or {},
+        )
+        return await self.document_repository._create_with_docling(
+            document, docling_document, chunks
+        )
+
     async def create_document(
         self,
         content: str,
@@ -67,12 +86,17 @@ class HaikuRAG:
         Returns:
             The created Document instance.
         """
+        # Convert content to DoclingDocument for processing
+        docling_document = text_to_docling_document(content)
+
         document = Document(
             content=content,
             uri=uri,
             metadata=metadata or {},
         )
-        return await self.document_repository.create(document, chunks)
+        return await self.document_repository._create_with_docling(
+            document, docling_document, chunks
+        )
 
     async def create_document_from_source(
         self, source: str | Path, metadata: dict = {}
@@ -101,9 +125,12 @@ class HaikuRAG:
         parsed_url = urlparse(source_str)
         if parsed_url.scheme in ("http", "https"):
             return await self._create_or_update_document_from_url(source_str, metadata)
-
+        elif parsed_url.scheme == "file":
-        # Handle as file path
+            # Handle file:// URI by converting to path
-        source_path = Path(source) if isinstance(source, str) else source
+            source_path = Path(parsed_url.path)
+        else:
+            # Handle as regular file path
+            source_path = Path(source) if isinstance(source, str) else source
         if source_path.suffix.lower() not in FileReader.extensions:
             raise ValueError(f"Unsupported file extension: {source_path.suffix}")
 
@@ -119,7 +146,7 @@ class HaikuRAG:
             # MD5 unchanged, return existing document
             return existing_doc
 
-        content = FileReader.parse_file(source_path)
+        docling_document = FileReader.parse_file(source_path)
 
         # Get content type from file extension
         content_type, _ = mimetypes.guess_type(str(source_path))
@@ -131,13 +158,15 @@ class HaikuRAG:
 
         if existing_doc:
             # Update existing document
-            existing_doc.content = content
+            existing_doc.content = docling_document.export_to_markdown()
             existing_doc.metadata = metadata
-            return await self.update_document(existing_doc)
+            return await self.document_repository._update_with_docling(
+                existing_doc, docling_document
+            )
         else:
-            # Create new document
+            # Create new document using DoclingDocument
-            return await self.create_document(
+            return await self._create_document_with_docling(
-                content=content, uri=uri, metadata=metadata
+                docling_document=docling_document, uri=uri, metadata=metadata
             )
 
     async def _create_or_update_document_from_url(
@@ -193,18 +222,20 @@ class HaikuRAG:
                 temp_path = Path(temp_file.name)
 
                 # Parse the content using FileReader
-                content = FileReader.parse_file(temp_path)
+                docling_document = FileReader.parse_file(temp_path)
 
             # Merge metadata with contentType and md5
             metadata.update({"contentType": content_type, "md5": md5_hash})
 
             if existing_doc:
-                existing_doc.content = content
+                existing_doc.content = docling_document.export_to_markdown()
                 existing_doc.metadata = metadata
-                return await self.update_document(existing_doc)
+                return await self.document_repository._update_with_docling(
+                    existing_doc, docling_document
+                )
             else:
-                return await self.create_document(
+                return await self._create_document_with_docling(
-                    content=content, uri=url, metadata=metadata
+                    docling_document=docling_document, uri=url, metadata=metadata
                 )
 
     def _get_extension_from_content_type_or_url(
@@ -262,7 +293,12 @@ class HaikuRAG:
 
     async def update_document(self, document: Document) -> Document:
         """Update an existing document."""
-        return await self.document_repository.update(document)
+        # Convert content to DoclingDocument
+        docling_document = text_to_docling_document(document.content)
+
+        return await self.document_repository._update_with_docling(
+            document, docling_document
+        )
 
     async def delete_document(self, document_id: int) -> bool:
         """Delete a document by its ID."""
@@ -328,6 +364,13 @@ class HaikuRAG:
     async def rebuild_database(self) -> AsyncGenerator[int, None]:
         """Rebuild the database by deleting all chunks and re-indexing all documents.
 
+        For documents with URIs:
+        - Deletes the document and re-adds it from source if source exists
+        - Skips documents where source no longer exists
+
+        For documents without URIs:
+        - Re-creates chunks from existing content
+
         Yields:
             int: The ID of the document currently being processed
         """
@@ -343,9 +386,36 @@ class HaikuRAG:
         documents = await self.list_documents()
 
         for doc in documents:
-            if doc.id is not None:
+            assert doc.id is not None, "Document ID should not be None"
+            if doc.uri:
+                # Document has a URI - delete and try to re-add from source
+                try:
+                    # Delete the old document first
+                    await self.delete_document(doc.id)
+
+                    # Try to re-create from source (this creates the document with chunks)
+                    new_doc = await self.create_document_from_source(
+                        doc.uri, doc.metadata or {}
+                    )
+
+                    assert new_doc.id is not None, "New document ID should not be None"
+                    yield new_doc.id
+
+                except (FileNotFoundError, ValueError, OSError) as e:
+                    # Source doesn't exist or can't be accessed - document already deleted, skip
+                    print(f"Skipping document with URI {doc.uri}: {e}")
+                    continue
+                except Exception as e:
+                    # Unexpected error - log it and skip
+                    print(
+                        f"Unexpected error processing document with URI {doc.uri}: {e}"
+                    )
+                    continue
+            else:
+                # Document without URI - re-create chunks from existing content
+                docling_document = text_to_docling_document(doc.content)
                 await self.chunk_repository.create_chunks_for_document(
-                    doc.id, doc.content, commit=False
+                    doc.id, docling_document, commit=False
                 )
                 yield doc.id
 

src/haiku/rag/config.py

@@ -10,7 +10,7 @@ load_dotenv()
 
 
 class AppConfig(BaseModel):
-    ENV: str = "development"
+    ENV: str = "production"
 
     DEFAULT_DATA_DIR: Path = get_default_data_dir()
     MONITOR_DIRECTORIES: list[Path] = []

src/haiku/rag/reader.py

@@ -2,6 +2,9 @@ from pathlib import Path
 from typing import ClassVar
 
 from docling.document_converter import DocumentConverter
+from docling_core.types.doc.document import DoclingDocument
+
+from haiku.rag.utils import text_to_docling_document
 
 
 class FileReader:
@@ -84,7 +87,7 @@ class FileReader:
     extensions: ClassVar[list[str]] = docling_extensions + text_extensions
 
     @staticmethod
-    def parse_file(path: Path) -> str:
+    def parse_file(path: Path) -> DoclingDocument:
         try:
             file_extension = path.suffix.lower()
 
@@ -92,7 +95,7 @@ class FileReader:
                 # Use docling for complex document formats
                 converter = DocumentConverter()
                 result = converter.convert(path)
-                return result.document.export_to_markdown()
+                return result.document
             elif file_extension in FileReader.text_extensions:
                 # Read plain text files directly
                 content = path.read_text(encoding="utf-8")
@@ -100,11 +103,13 @@ class FileReader:
                 # 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```"
+                    content = f"```{language}\n{content}\n```"
 
-                return content
+                # Convert text to DoclingDocument by wrapping as markdown
+                return text_to_docling_document(content, name=f"{path.stem}.md")
             else:
-                # Fallback: try to read as text
+                # Fallback: try to read as text and convert to DoclingDocument
-                return path.read_text(encoding="utf-8")
+                content = path.read_text(encoding="utf-8")
+                return text_to_docling_document(content, name=f"{path.stem}.md")
         except Exception:
             raise ValueError(f"Failed to parse file: {path}")

src/haiku/rag/store/repositories/chunk.py

@@ -1,6 +1,8 @@
 import json
 import re
 
+from docling_core.types.doc.document import DoclingDocument
+
 from haiku.rag.chunker import chunker
 from haiku.rag.embeddings import get_embedder
 from haiku.rag.store.models.chunk import Chunk
@@ -197,11 +199,11 @@ class ChunkRepository(BaseRepository[Chunk]):
         ]
 
     async def create_chunks_for_document(
-        self, document_id: int, content: str, commit: bool = True
+        self, document_id: int, document: DoclingDocument, commit: bool = True
     ) -> list[Chunk]:
-        """Create chunks and embeddings for a document."""
+        """Create chunks and embeddings for a document from DoclingDocument."""
         # Chunk the document content
-        chunk_texts = await chunker.chunk(content)
+        chunk_texts = await chunker.chunk(document)
         created_chunks = []
 
         # Create chunks with embeddings using the create method

src/haiku/rag/store/repositories/document.py

@@ -1,8 +1,11 @@
 import json
 from typing import TYPE_CHECKING
 
+from docling_core.types.doc.document import DoclingDocument
+
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.base import BaseRepository
+from haiku.rag.utils import text_to_docling_document
 
 if TYPE_CHECKING:
     from haiku.rag.store.models.chunk import Chunk
@@ -20,8 +23,11 @@ class DocumentRepository(BaseRepository[Document]):
             chunk_repository = ChunkRepository(store)
         self.chunk_repository = chunk_repository
 
-    async def create(
+    async def _create_with_docling(
-        self, entity: Document, chunks: list["Chunk"] | None = None
+        self,
+        entity: Document,
+        docling_document: DoclingDocument,
+        chunks: list["Chunk"] | None = None,
     ) -> Document:
         """Create a document with its chunks and embeddings."""
         if self.store._connection is None:
@@ -62,9 +68,9 @@ class DocumentRepository(BaseRepository[Document]):
                     chunk.metadata["order"] = order
                     await self.chunk_repository.create(chunk, commit=False)
             else:
-                # Create chunks and embeddings using ChunkRepository
+                # Create chunks and embeddings using DoclingDocument
                 await self.chunk_repository.create_chunks_for_document(
-                    document_id, entity.content, commit=False
+                    document_id, docling_document, commit=False
                 )
 
             cursor.execute("COMMIT")
@@ -74,6 +80,13 @@ class DocumentRepository(BaseRepository[Document]):
             cursor.execute("ROLLBACK")
             raise
 
+    async def create(self, entity: Document) -> Document:
+        """Create a document with its chunks and embeddings."""
+        # Convert content to DoclingDocument
+        docling_document = text_to_docling_document(entity.content)
+
+        return await self._create_with_docling(entity, docling_document)
+
     async def get_by_id(self, entity_id: int) -> Document | None:
         """Get a document by its ID."""
         if self.store._connection is None:
@@ -134,7 +147,9 @@ class DocumentRepository(BaseRepository[Document]):
             updated_at=updated_at,
         )
 
-    async def update(self, entity: Document) -> Document:
+    async def _update_with_docling(
+        self, entity: Document, docling_document: DoclingDocument
+    ) -> Document:
         """Update an existing document and regenerate its chunks and embeddings."""
         if self.store._connection is None:
             raise ValueError("Store connection is not available")
@@ -163,10 +178,10 @@ class DocumentRepository(BaseRepository[Document]):
                 },
             )
 
-            # Delete existing chunks and regenerate using ChunkRepository
+            # Delete existing chunks and regenerate using DoclingDocument
             await self.chunk_repository.delete_by_document_id(entity.id, commit=False)
             await self.chunk_repository.create_chunks_for_document(
-                entity.id, entity.content, commit=False
+                entity.id, docling_document, commit=False
             )
 
             cursor.execute("COMMIT")
@@ -176,6 +191,13 @@ class DocumentRepository(BaseRepository[Document]):
             cursor.execute("ROLLBACK")
             raise
 
+    async def update(self, entity: Document) -> Document:
+        """Update an existing document and regenerate its chunks and embeddings."""
+        # Convert content to DoclingDocument
+        docling_document = text_to_docling_document(entity.content)
+
+        return await self._update_with_docling(entity, docling_document)
+
     async def delete(self, entity_id: int) -> bool:
         """Delete a document and all its associated chunks and embeddings."""
         # Delete chunks and embeddings first

src/haiku/rag/utils.py

@@ -1,8 +1,12 @@
 import sys
 from importlib import metadata
+from io import BytesIO
 from pathlib import Path
 
 import httpx
+from docling.document_converter import DocumentConverter
+from docling_core.types.doc.document import DoclingDocument
+from docling_core.types.io import DocumentStream
 from packaging.version import Version, parse
 
 
@@ -77,3 +81,20 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
             # If no network connection, do not raise alarms.
             pypi_version = running_version
     return running_version >= pypi_version, running_version, pypi_version
+
+
+def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocument:
+    """Convert text content to a DoclingDocument.
+
+    Args:
+        text: The text content to convert.
+        name: The name to use for the document stream (defaults to "content.md").
+
+    Returns:
+        A DoclingDocument created from the text content.
+    """
+    bytes_io = BytesIO(text.encode("utf-8"))
+    doc_stream = DocumentStream(name=name, stream=bytes_io)
+    converter = DocumentConverter()
+    result = converter.convert(doc_stream)
+    return result.document

tests/test_chunk.py

@@ -6,6 +6,7 @@ from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 from haiku.rag.store.repositories.chunk import ChunkRepository
 from haiku.rag.store.repositories.document import DocumentRepository
+from haiku.rag.utils import text_to_docling_document
 
 
 @pytest.mark.asyncio
@@ -77,8 +78,11 @@ async def test_create_chunks_for_document(qa_corpus: Dataset):
 
     assert document_id is not None, "Document ID should not be None"
 
+    # Convert text to DoclingDocument
+    docling_document = text_to_docling_document(document_text, name="test.md")
+
     # Test creating chunks for the document
-    chunks = await chunk_repo.create_chunks_for_document(document_id, document_text)
+    chunks = await chunk_repo.create_chunks_for_document(document_id, docling_document)
 
     # Verify chunks were created
     assert len(chunks) > 0

tests/test_chunker.py

@@ -2,13 +2,18 @@ import pytest
 from datasets import Dataset
 
 from haiku.rag.chunker import Chunker
+from haiku.rag.utils import text_to_docling_document
 
 
 @pytest.mark.asyncio
 async def test_chunker(qa_corpus: Dataset):
     chunker = Chunker()
-    doc = qa_corpus[0]["document_extracted"]
+    doc_text = qa_corpus[0]["document_extracted"]
-    chunks = await Chunker().chunk(doc)
+
+    # Convert text to DoclingDocument
+    doc = text_to_docling_document(doc_text, name="test.md")
+
+    chunks = await chunker.chunk(doc)
 
     # Ensure that the text is split into multiple chunks
     assert len(chunks) > 1
@@ -27,7 +32,7 @@ async def test_chunker(qa_corpus: Dataset):
         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=()))
+    original_tokens = len(Chunker.encoder.encode(doc_text, disallowed_special=()))
 
     # Due to structure-aware chunking, we might have some variation in token count
     # but it should be reasonable

tests/test_reader.py

@@ -15,8 +15,9 @@ def test_code_file_wrapped_in_code_block():
         f.flush()
         temp_path = Path(f.name)
 
-        result = FileReader.parse_file(temp_path)
+        document = FileReader.parse_file(temp_path)
+        result = document.export_to_markdown()
 
-        assert result.startswith("```python\n")
+        assert result.startswith("```\n")
         assert result.endswith("\n```")
         assert "def hello_world():" in result

tests/test_utils.py

@@ -1,4 +1,8 @@
-from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int
+from haiku.rag.utils import (
+    int_to_semantic_version,
+    semantic_version_to_int,
+    text_to_docling_document,
+)
 
 
 def test_sqlite_user_version():
@@ -13,3 +17,117 @@ def test_sqlite_user_version():
     version = "255.255.255"
     assert semantic_version_to_int(version) == 16777215
     assert int_to_semantic_version(16777215) == version
+
+
+def test_text_to_docling_document():
+    """Test the text_to_docling_document utility function."""
+    # Test basic text conversion
+    simple_text = "This is a simple text document."
+    doc = text_to_docling_document(simple_text)
+
+    # Verify it returns a DoclingDocument
+    from docling_core.types.doc.document import DoclingDocument
+
+    assert isinstance(doc, DoclingDocument)
+
+    # Verify the content can be exported back to markdown
+    markdown = doc.export_to_markdown()
+    assert "This is a simple text document." in markdown
+
+
+def test_text_to_docling_document_with_custom_name():
+    """Test text_to_docling_document with custom name parameter."""
+    code_text = """# Python Code
+
+```python
+def hello():
+    print("Hello, World!")
+    return True
+```"""
+
+    doc = text_to_docling_document(code_text, name="hello.md")
+
+    # Verify it's a valid DoclingDocument
+    from docling_core.types.doc.document import DoclingDocument
+
+    assert isinstance(doc, DoclingDocument)
+
+    # Verify the content is preserved
+    markdown = doc.export_to_markdown()
+    assert "def hello():" in markdown
+    assert "Hello, World!" in markdown
+
+
+def test_text_to_docling_document_markdown_content():
+    """Test text_to_docling_document with markdown content."""
+    markdown_text = """# Test Document
+
+This is a test document with:
+
+- List item 1
+- List item 2
+
+## Code Example
+
+```python
+def test():
+    return "Hello"
+```
+
+**Bold text** and *italic text*."""
+
+    doc = text_to_docling_document(markdown_text, name="test.md")
+
+    # Verify it's a DoclingDocument
+    from docling_core.types.doc.document import DoclingDocument
+
+    assert isinstance(doc, DoclingDocument)
+
+    # Verify the markdown structure is preserved
+    result_markdown = doc.export_to_markdown()
+    assert "# Test Document" in result_markdown
+    assert "List item 1" in result_markdown
+    assert "def test():" in result_markdown
+
+
+def test_text_to_docling_document_empty_content():
+    """Test text_to_docling_document with empty content."""
+    doc = text_to_docling_document("")
+
+    # Should still create a valid DoclingDocument
+    from docling_core.types.doc.document import DoclingDocument
+
+    assert isinstance(doc, DoclingDocument)
+
+    # Export should work even with empty content
+    markdown = doc.export_to_markdown()
+    assert isinstance(markdown, str)
+
+
+def test_text_to_docling_document_unicode_content():
+    """Test text_to_docling_document with unicode content."""
+    unicode_text = """# 测试文档
+
+这是一个包含中文的测试文档。
+
+## Código en Español
+```javascript
+function saludar() {
+    return "¡Hola mundo!";
+}
+```
+
+Emoji test: 🚀 ✅ 📝"""
+
+    doc = text_to_docling_document(unicode_text, name="unicode.md")
+
+    # Verify it's a DoclingDocument
+    from docling_core.types.doc.document import DoclingDocument
+
+    assert isinstance(doc, DoclingDocument)
+
+    # Verify unicode content is preserved
+    result_markdown = doc.export_to_markdown()
+    assert "测试文档" in result_markdown
+    assert "¡Hola mundo!" in result_markdown
+    assert "🚀" in result_markdown