danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #460 from mcdonc/thread-processing-iterate

Browse source 
Thread picture chunk merging off the asyncio event loop
This commit is contained in:
Yiorgis Gozadinos 2026-06-23 09:05:45 +03:00 committed by GitHub
commit 0ee36a269d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 53 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

60
haiku_rag_slim/haiku/rag/client/processing.py
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: "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",

15
tests/test_processing.py
View file

@ -207,3 +207,18 @@ async def test_convert_text_path_also_warns(monkeypatch, caplog_warnings):
await convert(config, "<html><img src='...'/></html>")
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]