From ac9b2cbf817be1bfec30e0b860fd74e61606589a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 6 Aug 2026 11:24:47 +0300 Subject: [PATCH] Load document blobs only when asked --- CHANGELOG.md | 4 ++ docs/python.md | 8 +++ haiku_rag_slim/haiku/rag/client/documents.py | 6 +- haiku_rag_slim/haiku/rag/client/rebuild.py | 8 ++- .../haiku/rag/store/repositories/document.py | 18 +++-- tests/store/test_v0_58_0_migration.py | 2 +- tests/test_client.py | 38 +++++++++- tests/test_document.py | 69 +++++++++++++++++++ tests/test_rebuild.py | 26 ++++--- 9 files changed, 157 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f4d1e816..08483dfa 100644 --- a/CHANGELOG.md +++ b/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`. diff --git a/docs/python.md b/docs/python.md index abaecc1e..48075426 100644 --- a/docs/python.md +++ b/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) diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 8cb916e4..f0455198 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/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") diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 6ba9ca7b..639572de 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 8f332f0b..c31cdcae 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/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: - """Get a document by its ID.""" + _LIGHT_COLUMNS = ["id", "content"] + + 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: diff --git a/tests/store/test_v0_58_0_migration.py b/tests/store/test_v0_58_0_migration.py index 286b4d05..30fe510e 100644 --- a/tests/store/test_v0_58_0_migration.py +++ b/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" diff --git a/tests/test_client.py b/tests/test_client.py index 48de4ad5..db2da16f 100644 --- a/tests/test_client.py +++ b/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 ): diff --git a/tests/test_document.py b/tests/test_document.py index a45fd03d..4f10b9f2 100644 --- a/tests/test_document.py +++ b/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 diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index f76c5592..7e4a5ce6 100644 --- a/tests/test_rebuild.py +++ b/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)