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

Thread compress_docling_split in rebuild.py off the asyncio event loop

Browse source 
Extracts pydantic serialization and zstd compression into a sync
helper and wraps it with asyncio.to_thread so picture description
rebuilds don't block the event loop.

Fixes #457.
This commit is contained in:
Chris McDonough 2026-06-22 13:42:06 -04:00
parent 0fcf91fcf7
commit 22b4397148
1 changed files with 30 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

56
haiku_rag_slim/haiku/rag/client/rebuild.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
import logging
from collections.abc import AsyncGenerator
@ -618,6 +619,33 @@ async def _rebuild_rechunk(
await _flush_rebuild_batch(client, pending_docs, pending_chunks)
def _apply_descriptions_sync(docling_doc, doc, descriptions):
"""Patch picture descriptions into the docling document and re-compress.
Updates only docling_document — set_docling would also overwrite
docling_pages by routing through compress_docling_split, which
extracts pages from the in-memory JSON and finds none (the pages
blob is stored separately and is not loaded by get_docling_document).
That would silently destroy page rasters for every doc with at
least one undescribed picture.
"""
from docling_core.types.doc.document import DescriptionMetaField, PictureMeta
from haiku.rag.store.compression import compress_docling_split
for pic in docling_doc.pictures:
text = descriptions.get(pic.self_ref)
if not text:
continue
if pic.meta is None:
pic.meta = PictureMeta()
pic.meta.description = DescriptionMetaField(text=text)
structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json())
doc.docling_document = structure_bytes
doc.docling_version = docling_doc.version
return len(descriptions)
async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
"""Run the VLM against pictures lacking a description, patch the docling
blob in-place. Returns the number of newly described pictures.
@ -660,34 +688,10 @@ async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int:
if not descriptions:
return 0
# Patch the docling document in-place. PictureMeta + DescriptionMetaField
# are pydantic models; build them and assign.
from docling_core.types.doc.document import (
DescriptionMetaField,
PictureMeta,
return await asyncio.to_thread(
_apply_descriptions_sync, docling_doc, doc, descriptions
)
for pic in docling_doc.pictures:
text = descriptions.get(pic.self_ref)
if not text:
continue
if pic.meta is None:
pic.meta = PictureMeta()
pic.meta.description = DescriptionMetaField(text=text)
# Update only docling_document — set_docling would also overwrite
# docling_pages by routing through compress_docling_split, which
# extracts pages from the in-memory JSON and finds none (the pages
# blob is stored separately and is not loaded by get_docling_document).
# That would silently destroy page rasters for every doc with at
# least one undescribed picture.
from haiku.rag.store.compression import compress_docling_split
structure_bytes, _ = compress_docling_split(docling_doc.model_dump_json())
doc.docling_document = structure_bytes
doc.docling_version = docling_doc.version
return len(descriptions)
async def _rebuild_descriptions(
client: "HaikuRAG", documents: list[Document]