From ccab0127474b27b7008efd58c094355bd2758d49 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 14:01:35 -0400 Subject: [PATCH 1/3] Thread picture chunk merging off the asyncio event loop Extracts build_picture_chunks + iterate_items position mapping into a sync helper and wraps it with asyncio.to_thread so image-heavy documents don't block the event loop during chunking. Fixes #456. --- haiku_rag_slim/haiku/rag/client/processing.py | 60 ++++++++++++------- 1 file changed, 38 insertions(+), 22 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 3300e55e..75fed8db 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -1,3 +1,4 @@ +import asyncio import logging import tempfile from pathlib import Path @@ -167,6 +168,38 @@ async def convert( return doc +def _merge_picture_chunks( + docling_document, + text_chunks: list, + document_id: str | None, + existing_picture_data: dict | None, +) -> list: + picture_chunks = build_picture_chunks( + docling_document, + document_id=document_id, + existing_picture_data=existing_picture_data, + ) + + if not picture_chunks: + for i, c in enumerate(text_chunks): + c.order = i + return text_chunks + + positions = { + item.self_ref: pos + for pos, (item, _level) in enumerate(docling_document.iterate_items()) + } + + def first_pos(c): + refs = (c.metadata or {}).get("doc_item_refs") or [] + return positions.get(refs[0], len(positions)) if refs else len(positions) + + merged = sorted(text_chunks + picture_chunks, key=first_pos) + for i, c in enumerate(merged): + c.order = i + return merged + + async def chunk( config: AppConfig, docling_document: "DoclingDocument", @@ -196,31 +229,14 @@ async def chunk( c.order = i return text_chunks - picture_chunks = build_picture_chunks( + return await asyncio.to_thread( + _merge_picture_chunks, docling_document, - document_id=document_id, - existing_picture_data=existing_picture_data, + text_chunks, + document_id, + existing_picture_data, ) - if not picture_chunks: - for i, c in enumerate(text_chunks): - c.order = i - return text_chunks - - positions = { - item.self_ref: pos - for pos, (item, _level) in enumerate(docling_document.iterate_items()) - } - - def first_pos(c: Chunk) -> int: - refs = (c.metadata or {}).get("doc_item_refs") or [] - return positions.get(refs[0], len(positions)) if refs else len(positions) - - merged = sorted(text_chunks + picture_chunks, key=first_pos) - for i, c in enumerate(merged): - c.order = i - return merged - def build_picture_chunks( docling_document: "DoclingDocument", From c59f88bcd67e1893db514daa345c0409ace9ebd4 Mon Sep 17 00:00:00 2001 From: Chris McDonough Date: Mon, 22 Jun 2026 14:13:02 -0400 Subject: [PATCH 2/3] Add test for _merge_picture_chunks no-pictures branch --- tests/test_processing.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/test_processing.py b/tests/test_processing.py index 9216b3e4..571f9581 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -207,3 +207,18 @@ async def test_convert_text_path_also_warns(monkeypatch, caplog_warnings): await convert(config, "") assert any("0 described" in r.getMessage() for r in caplog_warnings) + + +def test_merge_picture_chunks_no_pictures_returns_text_chunks(): + """When there are no picture chunks, _merge_picture_chunks returns + text chunks with order set.""" + from haiku.rag.client.processing import _merge_picture_chunks + from haiku.rag.store.models.chunk import Chunk + + doc = _doc_without_pictures() + text_chunks = [Chunk(content="a"), Chunk(content="b")] + + result = _merge_picture_chunks(doc, text_chunks, None, None) + + assert result is text_chunks + assert [c.order for c in result] == [0, 1] From 141e288b834781d9fa9c5a3d71727174b290ad3f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 23 Jun 2026 08:56:51 +0300 Subject: [PATCH 3/3] Tighten _merge_picture_chunks type annotations Restore the typing the merge logic carried inline before extraction: text_chunks: list[Chunk], existing_picture_data: dict[str, bytes] | None, -> list[Chunk], and first_pos(c: Chunk) -> int. Co-Authored-By: Claude Opus 4.8 (1M context) --- haiku_rag_slim/haiku/rag/client/processing.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 75fed8db..198f7736 100644 --- a/haiku_rag_slim/haiku/rag/client/processing.py +++ b/haiku_rag_slim/haiku/rag/client/processing.py @@ -169,11 +169,11 @@ async def convert( def _merge_picture_chunks( - docling_document, - text_chunks: list, + docling_document: "DoclingDocument", + text_chunks: list[Chunk], document_id: str | None, - existing_picture_data: dict | None, -) -> list: + existing_picture_data: dict[str, bytes] | None, +) -> list[Chunk]: picture_chunks = build_picture_chunks( docling_document, document_id=document_id, @@ -190,7 +190,7 @@ def _merge_picture_chunks( for pos, (item, _level) in enumerate(docling_document.iterate_items()) } - def first_pos(c): + def first_pos(c: Chunk) -> int: refs = (c.metadata or {}).get("doc_item_refs") or [] return positions.get(refs[0], len(positions)) if refs else len(positions)