diff --git a/CHANGELOG.md b/CHANGELOG.md index 83352f45..faab4702 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,6 +13,7 @@ - **Context expansion performance**: Load only docling columns during expand_context (skip content blob), and only when doc_item_refs exist - **Chunk expansion performance**: Fetch only chunks in the needed order range during context expansion instead of all chunks for a document - **Embedding batching**: Batch embedding calls in groups of 512 to avoid request size limits and timeouts with large documents +- **DoclingDocument validation**: Strip page images before validation on the read path — pages are only needed for visualize_chunk and account for ~99% of the JSON size ## [0.36.3] - 2026-04-01 diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 2521b1a7..4217b0f5 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -1586,8 +1586,8 @@ class HaikuRAG: if not doc: return [] - # Get DoclingDocument - docling_doc = doc.get_docling_document() + # Get DoclingDocument with page images for rendering + docling_doc = doc.get_docling_document(include_pages=True) if not docling_doc: return [] diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 02406238..05d819f4 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/haiku_rag_slim/haiku/rag/store/models/document.py @@ -1,3 +1,4 @@ +import json from datetime import datetime from typing import TYPE_CHECKING @@ -13,21 +14,41 @@ if TYPE_CHECKING: _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100) -def _get_cached_docling_document( - document_id: str, compressed_data: bytes -) -> "DoclingDocument": - """Get or parse DoclingDocument with LRU caching by document ID.""" - if document_id in _docling_document_cache: - return _docling_document_cache[document_id] - +def _validate_without_pages(compressed_data: bytes) -> "DoclingDocument": + """Decompress and validate DoclingDocument, stripping page images.""" from docling_core.types.doc.document import DoclingDocument json_str = decompress_json(compressed_data) - doc = DoclingDocument.model_validate_json(json_str) + data = json.loads(json_str) + data.pop("pages", None) + return DoclingDocument.model_validate(data) + + +def _get_cached_docling_document( + document_id: str, compressed_data: bytes +) -> "DoclingDocument": + """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). + """ + if document_id in _docling_document_cache: + return _docling_document_cache[document_id] + + doc = _validate_without_pages(compressed_data) _docling_document_cache[document_id] = doc 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) @@ -48,22 +69,30 @@ class Document(BaseModel): created_at: datetime = Field(default_factory=datetime.now) updated_at: datetime = Field(default_factory=datetime.now) - def get_docling_document(self) -> "DoclingDocument | None": + def get_docling_document( + self, *, include_pages: bool = False + ) -> "DoclingDocument | None": """Parse and return the stored DoclingDocument. + 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. """ 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: - from docling_core.types.doc.document import DoclingDocument - - json_str = decompress_json(self.docling_document) - return DoclingDocument.model_validate_json(json_str) + return _validate_without_pages(self.docling_document) return _get_cached_docling_document(self.id, self.docling_document) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index b98f033a..1165950f 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -397,35 +397,6 @@ class ChunkRepository: for rec in results ] - async def get_adjacent_chunks(self, chunk: Chunk, num_adjacent: int) -> list[Chunk]: - """Get adjacent chunks before and after the given chunk within the same document.""" - assert chunk.document_id, "Document id is required for adjacent chunk finding" - - min_order = chunk.order - num_adjacent - max_order = chunk.order + num_adjacent - - where = ( - f"document_id = '{chunk.document_id}'" - f" AND `order` >= {min_order}" - f" AND `order` <= {max_order}" - f" AND id != '{chunk.id}'" - ) - results = list( - self.store.chunks_table.search() - .where(where) - .to_pydantic(self.store.ChunkRecord) - ) - return [ - Chunk( - id=rec.id, - document_id=rec.document_id, - content=rec.content, - metadata=json.loads(rec.metadata), - order=rec.order, - ) - for rec in results - ] - async def _process_search_results( self, query_result: "pd.DataFrame | LanceQueryBuilder" ) -> list[tuple[Chunk, float]]: