Add docling_pages column to DocumentRecord/Document for separate page image storage.

haiku_rag_slim/haiku/rag/client.py

@@ -17,7 +17,6 @@ import httpx
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.converters import get_converter
 from haiku.rag.reranking import get_reranker
-from haiku.rag.store.compression import compress_json
 from haiku.rag.store.engine import Store
 from haiku.rag.store.models.chunk import Chunk, SearchResult
 from haiku.rag.store.models.document import Document
@@ -495,9 +494,8 @@ class HaikuRAG:
             uri=uri,
             title=title,
             metadata=metadata or {},
-            docling_document=compress_json(docling_document.model_dump_json()),
-            docling_version=docling_document.version,
         )
+        document.set_docling(docling_document)
 
         # Store document and chunks
         return await self._store_document_with_chunks(document, embedded_chunks)
@@ -535,9 +533,8 @@ class HaikuRAG:
             uri=uri,
             title=title,
             metadata=metadata or {},
-            docling_document=compress_json(docling_document.model_dump_json()),
-            docling_version=docling_document.version,
         )
+        document.set_docling(docling_document)
 
         return await self._store_document_with_chunks(document, chunks)
 
@@ -672,10 +669,7 @@ class HaikuRAG:
             # Update existing document and rechunk
             existing_doc.content = stored_content
             existing_doc.metadata = metadata
-            existing_doc.docling_document = compress_json(
+            existing_doc.set_docling(docling_document)
-                docling_document.model_dump_json()
-            )
-            existing_doc.docling_version = docling_document.version
             if title is not None:
                 existing_doc.title = title
             elif existing_doc.title is None:
@@ -694,9 +688,8 @@ class HaikuRAG:
                 uri=uri,
                 title=title,
                 metadata=metadata,
-                docling_document=compress_json(docling_document.model_dump_json()),
-                docling_version=docling_document.version,
             )
+            document.set_docling(docling_document)
             return await self._store_document_with_chunks(document, embedded_chunks)
 
     async def _create_or_update_document_from_url(
@@ -789,10 +782,7 @@ class HaikuRAG:
                 # Update existing document and rechunk
                 existing_doc.content = stored_content
                 existing_doc.metadata = metadata
-                existing_doc.docling_document = compress_json(
+                existing_doc.set_docling(docling_document)
-                    docling_document.model_dump_json()
-                )
-                existing_doc.docling_version = docling_document.version
                 if title is not None:
                     existing_doc.title = title
                 elif existing_doc.title is None:
@@ -811,9 +801,8 @@ class HaikuRAG:
                     uri=url,
                     title=title,
                     metadata=metadata,
-                    docling_document=compress_json(docling_document.model_dump_json()),
-                    docling_version=docling_document.version,
                 )
