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.
This commit is contained in:
Chris McDonough 2026-06-22 14:01:35 -04:00
parent 0fcf91fcf7
commit ccab012747

View file

@ -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",