have rebuild --rechunk re-chunk from the stored docling blob instead of re-converting from the markdown export

This commit is contained in:
Yiorgis Gozadinos 2026-04-30 15:30:24 +03:00
parent 7831057339
commit 875c635a77
No known key found for this signature in database
2 changed files with 33 additions and 8 deletions

View file

@ -223,28 +223,27 @@ async def _flush_rebuild_batch(
async def _rebuild_rechunk(
client: "HaikuRAG", documents: list[Document]
) -> AsyncGenerator[str, None]:
"""Re-chunk and re-embed from existing document content."""
"""Re-chunk and re-embed each document from its stored docling blob."""
from haiku.rag.embeddings import embed_chunks
pending_chunks: list[Chunk] = []
pending_docs: list[Document] = []
pending_doc_ids: list[str] = []
converter = get_converter(client._config)
for doc in documents:
assert doc.id is not None
# Convert stored markdown to DoclingDocument
docling_document = await converter.convert_text(doc.content, format="md")
docling_document = doc.get_docling_document()
if docling_document is None:
raise ValueError(
f"Document {doc.id} has no stored docling document; rechunk "
"requires it. Run a full rebuild (without --rechunk) instead."
)
# Chunk and embed
chunks = await client.chunk(docling_document)
embedded_chunks = await embed_chunks(chunks, client._config)
# Update document fields
doc.set_docling(docling_document)
# Prepare chunks with document_id and order
for order, chunk in enumerate(embedded_chunks):
chunk.document_id = doc.id

View file

@ -166,6 +166,32 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat
assert "picture" in out.labels
@pytest.mark.asyncio
async def test_rechunk_preserves_picture_data(temp_db_path):
"""``rebuild --rechunk`` keeps ``picture_data`` for every picture row."""
from haiku.rag.client import RebuildMode
from haiku.rag.client.documents import _store_document_with_chunks
from haiku.rag.store.models.document import Document
from tests.store.test_document_items import _docling_doc_with_picture
docling_doc = _docling_doc_with_picture()
async with HaikuRAG(temp_db_path, create=True) as rag:
rag._config.processing.pictures = "image"
document = Document(content="x", uri="test://doc")
document.set_docling(docling_doc)
created = await _store_document_with_chunks(rag, document, [], docling_doc)
assert created.id is not None
before = await rag.document_item_repository.get_all_picture_data(created.id)
assert before.get("#/pictures/0") is not None
async for _ in rag.rebuild_database(mode=RebuildMode.RECHUNK):
pass
after = await rag.document_item_repository.get_all_picture_data(created.id)
assert after.get("#/pictures/0") == before.get("#/pictures/0")
@pytest.mark.asyncio
async def test_update_clears_picture_data_when_mode_none(temp_db_path):
"""Switching to ``pictures="none"`` and re-running update_document