+                document.set_docling(docling_document)
                 return await self._store_document_with_chunks(document, embedded_chunks)
 
     def _get_extension_from_content_type_or_url(
@@ -963,10 +952,7 @@ class HaikuRAG:
             # Store docling data if provided
             if docling_document is not None:
                 existing_doc.content = docling_document.export_to_markdown()
-                existing_doc.docling_document = compress_json(
+                existing_doc.set_docling(docling_document)
-                    docling_document.model_dump_json()
-                )
-                existing_doc.docling_version = docling_document.version
             elif content is not None:
                 existing_doc.content = content
 
@@ -975,10 +961,7 @@ class HaikuRAG:
         # DoclingDocument provided without chunks - chunk and embed using primitives
         if docling_document is not None:
             existing_doc.content = docling_document.export_to_markdown()
-            existing_doc.docling_document = compress_json(
+            existing_doc.set_docling(docling_document)
-                docling_document.model_dump_json()
-            )
-            existing_doc.docling_version = docling_document.version
 
             new_chunks = await self.chunk(docling_document)
             embedded_chunks = await embed_chunks(new_chunks, self._config)
@@ -990,10 +973,7 @@ class HaikuRAG:
         assert content is not None
         existing_doc.content = content
         converted_docling = await self.convert(existing_doc.content)
-        existing_doc.docling_document = compress_json(
+        existing_doc.set_docling(converted_docling)
-            converted_docling.model_dump_json()
-        )
-        existing_doc.docling_version = converted_docling.version
 
         new_chunks = await self.chunk(converted_docling)
         embedded_chunks = await embed_chunks(new_chunks, self._config)
@@ -1578,16 +1558,15 @@ class HaikuRAG:
 
         from PIL import ImageDraw
 
-        # Get the document
+        # Get the document structure (from cache if available)
         if not chunk.document_id:
             return []
 
-        doc = await self.document_repository.get_by_id(chunk.document_id)
+        doc = await self.document_repository.get_docling_data(chunk.document_id)
         if not doc:
             return []
 
-        # Get DoclingDocument with page images for rendering
+        docling_doc = doc.get_docling_document()
-        docling_doc = doc.get_docling_document(include_pages=True)
         if not docling_doc:
             return []
 
@@ -1604,13 +1583,19 @@ class HaikuRAG:
                 boxes_by_page[bbox.page_no] = []
             boxes_by_page[bbox.page_no].append(bbox)
 
+        # Load only the needed page images
+        pages_doc = await self.document_repository.get_pages_data(chunk.document_id)
+        if not pages_doc:
+            return []
+        page_images = pages_doc.get_page_images(list(boxes_by_page.keys()))
+
         # Render each page with its bounding boxes
         images = []
         for page_no in sorted(boxes_by_page.keys()):
-            if page_no not in docling_doc.pages:
+            if page_no not in page_images:
                 continue
 
-            page = docling_doc.pages[page_no]
+            page = page_images[page_no]
             if page.image is None or page.image.pil_image is None:
                 continue
 
@@ -1800,6 +1785,7 @@ class HaikuRAG:
                     title=doc.title,
                     metadata=json.dumps(doc.metadata),
                     docling_document=doc.docling_document,
+                    docling_pages=doc.docling_pages,
                     docling_version=doc.docling_version,
                     created_at=doc.created_at.isoformat() if doc.created_at else now,
                     updated_at=now,
@@ -1836,8 +1822,7 @@ class HaikuRAG:
             embedded_chunks = await embed_chunks(chunks, self._config)
 
             # Update document fields
-            doc.docling_document = compress_json(docling_document.model_dump_json())
+            doc.set_docling(docling_document)
-            doc.docling_version = docling_document.version
 
             # Prepare chunks with document_id and order
             for order, chunk in enumerate(embedded_chunks):
@@ -1915,8 +1900,7 @@ class HaikuRAG:
             chunks = await self.chunk(docling_document)
             embedded_chunks = await embed_chunks(chunks, self._config)
 
-            doc.docling_document = compress_json(docling_document.model_dump_json())
+            doc.set_docling(docling_document)
-            doc.docling_version = docling_document.version
 
             # Prepare chunks with document_id and order
             for order, chunk in enumerate(embedded_chunks):

haiku_rag_slim/haiku/rag/store/engine.py

@@ -26,6 +26,7 @@ class DocumentRecord(LanceModel):
     title: str | None = None
     metadata: str = Field(default="{}")
     docling_document: bytes | None = None
+    docling_pages: bytes | None = None
     docling_version: str | None = None
     created_at: str = Field(default_factory=lambda: "")
     updated_at: str = Field(default_factory=lambda: "")
@@ -43,10 +44,11 @@ def get_documents_arrow_schema() -> pa.Schema:
     which has 64-bit offsets and no practical size limit.
     """
     base_schema = DocumentRecord.to_arrow_schema()
+    large_binary_columns = {"docling_document", "docling_pages"}
     fields = []
     for field in base_schema:
-        if field.name == "docling_document":
+        if field.name in large_binary_columns:
-            fields.append(pa.field("docling_document", pa.large_binary()))
+            fields.append(pa.field(field.name, pa.large_binary()))
         else:
             fields.append(field)
     return pa.schema(fields)

haiku_rag_slim/haiku/rag/store/models/document.py

@@ -5,10 +5,10 @@ from typing import TYPE_CHECKING
 from cachetools import LRUCache
 from pydantic import BaseModel, Field
 
-from haiku.rag.store.compression import decompress_json
+from haiku.rag.store.compression import compress_docling_split, decompress_json
 
 if TYPE_CHECKING:
-    from docling_core.types.doc.document import DoclingDocument
+    from docling_core.types.doc.document import DoclingDocument, PageItem
 
 
 _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
@@ -30,8 +30,7 @@ def _get_cached_docling_document(
     """Get or parse DoclingDocument with LRU caching by document ID.
 
     Strips page images before validation for performance — cached documents
-    do not contain page data. Use _parse_full_docling_document for
+    do not contain page data.
-    operations that need page images (e.g. visualize_chunk).
     """
     if document_id in _docling_document_cache:
         return _docling_document_cache[document_id]
@@ -41,14 +40,6 @@ def _get_cached_docling_document(
     return doc
 
 
-def _parse_full_docling_document(compressed_data: bytes) -> "DoclingDocument":
-    """Parse DoclingDocument with full page data (no caching, no stripping)."""
-    from docling_core.types.doc.document import DoclingDocument
-
-    json_str = decompress_json(compressed_data)
-    return DoclingDocument.model_validate_json(json_str)
-
-
 def invalidate_docling_document_cache(document_id: str) -> None:
     """Remove a document from the DoclingDocument cache."""
     _docling_document_cache.pop(document_id, None)
@@ -65,34 +56,62 @@ class Document(BaseModel):
     title: str | None = None
     metadata: dict = {}
     docling_document: bytes | None = Field(default=None, exclude=True)
+    docling_pages: bytes | None = Field(default=None, exclude=True)
     docling_version: str | None = Field(default=None, exclude=True)
     created_at: datetime = Field(default_factory=datetime.now)
     updated_at: datetime = Field(default_factory=datetime.now)
 
-    def get_docling_document(
+    def set_docling(self, docling_doc: "DoclingDocument") -> None:
-        self, *, include_pages: bool = False
+        """Serialize and store a DoclingDocument, splitting structure and pages.
-    ) -> "DoclingDocument | None":
+
-        """Parse and return the stored DoclingDocument.
+        Sets docling_document (zstd-compressed structure without pages),
+        docling_pages (zstd-compressed page images), and docling_version.
+        """
+        structure, pages = compress_docling_split(docling_doc.model_dump_json())
+        self.docling_document = structure
+        self.docling_pages = pages
+        self.docling_version = docling_doc.version
+
+    def get_docling_document(self) -> "DoclingDocument | None":
+        """Parse and return the stored DoclingDocument (without page images).
 
-        By default, strips page images before parsing for performance.
         Uses LRU cache (keyed by document ID) to avoid repeated parsing.
 
-        Args:
-            include_pages: If True, parse with full page data (slower,
-                bypasses cache). Only needed for operations that access
-                page images (e.g. visualize_chunk).
-
         Returns:
-            The parsed DoclingDocument, or None if not stored or no ID.
+            The parsed DoclingDocument, or None if not stored.
         """
         if self.docling_document is None:
             return None
 
-        if include_pages:
-            return _parse_full_docling_document(self.docling_document)
-
         # No caching for documents without ID
         if self.id is None:
             return _validate_without_pages(self.docling_document)
 
         return _get_cached_docling_document(self.id, self.docling_document)
+
+    def get_page_images(self, page_numbers: list[int]) -> "dict[int, PageItem]":
+        """Decompress and return page images for the requested page numbers.
+
+        Loads only the docling_pages blob — does not need the structure.
+        Validates only the requested pages through Pydantic (for pil_image property).
+
+        Args:
+            page_numbers: Page numbers to retrieve.
+
+        Returns:
+            Dict mapping page number to validated PageItem.
+        """
+        if self.docling_pages is None:
+            return {}
+
+        from docling_core.types.doc.document import PageItem
+
+        pages_json = decompress_json(self.docling_pages)
+        all_pages = json.loads(pages_json)
+
+        result: dict[int, PageItem] = {}
+        for page_no in page_numbers:
+            page_data = all_pages.get(str(page_no))
+            if page_data is not None:
+                result[page_no] = PageItem.model_validate(page_data)
+        return result

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

@@ -36,6 +36,7 @@ class DocumentRepository:
             title=record.title,
             metadata=json.loads(record.metadata),
             docling_document=record.docling_document,
+            docling_pages=record.docling_pages,
             docling_version=record.docling_version,
             created_at=datetime.fromisoformat(record.created_at)
             if record.created_at
@@ -62,6 +63,7 @@ class DocumentRepository:
             title=entity.title,
             metadata=json.dumps(entity.metadata),
             docling_document=entity.docling_document,
+            docling_pages=entity.docling_pages,
             docling_version=entity.docling_version,
             created_at=now,
             updated_at=now,
@@ -114,6 +116,27 @@ class DocumentRepository:
             docling_version=row.get("docling_version"),
         )
 
+    async def get_pages_data(self, entity_id: str) -> Document | None:
+        """Get a document with only page image data loaded."""
+        safe_id = _escape_sql_string(entity_id)
+        results = list(
+            self.store.documents_table.search()
+            .select(["id", "docling_pages"])
+            .where(f"id = '{safe_id}'")
+            .limit(1)
+            .to_list()
+        )
+
+        if not results:
+            return None
+
+        row = results[0]
+        return Document(
+            id=row["id"],
+            content="",
+            docling_pages=row.get("docling_pages"),
+        )
+
     async def update(self, entity: Document) -> Document:
         """Update an existing document."""
         self.store._assert_writable()
@@ -138,6 +161,7 @@ class DocumentRepository:
                 "title": entity.title,
                 "metadata": json.dumps(entity.metadata),
                 "docling_document": entity.docling_document,
+                "docling_pages": entity.docling_pages,
                 "docling_version": entity.docling_version,
                 "updated_at": now,
             },

tests/test_client.py

@@ -1,3 +1,4 @@
+import json
 import tempfile
 from pathlib import Path
 from unittest.mock import AsyncMock, patch
@@ -864,8 +865,11 @@ async def test_client_import_document_stores_docling_data(temp_db_path):
         assert doc.id is not None
         assert "Content from docling document" in doc.content
         assert doc.docling_document is not None
