From 22b43971483ef79f9efcdcde076538592df4911e Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 13:42:06 -0400 Subject: [PATCH 1/6] Thread compress_docling_split in rebuild.py off the asyncio event loop Extracts pydantic serialization and zstd compression into a sync helper and wraps it with asyncio.to_thread so picture description rebuilds don't block the event loop. Fixes #457. --- haiku_rag_slim/haiku/rag/client/rebuild.py | 56 ++++++++++++---------- 1 file changed, 30 insertions(+), 26 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 6ea1986c..9b6a2c1a 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -1,3 +1,4 @@ +import asyncio import json import logging from collections.abc import AsyncGenerator @@ -618,6 +619,33 @@ async def _rebuild_rechunk( await _flush_rebuild_batch(client, pending_docs, pending_chunks) +def _apply_descriptions_sync(docling_doc, doc, descriptions): + """Patch picture descriptions into the docling document and re-compress. + + Updates only docling_document — set_docling would also overwrite + docling_pages by routing through compress_docling_split, which + extracts pages from the in-memory JSON and finds none (the pages + blob is stored separately and is not loaded by get_docling_document). + That would silently destroy page rasters for every doc with at + least one undescribed picture. + """ + from docling_core.types.doc.document import DescriptionMetaField, PictureMeta + from haiku.rag.store.compression import compress_docling_split + + for pic in docling_doc.pictures: + text = descriptions.get(pic.self_ref) + if not text: + continue + if pic.meta is None: + pic.meta = PictureMeta() + pic.meta.description = DescriptionMetaField(text=text) + + structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json()) + doc.docling_document = structure_bytes + doc.docling_version = docling_doc.version + return len(descriptions) + + async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int: """Run the VLM against pictures lacking a description, patch the docling blob in-place. Returns the number of newly described pictures. @@ -660,34 +688,10 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int: if not descriptions: return 0 - # Patch the docling document in-place. PictureMeta + DescriptionMetaField - # are pydantic models; build them and assign. - from docling_core.types.doc.document import ( - DescriptionMetaField, - PictureMeta, + return await asyncio.to_thread( + _apply_descriptions_sync, docling_doc, doc, descriptions ) - for pic in docling_doc.pictures: - text = descriptions.get(pic.self_ref) - if not text: - continue - if pic.meta is None: - pic.meta = PictureMeta() - pic.meta.description = DescriptionMetaField(text=text) - - # Update only docling_document — set_docling would also overwrite - # docling_pages by routing through compress_docling_split, which - # extracts pages from the in-memory JSON and finds none (the pages - # blob is stored separately and is not loaded by get_docling_document). - # That would silently destroy page rasters for every doc with at - # least one undescribed picture. - from haiku.rag.store.compression import compress_docling_split - - structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json()) - doc.docling_document = structure_bytes - doc.docling_version = docling_doc.version - return len(descriptions) - async def _rebuild_descriptions( client: "HaikuRAG", documents: list[Document] From 97748c2fee7b4d4fa71675798a8d5a925247df36 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 13:46:21 -0400 Subject: [PATCH 2/6] Fix import sorting lint error --- haiku_rag_slim/haiku/rag/client/rebuild.py | 1 + 1 file changed, 1 insertion(+) diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 9b6a2c1a..d672e1a0 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -630,6 +630,7 @@ def _apply_descriptions_sync(docling_doc, doc, descriptions): least one undescribed picture. """ from docling_core.types.doc.document import DescriptionMetaField, PictureMeta + from haiku.rag.store.compression import compress_docling_split for pic in docling_doc.pictures: From d160080ff8bd789627995972984dbd3dbeda0344 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 13:48:09 -0400 Subject: [PATCH 3/6] Move lazy imports to module scope --- haiku_rag_slim/haiku/rag/client/rebuild.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index d672e1a0..f06c6dad 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -5,6 +5,7 @@ from collections.abc import AsyncGenerator from datetime import datetime from typing import TYPE_CHECKING +from docling_core.types.doc.document import DescriptionMetaField, PictureMeta from lancedb.pydantic import LanceModel from haiku.rag.client.documents import check_source_accessible @@ -12,6 +13,7 @@ from haiku.rag.converters import get_converter from haiku.rag.store.engine import ChunkRecordBase from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document +from haiku.rag.store.compression import compress_docling_split from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.repositories.settings import SettingsRepository @@ -629,10 +631,6 @@ def _apply_descriptions_sync(docling_doc, doc, descriptions): That would silently destroy page rasters for every doc with at least one undescribed picture. """ - from docling_core.types.doc.document import DescriptionMetaField, PictureMeta - - from haiku.rag.store.compression import compress_docling_split - for pic in docling_doc.pictures: text = descriptions.get(pic.self_ref) if not text: From cee9484f35604f4ef908f4199c28fa561031a5c3 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 13:49:59 -0400 Subject: [PATCH 4/6] Fix import sort order --- haiku_rag_slim/haiku/rag/client/rebuild.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index f06c6dad..ff396884 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -10,10 +10,10 @@ from lancedb.pydantic import LanceModel from haiku.rag.client.documents import check_source_accessible from haiku.rag.converters import get_converter +from haiku.rag.store.compression import compress_docling_split from haiku.rag.store.engine import ChunkRecordBase from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document -from haiku.rag.store.compression import compress_docling_split from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.repositories.settings import SettingsRepository From a137b91654abc753aefbafe434e99feafe59eb66 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 14:08:23 -0400 Subject: [PATCH 5/6] Add test for _apply_descriptions_sync to cover skip branch --- tests/test_rebuild.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 9649a625..06e34877 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -970,6 +970,34 @@ async def test_rebuild_descriptions_raises_when_blob_is_missing( pass +def test_apply_descriptions_sync_skips_pictures_without_descriptions(): + """_apply_descriptions_sync skips pictures not in the descriptions dict.""" + from unittest.mock import MagicMock + + from haiku.rag.client.rebuild import _apply_descriptions_sync + + pic_with = MagicMock() + pic_with.self_ref = "pic-1" + pic_with.meta = None + + pic_without = MagicMock() + pic_without.self_ref = "pic-2" + + docling_doc = MagicMock() + docling_doc.pictures = [pic_with, pic_without] + docling_doc.model_dump_json.return_value = "{}" + docling_doc.version = "1.0" + + doc = MagicMock() + + descriptions = {"pic-1": "A red square."} + n = _apply_descriptions_sync(docling_doc, doc, descriptions) + + assert n == 1 + assert pic_with.meta is not None + assert pic_with.meta.description.text == "A red square." + + async def _add_chunk(client: HaikuRAG, vector: list[float]) -> str: """Insert a chunk row directly, bypassing the embedder.""" record = client.store.ChunkRecord( From 775dc58b53617e68a43f3e934fd56e77c5186b95 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 23 Jun 2026 09:14:36 +0300 Subject: [PATCH 6/6] Drop redundant _apply_descriptions_sync unit test, annotate the helper --- haiku_rag_slim/haiku/rag/client/rebuild.py | 6 ++++- tests/test_rebuild.py | 28 ---------------------- 2 files changed, 5 insertions(+), 29 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index ff396884..f7702f45 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -18,6 +18,8 @@ from haiku.rag.store.models.document_item import extract_items from haiku.rag.store.repositories.settings import SettingsRepository if TYPE_CHECKING: + from docling_core.types.doc.document import DoclingDocument + from haiku.rag.client import HaikuRAG, RebuildMode logger = logging.getLogger(__name__) @@ -621,7 +623,9 @@ async def _rebuild_rechunk( await _flush_rebuild_batch(client, pending_docs, pending_chunks) -def _apply_descriptions_sync(docling_doc, doc, descriptions): +def _apply_descriptions_sync( + docling_doc: "DoclingDocument", doc: Document, descriptions: dict[str, str] +) -> int: """Patch picture descriptions into the docling document and re-compress. Updates only docling_document — set_docling would also overwrite diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index 06e34877..9649a625 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -970,34 +970,6 @@ async def test_rebuild_descriptions_raises_when_blob_is_missing( pass -def test_apply_descriptions_sync_skips_pictures_without_descriptions(): - """_apply_descriptions_sync skips pictures not in the descriptions dict.""" - from unittest.mock import MagicMock - - from haiku.rag.client.rebuild import _apply_descriptions_sync - - pic_with = MagicMock() - pic_with.self_ref = "pic-1" - pic_with.meta = None - - pic_without = MagicMock() - pic_without.self_ref = "pic-2" - - docling_doc = MagicMock() - docling_doc.pictures = [pic_with, pic_without] - docling_doc.model_dump_json.return_value = "{}" - docling_doc.version = "1.0" - - doc = MagicMock() - - descriptions = {"pic-1": "A red square."} - n = _apply_descriptions_sync(docling_doc, doc, descriptions) - - assert n == 1 - assert pic_with.meta is not None - assert pic_with.meta.description.text == "A red square." - - async def _add_chunk(client: HaikuRAG, vector: list[float]) -> str: """Insert a chunk row directly, bypassing the embedder.""" record = client.store.ChunkRecord(