afc5b02285
- Change DocumentRecord.docling_document_json (str) to docling_document (bytes) - Use large_binary Arrow type (64-bit offsets) to avoid 2GB column limit - Decompress in Document.get_docling_document() using gzip - Update DocumentRepository field mappings
69 lines
2.1 KiB
Python
69 lines
2.1 KiB
Python
from datetime import datetime
|
|
from typing import TYPE_CHECKING
|
|
|
|
from cachetools import LRUCache
|
|
from pydantic import BaseModel, Field
|
|
|
|
from haiku.rag.store.compression import decompress_json
|
|
|
|
if TYPE_CHECKING:
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
|
|
_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]
|
|
|
|
from docling_core.types.doc.document import DoclingDocument
|
|
|
|
json_str = decompress_json(compressed_data)
|
|
doc = DoclingDocument.model_validate_json(json_str)
|
|
_docling_document_cache[document_id] = doc
|
|
return doc
|
|
|
|
|
|
def invalidate_docling_document_cache(document_id: str) -> None:
|
|
"""Remove a document from the DoclingDocument cache."""
|
|
_docling_document_cache.pop(document_id, None)
|
|
|
|
|
|
class Document(BaseModel):
|
|
"""
|
|
Represents a document with an ID, content, and metadata.
|
|
"""
|
|
|
|
id: str | None = None
|
|
content: str
|
|
uri: str | None = None
|
|
title: str | None = None
|
|
metadata: dict = {}
|
|
docling_document: bytes | None = None
|
|
docling_version: str | None = None
|
|
created_at: datetime = Field(default_factory=datetime.now)
|
|
updated_at: datetime = Field(default_factory=datetime.now)
|
|
|
|
def get_docling_document(self) -> "DoclingDocument | None":
|
|
"""Parse and return the stored DoclingDocument.
|
|
|
|
Uses LRU cache (keyed by document ID) to avoid repeated parsing.
|
|
|
|
Returns:
|
|
The parsed DoclingDocument, or None if not stored or no ID.
|
|
"""
|
|
if self.docling_document is None:
|
|
return None
|
|
|
|
# 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 _get_cached_docling_document(self.id, self.docling_document)
|