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