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

extract picture bytes to document_items.picture_data at ingest, strip them from the docling_document blob, and add 0.45.0 migration to backfill existing

Browse source 
databases
This commit is contained in:
Yiorgis Gozadinos 2026-04-27 16:09:19 +03:00
parent 191cdc636b
commit 6a77ce92a9
No known key found for this signature in database
13 changed files with 572 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -4,6 +4,11 @@
### Added
- **Storage column for embedded picture bytes.** `DocumentItemRecord` gains a `picture_data: bytes | None` column (Arrow `large_binary`) to hold per-`PictureItem` image bytes addressable by `(document_id, self_ref)`. New repository accessors `get_picture_bytes` and `get_pictures_for_chunk` expose them; the existing items read paths (`get_all_items`, `get_all_items_grouped`, `get_items_in_range`, `_record_to_item`) now project an explicit lightweight column set so context expansion and the analysis-sandbox `items.jsonl` build never pull picture bytes into memory. Existing databases pick up the column via the `0.45.0` migration alongside the picture-byte backfill (see below). Foundation for upcoming vision-in-context retrieval; not yet wired into ingestion or search.
- **Embedded picture bytes captured at ingestion.** `extract_items` now decodes each `PictureItem.image.uri` data URI into raw bytes and writes them to `document_items.picture_data` so per-figure lookups don't require decompressing the full docling blob. The same path also surfaces VLM-generated picture descriptions (`meta.description.text`) into `DocumentItem.text` so picture-only chunks survive `expand_with_items`' text filter. The `0.45.0` migration adds the `picture_data` column to existing databases and backfills it by extracting bytes out of `docling_document`, stripping picture URIs from the structure blob in the process; `compress_docling_split` does the same for new ingests so the structure stays lean. Rebuild and update flows snapshot picture bytes via the new `DocumentItemRepository.get_all_picture_data` accessor before re-extraction so a re-chunk doesn't drop them.
### Fixed
- **docling-serve picture-image extraction.** docling-serve never populated `PictureItem.image` when called with `image_export_mode="embedded"` (per upstream issue [docling-project/docling-serve#576](https://github.com/docling-project/docling-serve/issues/576) — picture-image generation is only triggered when the server's `image_export_mode == "referenced"`). The converter now switches to `image_export_mode="referenced"` + `target_type="zip"` whenever picture images are requested, parses the returned zip, and rehydrates `artifacts/<filename>` URIs back into `data:<mime>;base64,...` URIs so downstream code sees the same shape as docling-local. The previously-`xfail`ed `test_convert_pdf_with_picture_images` integration test now passes.
## [0.44.0] - 2026-04-29

2
evaluations/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-evals"
description = "Benchmarking and evaluation scripts for haiku.rag"
version = "0.44.0"
version = "0.45.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
requires-python = ">=3.12"

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

@ -96,10 +96,22 @@ async def _update_document_with_chunks(
await client.chunk_repository.create(chunks)
# Replace document items when a new DoclingDocument is provided
# Replace document items when a new DoclingDocument is provided.
# Snapshot existing picture bytes first so they survive the
# delete-and-re-extract cycle when the live docling has already had
# its picture URIs stripped (rebuild / round-trip scenarios).
if docling_document is not None:
existing_picture_data = (
await client.document_item_repository.get_all_picture_data(
updated_doc.id
)
)
await client.document_item_repository.delete_by_document_id(updated_doc.id)
items = extract_items(updated_doc.id, docling_document)
items = extract_items(
updated_doc.id,
docling_document,
existing_picture_data=existing_picture_data,
)
await client.document_item_repository.create_items(updated_doc.id, items)
if client._config.storage.auto_vacuum:

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

@ -195,13 +195,23 @@ async def _flush_rebuild_batch(
if chunks:
await client.chunk_repository.create(chunks)
# Repopulate document items from stored docling data
# Repopulate document items from stored docling data. The stored docling
# blob has had its picture URIs stripped (compress_docling_split), so
# re-extracting from it would lose picture_data; snapshot the existing
# bytes per document and merge them back during extraction.
for doc in documents:
assert doc.id is not None
docling_doc = doc.get_docling_document()
if docling_doc is not None:
existing_picture_data = (
await client.document_item_repository.get_all_picture_data(doc.id)
)
await client.document_item_repository.delete_by_document_id(doc.id)
items = extract_items(doc.id, docling_doc)
items = extract_items(
doc.id,
docling_doc,
existing_picture_data=existing_picture_data,
)
await client.document_item_repository.create_items(doc.id, items)

10
haiku_rag_slim/haiku/rag/store/compression.py
View file

@ -34,6 +34,12 @@ def decompress_json(data: bytes) -> str:
def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
"""Split a DoclingDocument JSON into structure and pages, compress both with zstd.
Picture image URIs are stripped from the structure blob — they are stored on
the corresponding ``document_items.picture_data`` rows and don't need to be
duplicated inside the structure JSON. ``ImageRef.uri`` is required when the
field is present, so each picture's ``image`` is set to ``None`` rather than
partially mutated to keep the JSON re-validating cleanly.
Returns:
Tuple of (structure_bytes, pages_bytes). pages_bytes is None if the
document has no page images.
@ -41,6 +47,10 @@ def compress_docling_split(json_str: str) -> tuple[bytes, bytes | None]:
data = json.loads(json_str)
pages = data.pop("pages", None)
for picture in data.get("pictures") or []:
if isinstance(picture, dict):
picture["image"] = None
structure_bytes = _zstd_compress(json.dumps(data).encode("utf-8"))
pages_bytes = None

81
haiku_rag_slim/haiku/rag/store/models/document_item.py
View file

@ -1,9 +1,11 @@
import base64
import binascii
from typing import TYPE_CHECKING
from pydantic import BaseModel
if TYPE_CHECKING:
from docling_core.types.doc.document import DoclingDocument, NodeItem
from docling_core.types.doc.document import DoclingDocument, NodeItem, PictureItem
class DocumentItem(BaseModel):
@ -16,13 +18,65 @@ class DocumentItem(BaseModel):
picture_data: bytes | None = None
def _picture_description_text(item: "PictureItem") -> str | None:
"""Return the VLM-generated description text for a PictureItem, if any.
Tries the modern ``meta.description`` location first (docling 2.91+) and
falls back to ``annotations`` entries that carry a ``text`` field
(PictureDescriptionData and similar).
"""
meta = getattr(item, "meta", None)
if meta is not None:
description = getattr(meta, "description", None)
if description is not None:
text = getattr(description, "text", None)
if isinstance(text, str) and text.strip():
return text
annotations = getattr(item, "annotations", None) or []
for ann in annotations:
text = getattr(ann, "text", None)
if isinstance(text, str) and text.strip():
return text
return None
def _decode_picture_bytes(item: "PictureItem") -> bytes | None:
"""Decode a PictureItem's embedded image into raw bytes.
Reads ``item.image.uri`` and base64-decodes it when it is a ``data:`` URI.
Returns None for items whose image is absent, stripped, or not a data URI
(e.g. file references). Tolerant of malformed data — returns None on any
decode failure rather than raising.
"""
image = getattr(item, "image", None)
if image is None:
return None
uri = getattr(image, "uri", None)
if uri is None:
return None
uri_str = str(uri)
if not uri_str.startswith("data:"):
return None
try:
_, encoded = uri_str.split(",", 1)
except ValueError:
return None
try:
return base64.b64decode(encoded, validate=False)
except (ValueError, binascii.Error):
return None
def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None:
"""Extract text content from a DocItem.
Handles different item types:
- TextItem, SectionHeaderItem, etc.: Use .text attribute
- TableItem: Use export_to_markdown() for table content
- PictureItem: Use export_to_markdown() with PLACEHOLDER mode to avoid base64
- PictureItem: Prefer the VLM description (when picture_description is on)
so pictures carry meaningful prose into chunk text and survive
``expand_with_items``' ``if item.text:`` filter; otherwise fall back to
a placeholder markdown export (no base64).
"""
from docling_core.types.doc.base import ImageRefMode
from docling_core.types.doc.document import PictureItem, TableItem
@ -31,6 +85,8 @@ def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str |
return text
if isinstance(item, PictureItem):
if description := _picture_description_text(item):
return description
return item.export_to_markdown(
docling_doc,
image_mode=ImageRefMode.PLACEHOLDER,
@ -51,7 +107,9 @@ def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str |
def extract_items(
document_id: str, docling_doc: "DoclingDocument"
document_id: str,
docling_doc: "DoclingDocument",
existing_picture_data: dict[str, bytes] | None = None,
) -> list[DocumentItem]:
"""Extract document items from a DoclingDocument for the items table.
@ -59,7 +117,17 @@ def extract_items(
self_ref, label, pre-rendered text, and page numbers from provenance.
Items are stored as docling produces them — container items (e.g., list_item)
may have empty text with content in their children.
For PictureItems, the embedded image is decoded from ``image.uri`` (a base64
data URI) and stored on ``DocumentItem.picture_data``. When the live docling
has already had its picture URIs stripped (rebuild / re-extract scenarios
where the docling structure round-trips through the compressed blob), a
fall-back lookup against ``existing_picture_data`` (keyed by ``self_ref``)
preserves the bytes that were captured at original ingest time.
"""
from docling_core.types.doc.document import PictureItem
existing = existing_picture_data or {}
items: list[DocumentItem] = []
for position, (item, _level) in enumerate(docling_doc.iterate_items()):
@ -75,6 +143,12 @@ def extract_items(
if page_no is not None and page_no not in page_numbers:
page_numbers.append(page_no)
picture_data: bytes | None = None
if isinstance(item, PictureItem):
picture_data = _decode_picture_bytes(item)
if picture_data is None:
picture_data = existing.get(item.self_ref)
items.append(
DocumentItem(
document_id=document_id,
@ -83,6 +157,7 @@ def extract_items(
label=label_str,
text=text,
page_numbers=sorted(page_numbers),
picture_data=picture_data,
)
)

22
haiku_rag_slim/haiku/rag/store/repositories/document_item.py
View file

@ -154,6 +154,28 @@ class DocumentItemRepository:
return None
return rows[0].get("picture_data")
async def get_all_picture_data(self, document_id: str) -> dict[str, bytes]:
"""Snapshot every picture row's bytes for a single document.
Returns ``{self_ref: picture_data}`` for every row whose
``picture_data`` is non-null. Used by rebuild / update flows to
preserve picture bytes across a delete-and-re-extract cycle when the
live docling document has already been stripped of its picture URIs.
"""
safe_id = escape_sql_string(document_id)
rows = await (
self.store.document_items_table.query()
.select(["self_ref", "picture_data"])
.where(f"document_id = '{safe_id}'")
.to_list()
)
result: dict[str, bytes] = {}
for row in rows:
data = row.get("picture_data")
if data:
result[row["self_ref"]] = data
return result
async def get_pictures_for_chunk(
self, document_id: str, refs: list[str]
) -> dict[str, bytes]:

5
haiku_rag_slim/haiku/rag/store/upgrades/__init__.py
View file

@ -84,8 +84,13 @@ from haiku.rag.store.upgrades.v0_38_0 import (
from haiku.rag.store.upgrades.v0_40_0 import (
upgrade_populate_document_items as upgrade_0_40_0_document_items,
)
from haiku.rag.store.upgrades.v0_45_0 import (
upgrade_extract_picture_bytes as upgrade_0_45_0_extract_picture_bytes,
)
upgrades.append(upgrade_0_20_0_docling)
upgrades.append(upgrade_0_23_1_contextualize)
upgrades.append(upgrade_0_25_0_compress)
upgrades.append(upgrade_0_38_0_split_pages)
upgrades.append(upgrade_0_40_0_document_items)
upgrades.append(upgrade_0_45_0_extract_picture_bytes)

197
haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py Normal file
View file

@ -0,0 +1,197 @@
import base64
import binascii
import json
import logging
import pyarrow as pa
from haiku.rag.store.engine import DocumentItemRecord, Store
from haiku.rag.store.upgrades import Upgrade
from haiku.rag.utils import escape_sql_string
logger = logging.getLogger(__name__)
BATCH_SIZE = 5
async def _ensure_picture_data_column(store: Store) -> None:
"""Add ``document_items.picture_data`` (large_binary) if it doesn't exist.
Idempotent: re-runs and fresh DBs that already declare the column via
``_init_tables`` see the column in the schema and skip.
"""
arrow_schema = await store.document_items_table.schema()
if any(field.name == "picture_data" for field in arrow_schema):
return
logger.info("Adding picture_data column to document_items table")
await store.document_items_table.add_columns(
pa.schema([pa.field("picture_data", pa.large_binary())])
)
async def _apply_extract_picture_bytes(store: Store) -> None: # pragma: no cover
"""Add the ``picture_data`` column to ``document_items`` (if missing),
backfill it from existing docling blobs, and strip the inline picture
URIs out of those blobs.
Idempotent: documents whose blob already has every ``pictures[i].image``
set to ``None`` (already migrated, or never had bytes) are skipped without
error. Documents missing matching items rows (legacy DBs that predated the
v0.40.0 migration) are also skipped — running the v0.40.0 migration first
populates the rows, then this one fills in the bytes.
"""
from haiku.rag.store.compression import compress_json, decompress_json
await _ensure_picture_data_column(store)
ids = (await store.documents_table.query().select(["id"]).to_arrow()).to_pylist()
ids = [row["id"] for row in ids]
if not ids:
logger.info("No documents to backfill picture_data for")
return
total = len(ids)
logger.info(
"Backfilling picture_data and stripping URIs across %d documents", total
)
backfilled = 0
blob_only = 0
skipped = 0
for batch_start in range(0, total, BATCH_SIZE):
batch_ids = ids[batch_start : batch_start + BATCH_SIZE]
for doc_id in batch_ids:
safe_id = escape_sql_string(doc_id)
rows = await (
store.documents_table.query()
.select(["id", "docling_document"])
.where(f"id = '{safe_id}'")
.limit(1)
.to_list()
)
if not rows:
skipped += 1
continue
blob = rows[0].get("docling_document")
if not isinstance(blob, bytes):
skipped += 1
continue
try:
data = json.loads(decompress_json(blob))
except Exception:
logger.warning(
"Could not decompress docling blob for document %s; skipping",
doc_id,
)
skipped += 1
continue
pictures = data.get("pictures") or []
updates: list[tuple[str, bytes]] = []
modified = False
for picture in pictures:
if not isinstance(picture, dict):
continue
self_ref = picture.get("self_ref")
image = picture.get("image")
if image is None or self_ref is None:
continue
# We're going to strip every picture image from the blob.
# Whether or not we successfully decode the URI to bytes, the
# blob gets normalised. modified=True for any picture that
# had a non-null image — that signals "blob will change".
modified = True
uri = image.get("uri") if isinstance(image, dict) else None
if isinstance(uri, str) and uri.startswith("data:"):
try:
_, encoded = uri.split(",", 1)
updates.append((self_ref, base64.b64decode(encoded)))
except (ValueError, binascii.Error):
pass
picture["image"] = None
if not modified:
# Already stripped (idempotent re-run) or no pictures present.
continue
wrote_items = False
if updates:
# Find the matching items rows so we can preserve their
# position/label/text/page_numbers and just attach bytes.
self_refs = [u[0] for u in updates]
ref_clause = ", ".join(f"'{escape_sql_string(r)}'" for r in self_refs)
existing_items = await (
store.document_items_table.query()
.where(f"document_id = '{safe_id}' AND self_ref IN ({ref_clause})")
.to_list()
)
existing_by_ref = {r["self_ref"]: r for r in existing_items}
new_records: list[DocumentItemRecord] = []
for ref, img_bytes in updates:
existing = existing_by_ref.get(ref)
if existing is None:
# No item row for this self_ref — likely a legacy DB
# where v0.40.0 didn't run. Skip; rerun migrations.
continue
new_records.append(
DocumentItemRecord(
document_id=doc_id,
position=existing["position"],
self_ref=ref,
label=existing.get("label", ""),
text=existing.get("text", ""),
page_numbers=existing.get("page_numbers", "[]"),
picture_data=img_bytes,
)
)
if new_records:
await (
store.document_items_table.merge_insert(
["document_id", "self_ref"]
)
.when_matched_update_all()
.when_not_matched_insert_all()
.execute(new_records)
)
wrote_items = True
new_structure_bytes = compress_json(json.dumps(data))
await store.documents_table.update(
{"docling_document": new_structure_bytes},
where=f"id = '{safe_id}'",
)
if wrote_items:
backfilled += 1
else:
blob_only += 1
done = batch_start + (batch_ids.index(doc_id) + 1)
if done % 10 == 0 or done == total:
logger.info(
"Progress: %d/%d (%d backfilled, %d blob-only, %d skipped)",
done,
total,
backfilled,
blob_only,
skipped,
)
logger.info(
"Picture backfill complete: %d backfilled, %d blob-only, %d skipped of %d",
backfilled,
blob_only,
skipped,
total,
)
upgrade_extract_picture_bytes = Upgrade(
version="0.45.0",
apply=_apply_extract_picture_bytes,
description="Add picture_data column, backfill from docling_document, strip picture URIs",
)

2
haiku_rag_slim/pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag-slim"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
version = "0.44.0"
version = "0.45.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }

4
pyproject.toml
View file

@ -2,7 +2,7 @@
name = "haiku.rag"
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
version = "0.44.0"
version = "0.45.0"
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
license = { text = "MIT" }
readme = { file = "README.md", content-type = "text/markdown" }
@ -30,7 +30,7 @@ classifiers = [
]
dependencies = [
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.44.0",
"haiku.rag-slim[docling,voyageai,mxbai,cohere,zeroentropy,tui]==0.45.0",
]
[project.scripts]

222
tests/store/test_document_items.py
View file

@ -431,3 +431,225 @@ class TestPictureDataStorage:
async with HaikuRAG(temp_db_path, create=True) as rag:
schema = await rag.store.document_items_table.schema()
assert "picture_data" in {f.name for f in schema}
def _docling_doc_with_picture():
"""Build a tiny DoclingDocument with one PictureItem carrying real PNG bytes
via ImageRef.from_pil. Used by the picture-extraction tests."""
from docling_core.types.doc.document import DoclingDocument, ImageRef
from docling_core.types.doc.labels import DocItemLabel
from PIL import Image as PilImageModule
img = PilImageModule.new("RGB", (8, 8), "red")
doc = DoclingDocument(name="test-with-picture")
doc.add_text(label=DocItemLabel.PARAGRAPH, text="Hello world")
doc.add_picture(image=ImageRef.from_pil(img, dpi=72))
return doc
class TestExtractItemsPictureBytes:
"""A2b: extract_items decodes picture image bytes from data URIs."""
def test_decodes_picture_bytes_from_live_doc(self):
doc = _docling_doc_with_picture()
items = extract_items("doc-1", doc)
picture_items = [i for i in items if i.label == "picture"]
assert len(picture_items) == 1
data = picture_items[0].picture_data
assert data is not None and len(data) > 0
# PNG magic header — confirms we round-tripped real bytes, not a mangled URI.
assert data.startswith(b"\x89PNG")
def test_existing_picture_data_used_when_image_stripped(self):
"""Rebuild round-trip: live docling has image=None, snapshot fills the gap."""
doc = _docling_doc_with_picture()
for picture in doc.pictures:
picture.image = None
snapshot = {"#/pictures/0": b"snapshot-picture-bytes"}
items = extract_items("doc-1", doc, existing_picture_data=snapshot)
picture_items = [i for i in items if i.label == "picture"]
assert picture_items[0].picture_data == b"snapshot-picture-bytes"
class TestExtractItemTextDescription:
"""A2b: extract_item_text returns VLM description text for PictureItems."""
def test_returns_description_text_when_present(self):
from docling_core.types.doc.document import (
DescriptionAnnotation,
DoclingDocument,
)
doc = DoclingDocument(name="t")
doc.add_picture(
annotations=[
DescriptionAnnotation(text="A small red square", provenance="test")
]
)
items = extract_items("doc-1", doc)
picture_items = [i for i in items if i.label == "picture"]
assert picture_items[0].text == "A small red square"
class TestCompressDoclingSplitStripsPictureUris:
"""A2b: compress_docling_split removes inline picture URIs from the structure."""
def test_picture_image_set_to_none_in_structure(self):
import json
from haiku.rag.store.compression import (
compress_docling_split,
decompress_json,
)
doc_json = {
"schema_name": "DoclingDocument",
"version": "1.10.0",
"name": "test",
"pictures": [
{
"self_ref": "#/pictures/0",
"image": {
"mimetype": "image/png",
"uri": "data:image/png;base64,abc",
},
},
{"self_ref": "#/pictures/1", "image": None},
],
"pages": {},
}
structure_bytes, pages_bytes = compress_docling_split(json.dumps(doc_json))
decoded = json.loads(decompress_json(structure_bytes))
for pic in decoded["pictures"]:
assert pic["image"] is None
assert pages_bytes is None # no pages in this fixture
@pytest.mark.asyncio
class TestPictureDataMigrationBackfill:
"""A2b: v0.45.0 migration backfills picture_data and strips URIs from blobs."""
async def test_backfill_populates_column_and_strips_blob(self, temp_db_path):
import base64
import json
from haiku.rag.store.compression import compress_json, decompress_json
from haiku.rag.store.engine import DocumentItemRecord, DocumentRecord
fake_png = b"\x89PNG\r\n\x1a\nlegacy-picture-bytes-for-test"
data_uri = "data:image/png;base64," + base64.b64encode(fake_png).decode("ascii")
blob_data = {
"schema_name": "DoclingDocument",
"version": "1.10.0",
"name": "legacy",
"pictures": [
{
"self_ref": "#/pictures/0",
"label": "picture",
"image": {"mimetype": "image/png", "uri": data_uri},
},
],
}
blob_bytes = compress_json(json.dumps(blob_data))
# Build a legacy-state DB at v0.44.0 *without* the picture_data column,
# mirroring users coming from main's 0.44.0 release. The 0.45.0
# migration must add the column AND backfill it from the blob in one
# pass.
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
await store.set_haiku_version("0.44.0")
await store.documents_table.add(
[
DocumentRecord(
id="legacy-doc",
content="legacy",
docling_document=blob_bytes,
)
]
)
# Items row exists (v0.40.0 would have placed it) but no picture_data yet.
await store.document_items_table.add(
[
DocumentItemRecord(
document_id="legacy-doc",
position=0,
self_ref="#/pictures/0",
label="picture",
text="",
page_numbers="[]",
)
]
)
# Drop the column so the migration's column-add path is exercised.
await store.document_items_table.drop_columns(["picture_data"])
schema_before = await store.document_items_table.schema()
assert "picture_data" not in {f.name for f in schema_before}
async with Store(temp_db_path, skip_migration_check=True) as store:
applied = await store.migrate()
assert any("picture" in d.lower() for d in applied)
# Column was added by the migration
schema_after = await store.document_items_table.schema()
assert "picture_data" in {f.name for f in schema_after}
# picture_data backfilled with the legacy bytes
rows = await (
store.document_items_table.query()
.select(["self_ref", "picture_data"])
.where("document_id = 'legacy-doc'")
.to_list()
)
picture_rows = [r for r in rows if r["self_ref"] == "#/pictures/0"]
assert len(picture_rows) == 1
assert picture_rows[0]["picture_data"] == fake_png
# docling_document blob has been re-compressed with image=None
doc_rows = await (
store.documents_table.query()
.select(["docling_document"])
.where("id = 'legacy-doc'")
.to_list()
)
decoded = json.loads(decompress_json(doc_rows[0]["docling_document"]))
assert decoded["pictures"][0]["image"] is None
@pytest.mark.asyncio
class TestPictureDataPreservedThroughRoundTrip:
"""A2b: snapshot/merge keeps picture bytes through update / rebuild cycles."""
async def test_update_preserves_picture_data_when_blob_round_tripped(
self, temp_db_path
):
"""update_document on a docling pulled from the (stripped) blob must
not clobber picture_data — the snapshot/merge in
_update_document_with_chunks handles it."""
from haiku.rag.client.documents import (
_store_document_with_chunks,
_update_document_with_chunks,
)
from haiku.rag.store.models.document import Document
docling_doc = _docling_doc_with_picture()
async with HaikuRAG(temp_db_path, create=True) as rag:
document = Document(content="Hello world", uri="test://doc")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert created.id is not None
original = await rag.document_item_repository.get_all_picture_data(
created.id
)
assert original.get("#/pictures/0") is not None
# Re-load the docling from the stored blob — pictures now have image=None
from_blob = created.get_docling_document()
assert from_blob is not None
assert all(p.image is None for p in from_blob.pictures)
await _update_document_with_chunks(rag, created, [], from_blob)
after = await rag.document_item_repository.get_all_picture_data(created.id)
assert after.get("#/pictures/0") == original.get("#/pictures/0")

6
uv.lock
View file

@ -1418,7 +1418,7 @@ wheels = [
[[package]]
name = "haiku-rag"
version = "0.44.0"
version = "0.45.0"
source = { editable = "." }
dependencies = [
{ name = "haiku-rag-slim", extra = ["cohere", "docling", "mxbai", "tui", "voyageai", "zeroentropy"] },
@ -1475,7 +1475,7 @@ dev = [
[[package]]
name = "haiku-rag-evals"
version = "0.44.0"
version = "0.45.0"
source = { editable = "evaluations" }
dependencies = [
{ name = "datasets" },
@ -1500,7 +1500,7 @@ requires-dist = [
[[package]]
name = "haiku-rag-slim"
version = "0.44.0"
version = "0.45.0"
source = { editable = "haiku_rag_slim" }
dependencies = [
{ name = "docling-core" },