Strip page images from DoclingDocument before validation, unless we use visualize_chunk()

This commit is contained in:
Yiorgis Gozadinos 2026-04-07 12:34:39 +03:00
parent 32258d4e2a
commit 8a24606784
No known key found for this signature in database
4 changed files with 45 additions and 44 deletions

View file

@ -13,6 +13,7 @@
- **Context expansion performance**: Load only docling columns during expand_context (skip content blob), and only when doc_item_refs exist - **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 - **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 - **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 ## [0.36.3] - 2026-04-01

View file

@ -1586,8 +1586,8 @@ class HaikuRAG:
if not doc: if not doc:
return [] return []
# Get DoclingDocument # 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: if not docling_doc:
return [] return []

View file

@ -1,3 +1,4 @@
import json
from datetime import datetime from datetime import datetime
from typing import TYPE_CHECKING from typing import TYPE_CHECKING
@ -13,21 +14,41 @@ if TYPE_CHECKING:
_docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100) _docling_document_cache: LRUCache[str, "DoclingDocument"] = LRUCache(maxsize=100)
def _get_cached_docling_document( def _validate_without_pages(compressed_data: bytes) -> "DoclingDocument":
document_id: str, compressed_data: bytes """Decompress and validate DoclingDocument, stripping page images."""
) -> "DoclingDocument":
"""Get or parse DoclingDocument with LRU caching by document ID."""
if document_id in _docling_document_cache:
return _docling_document_cache[document_id]
from docling_core.types.doc.document import DoclingDocument from docling_core.types.doc.document import DoclingDocument
json_str = decompress_json(compressed_data) 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 _docling_document_cache[document_id] = doc
return 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: def invalidate_docling_document_cache(document_id: str) -> None:
"""Remove a document from the DoclingDocument cache.""" """Remove a document from the DoclingDocument cache."""
_docling_document_cache.pop(document_id, None) _docling_document_cache.pop(document_id, None)
@ -48,22 +69,30 @@ class Document(BaseModel):
created_at: datetime = Field(default_factory=datetime.now) created_at: datetime = Field(default_factory=datetime.now)
updated_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. """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. 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: Returns:
The parsed DoclingDocument, or None if not stored or no ID. The parsed DoclingDocument, or None if not stored or no ID.
""" """
if self.docling_document is None: if self.docling_document is None:
return None return None
if include_pages:
return _parse_full_docling_document(self.docling_document)
# No caching for documents without ID # No caching for documents without ID
if self.id is None: if self.id is None:
from docling_core.types.doc.document import DoclingDocument return _validate_without_pages(self.docling_document)
json_str = decompress_json(self.docling_document)
return DoclingDocument.model_validate_json(json_str)
return _get_cached_docling_document(self.id, self.docling_document) return _get_cached_docling_document(self.id, self.docling_document)

View file

@ -397,35 +397,6 @@ class ChunkRepository:
for rec in results 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( async def _process_search_results(
self, query_result: "pd.DataFrame | LanceQueryBuilder" self, query_result: "pd.DataFrame | LanceQueryBuilder"
) -> list[tuple[Chunk, float]]: ) -> list[tuple[Chunk, float]]: