From 53445b872200d07de34e2bf2396816b159ed8e07 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 7 Apr 2026 16:22:45 +0300 Subject: [PATCH] Add docling_pages column to DocumentRecord/Document for separate page image storage. --- haiku_rag_slim/haiku/rag/client.py | 62 +++++-------- haiku_rag_slim/haiku/rag/store/engine.py | 6 +- .../haiku/rag/store/models/document.py | 71 +++++++++------ .../haiku/rag/store/repositories/document.py | 24 +++++ tests/test_client.py | 20 +++-- tests/test_document.py | 88 +++++++++++++++++++ 6 files changed, 195 insertions(+), 76 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 4217b0f5..01e88059 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/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( - docling_document.model_dump_json() - ) - existing_doc.docling_version = docling_document.version + existing_doc.set_docling(docling_document) 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( - docling_document.model_dump_json() - ) - existing_doc.docling_version = docling_document.version + existing_doc.set_docling(docling_document) 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( - docling_document.model_dump_json() - ) - existing_doc.docling_version = docling_document.version + existing_doc.set_docling(docling_document) 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( - docling_document.model_dump_json() - ) - existing_doc.docling_version = docling_document.version + existing_doc.set_docling(docling_document) 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( - converted_docling.model_dump_json() - ) - existing_doc.docling_version = converted_docling.version + existing_doc.set_docling(converted_docling) 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(include_pages=True) + docling_doc = doc.get_docling_document() 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.docling_version = docling_document.version + doc.set_docling(docling_document) # 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.docling_version = docling_document.version + doc.set_docling(docling_document) # Prepare chunks with document_id and order for order, chunk in enumerate(embedded_chunks): diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index aeb6748a..f97e35de 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/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": - fields.append(pa.field("docling_document", pa.large_binary())) + if field.name in large_binary_columns: + fields.append(pa.field(field.name, pa.large_binary())) else: fields.append(field) return pa.schema(fields) diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 05d819f4..024205cd 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/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 - operations that need page images (e.g. visualize_chunk). + do not contain page data. """ 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( - self, *, include_pages: bool = False - ) -> "DoclingDocument | None": - """Parse and return the stored DoclingDocument. + def set_docling(self, docling_doc: "DoclingDocument") -> None: + """Serialize and store a DoclingDocument, splitting structure and pages. + + 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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index c829bd57..179d6bca 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/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, }, diff --git a/tests/test_client.py b/tests/test_client.py index c53c00ac..40a1ee73 100644 --- a/tests/test_client.py +++ b/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 ( - decompress_json(updated_doc.docling_document) - == docling_doc.model_dump_json() - ) + structure = json.loads(decompress_json(updated_doc.docling_document)) + assert "pages" not in structure # Custom chunks should be used (not rechunked from docling) chunks = await client.chunk_repository.get_by_document_id(doc.id) diff --git a/tests/test_document.py b/tests/test_document.py index 9e4c92bc..1ce394b8 100644 --- a/tests/test_document.py +++ b/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