From 9f4dda9254f6ec4d7f0621801a22885ef2c4950b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 24 Jun 2026 12:02:13 +0300 Subject: [PATCH] Reuse one MarkdownDocSerializer per document in item extraction --- CHANGELOG.md | 1 + .../haiku/rag/store/models/document_item.py | 56 +++++++-- tests/store/test_document_items.py | 106 ++++++++++++++++++ 3 files changed, 151 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9612b7fb..c3584c65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ - zstd compression uses a fresh `zstandard` compressor/decompressor per call instead of shared module-level singletons, fixing process segfaults when ingester workers compress documents concurrently (Python < 3.14). - Ingestion compresses the DoclingDocument structure directly from the in-memory dict, removing one full-size serialized copy from per-document peak memory. +- Document item extraction builds one `MarkdownDocSerializer` per document (reused across tables) and derives description-less picture text from captions, instead of constructing a serializer per picture and table. ## [0.61.0] - 2026-06-23 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 f86080e7..e05506cc 100644 --- a/haiku_rag_slim/haiku/rag/store/models/document_item.py +++ b/haiku_rag_slim/haiku/rag/store/models/document_item.py @@ -1,5 +1,6 @@ import base64 -from typing import TYPE_CHECKING +from collections.abc import Callable +from typing import TYPE_CHECKING, Any from pydantic import BaseModel @@ -53,18 +54,33 @@ def _decode_picture_bytes(item: "PictureItem") -> bytes | None: return base64.b64decode(encoded, validate=False) -def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | None: +def _picture_caption_text(item: "PictureItem", docling_doc: "DoclingDocument") -> str: + """Join a picture's caption texts with spaces, preserving word boundaries.""" + return " ".join( + text + for caption in item.captions + if (text := caption.resolve(docling_doc).text.strip()) + ) + + +def extract_item_text( + item: "NodeItem", + docling_doc: "DoclingDocument", + *, + get_serializer: Callable[[], Any] | None = None, +) -> 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 + - TableItem: serialize to markdown. ``get_serializer`` supplies a reused + ``MarkdownDocSerializer`` (see ``extract_items``); when absent a one-off + serializer is built so direct calls keep working. - 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). + the picture's caption text. """ - from docling_core.types.doc.base import ImageRefMode from docling_core.types.doc.document import PictureItem, TableItem if text := getattr(item, "text", None): @@ -73,15 +89,19 @@ def extract_item_text(item: "NodeItem", docling_doc: "DoclingDocument") -> str | if isinstance(item, PictureItem): if description := _picture_description_text(item): return description - return item.export_to_markdown( - docling_doc, - image_mode=ImageRefMode.PLACEHOLDER, - image_placeholder="", - ) + return _picture_caption_text(item, docling_doc) if isinstance(item, TableItem): try: - return item.export_to_markdown(docling_doc) + if get_serializer is None: + from docling_core.transforms.serializer.markdown import ( + MarkdownDocSerializer, + ) + + serializer = MarkdownDocSerializer(doc=docling_doc) + else: + serializer = get_serializer() + return serializer.serialize(item=item).text except Exception: pass @@ -116,11 +136,23 @@ def extract_items( existing = existing_picture_data or {} items: list[DocumentItem] = [] + serializer: Any = None + + def get_serializer() -> Any: + nonlocal serializer + if serializer is None: + from docling_core.transforms.serializer.markdown import ( + MarkdownDocSerializer, + ) + + serializer = MarkdownDocSerializer(doc=docling_doc) + return serializer + for position, (item, level) in enumerate(docling_doc.iterate_items()): label = getattr(item, "label", None) label_str = str(label.value) if hasattr(label, "value") else str(label or "") - text = extract_item_text(item, docling_doc) or "" + text = extract_item_text(item, docling_doc, get_serializer=get_serializer) or "" page_numbers: list[int] = [] if prov := getattr(item, "prov", None): diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index 37700f01..247e4223 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -140,6 +140,112 @@ class TestExtractItemText: assert items == [] +def _doc_with_captioned_picture(*captions: str): + from docling_core.types.doc.document import DoclingDocument, ImageRef + from docling_core.types.doc.labels import DocItemLabel + from PIL import Image + + doc = DoclingDocument(name="pics") + caption_items = [doc.add_text(label=DocItemLabel.CAPTION, text=c) for c in captions] + img = ImageRef.from_pil(Image.new("RGB", (8, 8), "red"), dpi=72) + doc.add_picture(image=img, caption=caption_items[0] if caption_items else None) + pic = doc.pictures[0] + for extra in caption_items[1:]: + pic.captions.append(extra.get_ref()) + return doc, pic + + +def _doc_with_tables(n: int): + from docling_core.types.doc.document import DoclingDocument, TableCell, TableData + + doc = DoclingDocument(name="tables") + for _ in range(n): + cells = [ + TableCell( + text=f"r{r}c{c}", + row_span=1, + col_span=1, + start_row_offset_idx=r, + end_row_offset_idx=r + 1, + start_col_offset_idx=c, + end_col_offset_idx=c + 1, + ) + for r in range(2) + for c in range(2) + ] + doc.add_table(data=TableData(num_rows=2, num_cols=2, table_cells=cells)) + return doc + + +class TestExtractItemTextPictures: + """Description-less pictures derive text from captions, without a serializer.""" + + def test_multiple_captions_space_joined(self): + doc, pic = _doc_with_captioned_picture("First", "Second") + assert extract_item_text(pic, doc) == "First Second" + + def test_description_wins_over_captions(self): + from docling_core.types.doc.document import DescriptionMetaField, PictureMeta + + doc, pic = _doc_with_captioned_picture("A caption") + pic.meta = PictureMeta(description=DescriptionMetaField(text="A red square.")) + assert extract_item_text(pic, doc) == "A red square." + + def test_picture_path_does_not_export_markdown(self, monkeypatch): + from docling_core.types.doc.document import PictureItem + + doc, pic = _doc_with_captioned_picture("Only caption") + + def _boom(*args, **kwargs): + raise AssertionError("picture path must not build a serializer") + + monkeypatch.setattr(PictureItem, "export_to_markdown", _boom) + assert extract_item_text(pic, doc) == "Only caption" + + def test_picture_path_never_requests_serializer(self): + doc, pic = _doc_with_captioned_picture("Only caption") + + def _explode(): + raise AssertionError("picture path must not request a serializer") + + assert extract_item_text(pic, doc, get_serializer=_explode) == "Only caption" + + +class TestExtractItemsTableSerializer: + """Table text is unchanged; one serializer is reused across the whole pass.""" + + def test_table_text_matches_export_to_markdown(self): + doc = _doc_with_tables(1) + expected = doc.tables[0].export_to_markdown(doc) + items = extract_items("doc-1", doc) + table_items = [i for i in items if i.label == "table"] + assert len(table_items) == 1 + assert table_items[0].text == expected + + def test_one_serializer_built_for_multiple_tables(self, monkeypatch): + import docling_core.transforms.serializer.markdown as md + + count = {"n": 0} + base = md.MarkdownDocSerializer + + class Counting(base): + def __init__(self, *args, **kwargs): + count["n"] += 1 + super().__init__(*args, **kwargs) + + doc = _doc_with_tables(3) + monkeypatch.setattr(md, "MarkdownDocSerializer", Counting) + + items = extract_items("doc-1", doc) + assert sum(1 for i in items if i.label == "table") == 3 + assert count["n"] == 1 + + def test_direct_table_call_builds_one_off_serializer(self): + doc = _doc_with_tables(1) + expected = doc.tables[0].export_to_markdown(doc) + assert extract_item_text(doc.tables[0], doc) == expected + + @pytest.mark.asyncio class TestDocumentItemRepository: async def test_create_and_get_range(self, temp_db_path):