diff --git a/haiku_rag_slim/haiku/rag/client/processing.py b/haiku_rag_slim/haiku/rag/client/processing.py index 3300e55e..198f7736 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: "DoclingDocument", + text_chunks: list[Chunk], + document_id: str | None, + existing_picture_data: dict[str, bytes] | None, +) -> list[Chunk]: + 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: 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 + + 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", 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]