From 875c635a772016b2f936946b072b8c16dbd14c75 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Apr 2026 15:30:24 +0300
Subject: [PATCH] have rebuild --rechunk re-chunk from the stored docling blob
 instead of re-converting from the markdown export

---
 haiku_rag_slim/haiku/rag/client/rebuild.py | 15 ++++++-------
 tests/test_picture_in_context.py           | 26 ++++++++++++++++++++++
 2 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py
index 36448377..92fc1630 100644
--- a/haiku_rag_slim/haiku/rag/client/rebuild.py
+++ b/haiku_rag_slim/haiku/rag/client/rebuild.py
@@ -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
diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py
index 749812cd..e36085f3 100644
--- a/tests/test_picture_in_context.py
+++ b/tests/test_picture_in_context.py
@@ -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