diff --git a/haiku_rag_slim/haiku/rag/store/models/document.py b/haiku_rag_slim/haiku/rag/store/models/document.py index 7f73f904..5dd6105a 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document.py +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 3320c3ec..d35297cc 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/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) diff --git a/tests/test_document.py b/tests/test_document.py index f2ccb4c8..3b2c9fa3 100644 --- a/tests/test_document.py +++ b/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