From 6fdb15e3b23ddcc2efc4e05f153e656ee9477902 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Tue, 7 Apr 2026 11:19:04 +0300
Subject: [PATCH] Add DocumentRepository.get_docling_data() for lazy loading
 docling in expand_context

---
 haiku_rag_slim/haiku/rag/client.py            | 19 +++----
 .../haiku/rag/store/repositories/document.py  | 24 +++++++++
 tests/test_document.py                        | 50 +++++++++++++++++++
 3 files changed, 82 insertions(+), 11 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py
index aff7e91c..7b7d0dfc 100644
--- a/haiku_rag_slim/haiku/rag/client.py
+++ b/haiku_rag_slim/haiku/rag/client.py
@@ -1114,19 +1114,16 @@ class HaikuRAG:
                 expanded_results.extend(doc_results)
                 continue
 
-            # Fetch the document to get DoclingDocument
-            doc = await self.get_document_by_id(doc_id)
-            if doc is None:
-                expanded_results.extend(doc_results)
-                continue
-
-            docling_doc = doc.get_docling_document()
-
-            # Check if we can use DoclingDocument-based expansion
-            has_docling = docling_doc is not None
             has_refs = any(r.doc_item_refs for r in doc_results)
+            docling_doc = None
 
-            if has_docling and has_refs:
+            if has_refs:
+                # Only load docling data when refs exist (skips content blob)
+                doc = await self.document_repository.get_docling_data(doc_id)
+                if doc is not None:
+                    docling_doc = doc.get_docling_document()
+
+            if docling_doc is not None and has_refs:
                 # Use DoclingDocument-based expansion
                 expanded = await self._expand_with_docling(
                     doc_results,
diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py
index f40a6ed3..c829bd57 100644
--- a/haiku_rag_slim/haiku/rag/store/repositories/document.py
+++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py
@@ -90,6 +90,30 @@ class DocumentRepository:
 
         return self._record_to_document(results[0])
 
+    _DOCLING_COLUMNS = ["id", "docling_document", "docling_version"]
+
+    async def get_docling_data(self, entity_id: str) -> Document | None:
+        """Get a document with only docling data loaded (skips content blob)."""
+        safe_id = _escape_sql_string(entity_id)
+        results = list(
+            self.store.documents_table.search()
+            .select(self._DOCLING_COLUMNS)
+            .where(f"id = '{safe_id}'")
+            .limit(1)
+            .to_list()
+        )
+
+        if not results:
+            return None
+
+        row = results[0]
+        return Document(
+            id=row["id"],
+            content="",
+            docling_document=row.get("docling_document"),
+            docling_version=row.get("docling_version"),
+        )
+
     async def update(self, entity: Document) -> Document:
         """Update an existing document."""
         self.store._assert_writable()
diff --git a/tests/test_document.py b/tests/test_document.py
index 3f43e12b..9e4c92bc 100644
--- a/tests/test_document.py
+++ b/tests/test_document.py
@@ -229,6 +229,56 @@ def test_document_get_docling_document_no_id_no_cache():
     assert doc1 is not doc2
 
 
+@pytest.mark.asyncio
+async def test_get_docling_data_loads_only_docling_columns(
+    qa_corpus: Dataset, temp_db_path
+):
+    """get_docling_data returns docling blob without loading content."""
+    import json
+
+    from haiku.rag.store.compression import compress_json
+
+    store = Store(temp_db_path, create=True)
+    doc_repo = DocumentRepository(store)
+
+    doc_json = {
+        "name": "test_doc",
+        "texts": [],
+        "tables": [],
+        "pictures": [],
+        "groups": [],
+        "body": {"self_ref": "#/body", "children": []},
+        "furniture": {"self_ref": "#/furniture", "children": []},
+    }
+    compressed = compress_json(json.dumps(doc_json))
+
+    doc = Document(
+        content=qa_corpus[0]["document_extracted"],
+        uri="https://example.com/doc.txt",
+        docling_document=compressed,
+        docling_version="2.1.0",
+    )
+    created = await doc_repo.create(doc)
+    assert created.id is not None
+
+    result = await doc_repo.get_docling_data(created.id)
+    assert result is not None
+    assert result.id == created.id
+    assert result.content == ""
+    assert result.docling_document == compressed
+    assert result.docling_version == "2.1.0"
+
+    # Verify docling document can be parsed
+    docling_doc = result.get_docling_document()
+    assert docling_doc is not None
+    assert docling_doc.name == "test_doc"
+
+    # Non-existent ID returns None
+    assert await doc_repo.get_docling_data("nonexistent-id") is None
+
+    store.close()
+
+
 @pytest.mark.asyncio
 async def test_document_get_by_uri_with_special_characters(
     qa_corpus: Dataset, temp_db_path