Load document blobs only when asked

CHANGELOG.md

@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Changed
+
+- `get_document_by_id` / `get_document_by_uri` return content and the mutable attributes only; the docling structure and page-image blobs are no longer loaded. Load them with `DocumentRepository.get_docling_data` / `get_pages_data`, or pass `include_blobs=True` to the repository method.
+
 ### Fixed
 
 - Ingester jobs failing with an `obstore` `PermissionDeniedError`, `UnauthenticatedError`, `UnknownConfigurationKeyError` or `InvalidPathError` are dead-lettered instead of retried to `max_attempts`.

docs/python.md

@@ -94,6 +94,14 @@ By URI:
 doc = await client.get_document_by_uri("file:///path/to/document.pdf")
 ```
 
+Both return content, uri, title and metadata. The multi-MB docling blobs are
+loaded separately:
+
+```python
+docling = await client.document_repository.get_docling_data(doc.id)
+pages = await client.document_repository.get_pages_data(doc.id)
+```
+
 List all documents:
 ```python
 docs = await client.list_documents(limit=10, offset=0)

haiku_rag_slim/haiku/rag/client/documents.py

@@ -856,7 +856,11 @@ async def update_document(
             "Provide one or the other, not both."
         )
 
-    existing_doc = await client.get_document_by_id(document_id)
+    # An update that only replaces content writes the record back as-is, so
+    # the blobs have to make the round trip.
+    existing_doc = await client.document_repository.get_by_id(
+        document_id, include_blobs=True
+    )
     if existing_doc is None:
         raise ValueError(f"Document with ID {document_id} not found")
 

haiku_rag_slim/haiku/rag/client/rebuild.py