-        assert decompress_json(doc.docling_document) == docling_doc.model_dump_json()
         assert doc.docling_version == docling_doc.version
+        # Structure is stored without pages
+        structure = json.loads(decompress_json(doc.docling_document))
+        assert "pages" not in structure
+        assert structure["name"] == "test"
 
 
 @pytest.mark.vcr()
@@ -983,11 +987,11 @@ async def test_client_update_document_with_docling_rechunks(temp_db_path):
         # Content should be extracted from docling document
         assert "Completely different text" in updated_doc.content
         assert updated_doc.docling_document is not None
-        assert (
-            decompress_json(updated_doc.docling_document)
-            == docling_doc.model_dump_json()
-        )
         assert updated_doc.docling_version == docling_doc.version
+        # Structure is stored without pages
+        structure = json.loads(decompress_json(updated_doc.docling_document))
+        assert "pages" not in structure
+        assert structure["name"] == "updated"
 
         # Chunks should be regenerated
         new_chunks = await client.chunk_repository.get_by_document_id(doc.id)
@@ -1026,10 +1030,8 @@ async def test_client_update_document_docling_with_chunks(temp_db_path):
         # Content should be extracted from docling (since content wasn't provided)
         assert "Text from docling" in updated_doc.content
         assert updated_doc.docling_document is not None
-        assert (
+        structure = json.loads(decompress_json(updated_doc.docling_document))
-            decompress_json(updated_doc.docling_document)
+        assert "pages" not in structure
-            == docling_doc.model_dump_json()
-        )
 
         # Custom chunks should be used (not rechunked from docling)
         chunks = await client.chunk_repository.get_by_document_id(doc.id)

tests/test_document.py

@@ -229,6 +229,94 @@ def test_document_get_docling_document_no_id_no_cache():
     assert doc1 is not doc2
 
 
+def test_set_docling_splits_structure_and_pages():
+    """set_docling stores structure and pages separately."""
+    import json
+
+    from docling_core.types.doc.document import DoclingDocument
+    from docling_core.types.doc.labels import DocItemLabel
+
+    from haiku.rag.store.compression import decompress_json
+
+    docling_doc = DoclingDocument(name="split_test")
+    docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
+
+    document = Document(content="test")
+    document.set_docling(docling_doc)
+
+    assert document.docling_document is not None
+    assert document.docling_version == docling_doc.version
+
+    # Structure should not contain pages
+    structure = json.loads(decompress_json(document.docling_document))
+    assert "pages" not in structure
+    assert structure["name"] == "split_test"
+
+    # get_docling_document should work from the split structure
+    parsed = document.get_docling_document()
+    assert parsed is not None
+    assert parsed.name == "split_test"
+    assert len(list(parsed.iterate_items())) > 0
+
+
+def test_set_docling_with_page_images():
+    """set_docling stores page images in docling_pages."""
+    import json
+
+    from docling_core.types.doc.base import Size
+    from docling_core.types.doc.document import DoclingDocument, PageItem
+    from docling_core.types.doc.labels import DocItemLabel
+
+    from haiku.rag.store.compression import decompress_json
+
+    docling_doc = DoclingDocument(name="pages_test")
+    docling_doc.add_text(label=DocItemLabel.PARAGRAPH, text="Content")
+    docling_doc.pages[1] = PageItem(
+        size=Size(width=612, height=792),
+        page_no=1,
+    )
+
+    document = Document(content="test")
+    document.set_docling(docling_doc)
+
+    assert document.docling_pages is not None
+
+    # Pages blob should contain page data
+    pages = json.loads(decompress_json(document.docling_pages))
+    assert "1" in pages
+
+
+def test_get_page_images():
+    """get_page_images returns requested pages from docling_pages blob."""
+    import json
+
+    from haiku.rag.store.compression import compress_json
+
+    pages_data = {
+        "1": {"size": {"width": 612, "height": 792}, "page_no": 1},
+        "2": {"size": {"width": 612, "height": 792}, "page_no": 2},
+        "3": {"size": {"width": 612, "height": 792}, "page_no": 3},
+    }
+    document = Document(
+        content="test",
+        docling_pages=compress_json(json.dumps(pages_data)),
+    )
+
+    result = document.get_page_images([1, 3])
+    assert len(result) == 2
+    assert 1 in result
+    assert 3 in result
+    assert 2 not in result
+
+    # Missing pages are skipped
+    result = document.get_page_images([99])
+    assert len(result) == 0
+
+    # None docling_pages returns empty
+    doc_no_pages = Document(content="test")
+    assert doc_no_pages.get_page_images([1]) == {}
+
+
 @pytest.mark.asyncio
 async def test_get_docling_data_loads_only_docling_columns(
     qa_corpus: Dataset, temp_db_path