danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Load document blobs only when asked

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-08-06 11:24:47 +03:00
parent 788a2fe731
commit ac9b2cbf81
No known key found for this signature in database
9 changed files with 157 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
CHANGELOG.md
View file

@ -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`.

8
docs/python.md
View file

@ -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)

6
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -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")

8
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -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

18
haiku_rag_slim/haiku/rag/store/repositories/document.py
View file

@ -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:

2
tests/store/test_v0_58_0_migration.py
View file

@ -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"

38
tests/test_client.py
View file

@ -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
):

69
tests/test_document.py
View file

@ -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

26
tests/test_rebuild.py
View file

@ -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)