LRU cache for DoclingDocument

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

@@ -1,12 +1,33 @@
 from datetime import datetime
 from typing import TYPE_CHECKING
 
+from cachetools import LRUCache
 from pydantic import BaseModel, Field
 
 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, json_str: str) -> "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
+
+    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.
@@ -25,12 +46,18 @@ class Document(BaseModel):
     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.
+            The parsed DoclingDocument, or None if not stored or no ID.
         """
         if self.docling_document_json is None:
             return None
 
-        from docling_core.types.doc.document import DoclingDocument
+        # No caching for documents without ID
+        if self.id is None:
+            from docling_core.types.doc.document import DoclingDocument
 
-        return DoclingDocument.model_validate_json(self.docling_document_json)
+            return DoclingDocument.model_validate_json(self.docling_document_json)
+
+        return _get_cached_docling_document(self.id, self.docling_document_json)

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

@@ -92,8 +92,13 @@ class DocumentRepository:
 
     async def update(self, entity: Document) -> Document:
         """Update an existing document."""
+        from haiku.rag.store.models.document import invalidate_docling_document_cache
+
         assert entity.id, "Document ID is required for update"
 
+        # Invalidate cache before update
+        invalidate_docling_document_cache(entity.id)
+
         # Update timestamp
         now = datetime.now().isoformat()
         entity.updated_at = datetime.fromisoformat(now)
@@ -116,11 +121,16 @@ class DocumentRepository:
 
     async def delete(self, entity_id: str) -> bool:
         """Delete a document by its ID."""
+        from haiku.rag.store.models.document import invalidate_docling_document_cache
+
         # Check if document exists
         doc = await self.get_by_id(entity_id)
         if doc is None:
             return False
 
+        # Invalidate cache before delete
+        invalidate_docling_document_cache(entity_id)
+
         # Delete associated chunks first
         await self.chunk_repository.delete_by_document_id(entity_id)
 

tests/test_document.py

@@ -181,3 +181,86 @@ def test_document_get_docling_document_none():
 
     assert document.docling_document_json is None
     assert document.get_docling_document() is None
+
+
+def test_document_get_docling_document_caching():
+    """Test that get_docling_document uses LRU cache keyed by document ID."""
+    from haiku.rag.store.models.document import (
+        _docling_document_cache,
+        invalidate_docling_document_cache,
+    )
+
+    doc_json = {
+        "name": "test_doc",
+        "texts": [
+            {
+                "self_ref": "#/texts/0",
+                "text": "Test text",
+                "orig": "Test text",
+                "label": "paragraph",
+            },
+        ],
+        "tables": [],
+        "pictures": [],
+        "groups": [],
+        "body": {"self_ref": "#/body", "children": []},
+        "furniture": {"self_ref": "#/furniture", "children": []},
+    }
+
+    import json
+
+    json_str = json.dumps(doc_json)
+
+    # Clear cache to get clean state
+    _docling_document_cache.clear()
+
+    document = Document(
+        id="test-doc-id", content="Test content", docling_document_json=json_str
+    )
+
+    # First call - not in cache
+    assert "test-doc-id" not in _docling_document_cache
+    doc1 = document.get_docling_document()
+    assert "test-doc-id" in _docling_document_cache
+
+    # Second call - cache hit, same object
+    doc2 = document.get_docling_document()
+    assert doc1 is doc2
+
+    # Invalidation removes from cache
+    invalidate_docling_document_cache("test-doc-id")
+    assert "test-doc-id" not in _docling_document_cache
+
+
+def test_document_get_docling_document_no_id_no_cache():
+    """Test that documents without ID don't use cache."""
+    from haiku.rag.store.models.document import _docling_document_cache
+
+    doc_json = {
+        "name": "test_doc",
+        "texts": [],
+        "tables": [],
+        "pictures": [],
+        "groups": [],
+        "body": {"self_ref": "#/body", "children": []},
+        "furniture": {"self_ref": "#/furniture", "children": []},
+    }
+
+    import json
+
+    json_str = json.dumps(doc_json)
+
+    # Clear cache
+    _docling_document_cache.clear()
+
+    # Document without ID
+    document = Document(content="Test content", docling_document_json=json_str)
+
+    doc1 = document.get_docling_document()
+    doc2 = document.get_docling_document()
+
+    # Cache should remain empty (no ID to cache by)
+    assert len(_docling_document_cache) == 0
+
+    # Each call parses fresh (different objects)
+    assert doc1 is not doc2