@@ -187,7 +187,9 @@ async def _hydrate(
     """
     for light_doc in light_docs:
         assert light_doc.id is not None
-        doc = await client.get_document_by_id(light_doc.id)
+        doc = await client.document_repository.get_by_id(
+            light_doc.id, include_blobs=True
+        )
         if doc is None:
             continue
         assert doc.id is not None
@@ -822,7 +824,9 @@ async def _rebuild_full(
 
         # Fallback: rebuild from stored content. Now we need the full
         # record (content + docling_pages for the round-trip write).
-        doc = await client.get_document_by_id(light_doc.id)
+        doc = await client.document_repository.get_by_id(
+            light_doc.id, include_blobs=True
+        )
         if doc is None:
             continue
         assert doc.id is not None

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

@@ -182,11 +182,16 @@ class DocumentRepository:
             raise
         return documents
 
-    async def get_by_id(self, entity_id: str) -> Document | None:
+    _LIGHT_COLUMNS = ["id", "content"]
-        """Get a document by its ID."""
+
+    async def get_by_id(
+        self, entity_id: str, include_blobs: bool = False
+    ) -> Document | None:
+        """Get a document by its ID. `include_blobs` adds the docling blobs."""
         safe_id = escape_sql_string(entity_id)
+        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
         results = await query_to_pydantic(
-            self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
+            query if include_blobs else query.select(self._LIGHT_COLUMNS),
             DocumentRecord,
         )
 
@@ -373,7 +378,9 @@ class DocumentRepository:
         """Count documents with optional filtering (over document_meta columns)."""
         return await self.store.document_meta_table.count_rows(filter=filter)
 
-    async def get_by_uri(self, uri: str) -> Document | None:
+    async def get_by_uri(
+        self, uri: str, include_blobs: bool = False
+    ) -> Document | None:
         """Get a document by its URI (resolved via document_meta)."""
         escaped_uri = escape_sql_string(uri)
         meta_results = await query_to_pydantic(
@@ -388,8 +395,9 @@ class DocumentRepository:
 
         meta = meta_results[0]
         safe_id = escape_sql_string(meta.id)
+        query = self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1)
         doc_results = await query_to_pydantic(
-            self.store.documents_table.query().where(f"id = '{safe_id}'").limit(1),
+            query if include_blobs else query.select(self._LIGHT_COLUMNS),
             DocumentRecord,
         )
         if not doc_results:

tests/store/test_v0_58_0_migration.py

@@ -59,7 +59,7 @@ class TestV0_58_0Migration:
 
             # Full hydration still works (content + metadata + blobs intact).
             repo = DocumentRepository(store)
-            doc = await repo.get_by_id("doc-1")
+            doc = await repo.get_by_id("doc-1", include_blobs=True)
             assert doc is not None
             assert doc.content == "body one"
             assert doc.uri == "s3://b/one"

tests/test_client.py

@@ -984,13 +984,17 @@ async def test_metadata_only_update_does_not_advance_documents_table(temp_db_pat
         # The light document_meta table absorbs the updates.
         assert await client.store.document_meta_table.version() > meta_v0
 
-        # Reads still hydrate the full document (content + blobs + metadata).
+        # Reads still hydrate content and the mutable attributes together.
         fetched = await client.get_document_by_id(created.id)
         assert fetched is not None
         assert fetched.metadata["source_revision"] == "rev-5"
         assert fetched.title == "Title 5"
         assert fetched.content == "Body text"
-        assert fetched.get_docling_document() is not None
+
+        # And the untouched docling blob is still there.
+        docling = await client.document_repository.get_docling_data(created.id)
+        assert docling is not None
+        assert docling.get_docling_document() is not None
 
 
 async def test_delete_marks_vacuum_dirty(temp_db_path):
@@ -1202,7 +1206,9 @@ async def test_client_create_document_from_file_stores_docling_json(temp_db_path
             assert doc.docling_version is not None
 
             # Verify the stored document also has the JSON
-            retrieved = await client.get_document_by_id(doc.id)
+            retrieved = await client.document_repository.get_by_id(
+                doc.id, include_blobs=True
+            )
             assert retrieved is not None
             assert retrieved.docling_document == doc.docling_document
             assert retrieved.docling_version == doc.docling_version
@@ -2200,6 +2206,32 @@ async def test_update_document_with_url_prefixed_content(temp_db_path, monkeypat
         assert "New heading" in updated.content
 
 
+async def test_update_document_with_chunks_keeps_page_images(temp_db_path, monkeypatch):
+    """Replacing content and chunks without a docling document writes the stored
+    record back as-is, so its page rasters must survive the round trip."""
+    _patch_embed_chunks(monkeypatch)
+
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.create_document(content="initial body", uri="test://pages")
+        assert doc.id is not None
+
+        sentinel_pages = b"\x80SENTINEL_PAGE_BYTES"
+        await client.store.documents_table.update(
+            {"docling_pages": sentinel_pages}, where=f"id = '{doc.id}'"
+        )
+
+        await client.update_document(
+            doc.id,
+            content="replacement body",
+            chunks=[Chunk(content="replacement body")],
+        )
+
+        stored = await client.document_repository.get_by_id(doc.id, include_blobs=True)
+        assert stored is not None
+        assert stored.content == "replacement body"
+        assert stored.docling_pages == sentinel_pages
+
+
 async def test_rebuild_rechunk_with_url_prefixed_stored_content(
     temp_db_path, monkeypatch
 ):

tests/test_document.py

@@ -457,6 +457,75 @@ async def test_get_pages_data_loads_only_pages_column(
         assert await doc_repo.get_pages_data("nonexistent-id") is None
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_blobs", [False, True])
+async def test_document_get_by_id_docling_blobs(temp_db_path, include_blobs):
+    """get_by_id leaves the docling blobs out unless asked for them: a single
+    document's page rasters run to hundreds of MB."""
+    async with Store(temp_db_path, create=True) as store:
+        doc_repo = DocumentRepository(store)
+
+        created = await doc_repo.create(
+            Document(
+                content="the text",
+                uri="https://example.com/doc.pdf",
+                title="Test Document",
+                metadata={"key": "value"},
+                docling_document=b"structure-blob",
+                docling_pages=b"page-raster-blob",
+                docling_version="2.1.0",
+            )
+        )
+        assert created.id is not None
+
+        doc = await doc_repo.get_by_id(created.id, include_blobs=include_blobs)
+
+        assert doc is not None
+        assert doc.id == created.id
+        assert doc.content == "the text"
+        assert doc.uri == "https://example.com/doc.pdf"
+        assert doc.title == "Test Document"
+        assert doc.metadata == {"key": "value"}
+        if include_blobs:
+            assert doc.docling_document == b"structure-blob"
+            assert doc.docling_pages == b"page-raster-blob"
+            assert doc.docling_version == "2.1.0"
+        else:
+            assert doc.docling_document is None
+            assert doc.docling_pages is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("include_blobs", [False, True])
+async def test_document_get_by_uri_docling_blobs(temp_db_path, include_blobs):
+    """get_by_uri has the same projection as get_by_id."""
+    async with Store(temp_db_path, create=True) as store:
+        doc_repo = DocumentRepository(store)
+
+        created = await doc_repo.create(
+            Document(
+                content="the text",
+                uri="https://example.com/doc.pdf",
+                docling_document=b"structure-blob",
+                docling_pages=b"page-raster-blob",
+            )
+        )
+
+        doc = await doc_repo.get_by_uri(
+            "https://example.com/doc.pdf", include_blobs=include_blobs
+        )
+
+        assert doc is not None
+        assert doc.id == created.id
+        assert doc.content == "the text"
+        if include_blobs:
+            assert doc.docling_document == b"structure-blob"
+            assert doc.docling_pages == b"page-raster-blob"
+        else:
+            assert doc.docling_document is None
+            assert doc.docling_pages is None
+
+
 @pytest.mark.asyncio
 async def test_document_get_by_uri_with_special_characters(
     qa_corpus: list[dict[str, str]], temp_db_path

tests/test_rebuild.py

@@ -36,7 +36,9 @@ async def test_rebuild_full(qa_corpus: list[dict[str, str]], temp_db_path):
         assert doc.id in processed_ids
 
         # Verify DoclingDocument JSON is preserved after rebuild
-        doc_after = await client.document_repository.get_by_id(doc.id)
+        doc_after = await client.document_repository.get_by_id(
+            doc.id, include_blobs=True
+        )
         assert doc_after is not None
         assert doc_after.docling_document is not None
         assert doc_after.docling_version is not None
@@ -70,7 +72,9 @@ async def test_rebuild_embed_only(qa_corpus: list[dict[str, str]], temp_db_path)
         assert doc.id in processed_ids
 
         # DoclingDocument JSON should be unchanged (embed-only doesn't touch documents)
-        doc_after = await client.document_repository.get_by_id(doc.id)
+        doc_after = await client.document_repository.get_by_id(
+            doc.id, include_blobs=True
+        )
         assert doc_after is not None
         assert doc_after.docling_document == original_docling_json
 
@@ -498,7 +502,9 @@ async def test_rebuild_rechunk(qa_corpus: list[dict[str, str]], temp_db_path):
         assert doc.id in processed_ids
 
         # Document content should be unchanged, but docling JSON should be updated
-        doc_after = await client.document_repository.get_by_id(doc.id)
+        doc_after = await client.document_repository.get_by_id(
+            doc.id, include_blobs=True
+        )
         assert doc_after is not None
         assert doc_after.content == content_before
         assert doc_after.docling_document is not None
@@ -707,7 +713,7 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
         )
 
         from_blob = (
-            await rag.document_repository.get_by_id(created.id)
+            await rag.document_repository.get_by_id(created.id, include_blobs=True)
         ).get_docling_document()  # type: ignore[union-attr]
         assert from_blob is not None and from_blob.pictures
         # No description in the freshly-ingested doc
@@ -737,7 +743,7 @@ async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkey
         assert created.id in processed
 
         # The stored docling blob now has the description
-        after = await rag.document_repository.get_by_id(created.id)
+        after = await rag.document_repository.get_by_id(created.id, include_blobs=True)
         assert after is not None
         after_doc = after.get_docling_document()
         assert after_doc is not None and after_doc.pictures
@@ -799,7 +805,7 @@ async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkey
         # VLM was never called for this picture (it already had a description)
         assert called_with == [] or all(not d for d in called_with)
 
-        after = await rag.document_repository.get_by_id(created.id)
+        after = await rag.document_repository.get_by_id(created.id, include_blobs=True)
         assert after is not None
         after_doc = after.get_docling_document()
         assert after_doc is not None
@@ -1142,10 +1148,10 @@ async def test_hydrate_skips_documents_deleted_mid_rebuild(temp_db_path):
             Document(content="body", uri="test://gone")
         )
 
-        async def vanished(_document_id):
+        async def vanished(_document_id, include_blobs=False):
             return None
 
-        client.get_document_by_id = vanished  # type: ignore[method-assign]
+        client.document_repository.get_by_id = vanished  # type: ignore[method-assign]
 
         assert [doc async for doc in _hydrate(client, [stored])] == []
 
@@ -1438,10 +1444,10 @@ async def test_rebuild_full_skips_document_deleted_mid_rebuild(temp_db_path):
         doc = await client.create_document(content="doc that disappears")
         assert doc.id is not None
 
-        async def vanished(_document_id):
+        async def vanished(_document_id, include_blobs=False):
             return None
 
-        client.get_document_by_id = vanished  # type: ignore[method-assign]
+        client.document_repository.get_by_id = vanished  # type: ignore[method-assign]
 
         processed = [
             doc_id async for doc_id in client.rebuild_database(mode=RebuildMode.FULL)