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

Thread extract_items off the asyncio event loop

Browse source 
Moves CPU-bound docling document item extraction into worker threads
via asyncio.to_thread so the event loop stays responsive during
ingestion. The extract_items calls are hoisted out of the write lock
(they are pure computation with no DB I/O) and run with placeholder
document IDs that are patched after the DB create returns.

Fixes #453.
This commit is contained in:
Chris McDonough 2026-06-22 13:19:38 -04:00
parent 0fcf91fcf7
commit 09e78ef46a
1 changed files with 20 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
haiku_rag_slim/haiku/rag/client/documents.py
View file

@ -118,6 +118,7 @@ async def _store_document_with_chunks(
Handles versioning/rollback on failure. Handles versioning/rollback on failure.
""" """
chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder) chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
items = await asyncio.to_thread(extract_items, "", docling_document)
async with client.store._write_lock: async with client.store._write_lock:
versions = await client.store.current_table_versions() versions = await client.store.current_table_versions()
@ -134,7 +135,8 @@ async def _store_document_with_chunks(
await client.chunk_repository.create(chunks) await client.chunk_repository.create(chunks)
items = extract_items(created_doc.id, docling_document) for item in items:
item.document_id = created_doc.id
await client.document_item_repository.create_items(created_doc.id, items) await client.document_item_repository.create_items(created_doc.id, items)
if client._config.storage.auto_vacuum: if client._config.storage.auto_vacuum:
@ -170,6 +172,12 @@ async def _update_document_with_chunks(
chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder) chunks = await ensure_chunks_embedded(client._config, chunks, client.embedder)
items: list | None = None
if docling_document is not None:
items = await asyncio.to_thread(
extract_items, document.id, docling_document, existing_picture_data
)
async with client.store._write_lock: async with client.store._write_lock:
versions = await client.store.current_table_versions() versions = await client.store.current_table_versions()
@ -183,12 +191,7 @@ async def _update_document_with_chunks(
await client.chunk_repository.replace_for_document(updated_doc.id, chunks) await client.chunk_repository.replace_for_document(updated_doc.id, chunks)
if docling_document is not None: if items is not None:
items = extract_items(
updated_doc.id,
docling_document,
existing_picture_data=existing_picture_data,
)
await client.document_item_repository.replace_for_document( await client.document_item_repository.replace_for_document(
updated_doc.id, items updated_doc.id, items
) )
@ -279,6 +282,11 @@ async def _store_documents_with_chunks(
for _, chunks, _ in prepared for _, chunks, _ in prepared
] ]
def _extract_all_items():
return [extract_items("", d) for _, _, d in prepared]
all_item_lists = await asyncio.to_thread(_extract_all_items)
async with client.store._write_lock: async with client.store._write_lock:
versions = await client.store.current_table_versions() versions = await client.store.current_table_versions()
@ -289,15 +297,17 @@ async def _store_documents_with_chunks(
try: try:
all_chunks: list[Chunk] = [] all_chunks: list[Chunk] = []
all_items = [] all_items = []
for doc, doc_chunks, docling_document in zip( for doc, doc_chunks, item_list in zip(
created, embedded, (d for _, _, d in prepared) created, embedded, all_item_lists
): ):
assert doc.id is not None assert doc.id is not None
for order, chunk in enumerate(doc_chunks): for order, chunk in enumerate(doc_chunks):
chunk.document_id = doc.id chunk.document_id = doc.id
chunk.order = order chunk.order = order
all_chunks.extend(doc_chunks) all_chunks.extend(doc_chunks)
all_items.extend(extract_items(doc.id, docling_document)) for item in item_list:
item.document_id = doc.id
all_items.extend(item_list)
await client.chunk_repository.create(all_chunks) await client.chunk_repository.create(all_chunks)
await client.document_item_repository.create_all(all_items) await client.document_item_repository.create_all(all_items)