diff --git a/CHANGELOG.md b/CHANGELOG.md index e3f88ab3..f563c08c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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/` URIs back into `data:;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 diff --git a/evaluations/pyproject.toml b/evaluations/pyproject.toml index 9a08d5f5..3a0cc8d6 100644 --- a/evaluations/pyproject.toml +++ b/evaluations/pyproject.toml @@ -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" diff --git a/haiku_rag_slim/haiku/rag/client/documents.py b/haiku_rag_slim/haiku/rag/client/documents.py index 1cd0861e..03098a4c 100644 --- a/haiku_rag_slim/haiku/rag/client/documents.py +++ b/haiku_rag_slim/haiku/rag/client/documents.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 9c8280f5..44c84956 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/store/compression.py b/haiku_rag_slim/haiku/rag/store/compression.py index 80abf712..5aecf464 100644 --- a/haiku_rag_slim/haiku/rag/store/compression.py +++ b/haiku_rag_slim/haiku/rag/store/compression.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/store/models/document_item.py b/haiku_rag_slim/haiku/rag/store/models/document_item.py index 7c781b41..744b64d8 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/models/document_item.py @@ -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, ) ) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py index 64c0d0e0..1e91d90b 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document_item.py @@ -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]: diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py index 82ef610b..faa8c05f 100644 --- a/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/upgrades/__init__.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py b/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py new file mode 100644 index 00000000..b1484490 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/upgrades/v0_45_0.py @@ -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", +) diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 8cd74aa9..70b71948 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -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" } diff --git a/pyproject.toml b/pyproject.toml index b96841d2..fe94c4c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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] diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index 68cf4a3a..862f401c 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -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") diff --git a/uv.lock b/uv.lock index 8dcdc94f..7f8f2d33 100644 --- a/uv.lock +++ b/uv.lock @@ -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" },