Operate on DoclingDocument instead of Document

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/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 = {}
@@ -119,8 +143,7 @@ class HaikuRAG:
             # MD5 unchanged, return existing document
             return existing_doc
 
-        document = FileReader.parse_file(source_path)
+        docling_document = FileReader.parse_file(source_path)
-        content = document.export_to_markdown()
 
         # Get content type from file extension
         content_type, _ = mimetypes.guess_type(str(source_path))
@@ -132,13 +155,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(
@@ -194,19 +219,20 @@ class HaikuRAG:
                 temp_path = Path(temp_file.name)
 
                 # Parse the content using FileReader
-                document = FileReader.parse_file(temp_path)
+                docling_document = FileReader.parse_file(temp_path)
-                content = document.export_to_markdown()
 
             # 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(
@@ -264,7 +290,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."""
@@ -346,8 +377,11 @@ class HaikuRAG:
 
         for doc in documents:
             if doc.id is not None:
+                # Convert content to DoclingDocument for rebuild
+                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/reader.py

@@ -1,10 +1,10 @@
-from io import BytesIO
 from pathlib import Path
 from typing import ClassVar
 
 from docling.document_converter import DocumentConverter
 from docling_core.types.doc.document import DoclingDocument
-from docling_core.types.io import DocumentStream
+
+from haiku.rag.utils import text_to_docling_document
 
 
 class FileReader:
@@ -106,18 +106,10 @@ class FileReader:
                     content = f"```{language}\n{content}\n```"
 
                 # Convert text to DoclingDocument by wrapping as markdown
-                bytes_io = BytesIO(content.encode("utf-8"))
+                return text_to_docling_document(content, name=f"{path.stem}.md")
-                doc_stream = DocumentStream(name=f"{path.stem}.md", stream=bytes_io)
-                converter = DocumentConverter()
-                result = converter.convert(doc_stream)
-                return result.document
             else:
                 # Fallback: try to read as text and convert to DoclingDocument
                 content = path.read_text(encoding="utf-8")
-                bytes_io = BytesIO(content.encode("utf-8"))
+                return text_to_docling_document(content, name=f"{path.stem}.md")
-                doc_stream = DocumentStream(name=f"{path.stem}.md", stream=bytes_io)
-                converter = DocumentConverter()
-                result = converter.convert(doc_stream)
-                return result.document
         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

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