Merge pull request #416 from ggozad/fix/rebuild-embed-images

Fix rebuild --embed-only corrupting picture embeddings

.github/workflows/test.yml

@@ -62,15 +62,21 @@ jobs:
       - name: Install dependencies
         run: uv sync --all-extras
       - name: Cache HuggingFace models
+        id: hf-cache
         uses: actions/cache@v4
         with:
           path: ~/.cache/huggingface
-          key: huggingface-${{ runner.os }}-qwen-tokenizer-v1
+          key: huggingface-${{ runner.os }}-test-models-v1
       - name: Pre-download tokenizer
+        if: steps.hf-cache.outputs.cache-hit != 'true'
         run: uv run python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('Qwen/Qwen3-Embedding-0.6B')"
       - name: Pre-download cross-encoder test model
+        if: steps.hf-cache.outputs.cache-hit != 'true'
         run: uv run python -c "from sentence_transformers import CrossEncoder; CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')"
       - name: Run tests with coverage
+        env:
+          HF_HUB_OFFLINE: ${{ steps.hf-cache.outputs.cache-hit == 'true' && '1' || '0' }}
+          TRANSFORMERS_OFFLINE: ${{ steps.hf-cache.outputs.cache-hit == 'true' && '1' || '0' }}
         run: uv run pytest -m "not integration" --cov=haiku --cov-report=xml
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v5

CHANGELOG.md

@@ -5,6 +5,10 @@
 
 - `ingester.queue.retention_days` (default 30): the reaper deletes succeeded/dead jobs whose `completed_at` is older than the window. `null` disables pruning.
 
+### Fixed
+
+- `rebuild --embed-only` re-embeds picture chunks through the image path instead of overwriting their vectors with a text embedding of the caption.
+
 ## [0.52.0] - 2026-06-01
 
 ### Added

haiku_rag_slim/haiku/rag/client/rebuild.py

@@ -357,10 +357,10 @@ async def _rebuild_embed_only(
     treats it as a partial phase 1, which is harmless because phase 2 has
     already finished writing the new chunks table.
     """
-    from haiku.rag.embeddings import contextualize
+    from haiku.rag.embeddings import contextualize, embed_chunks
 
     db = client.store.db
-    batch_size = client._config.embeddings.batch_size
+    embedder = client.chunk_repository.embedder
 
     if not resume_from_staging:
         # Phase 1: copy chunks into staging, then mark it complete. After the
@@ -384,17 +384,37 @@ async def _rebuild_embed_only(
         if not chunks:
             continue
 
-        texts = contextualize(chunks)
+        # Re-attach picture bytes stripped by the staging copy so picture
-        embeddings: list[list[float]] = []
+        # chunks route through embed_image rather than being text-embedded.
-        for i in range(0, len(texts), batch_size):
+        # Bytes live in document_items (embed-only never touches that table).
-            batch_embeddings = await client.chunk_repository.embedder.embed_documents(
+        if embedder.supports_images:
-                texts[i : i + batch_size]
+            picture_data = await client.document_item_repository.get_all_picture_data(
+                doc.id
             )
-            embeddings.extend(batch_embeddings)
+            for chunk in chunks:
+                if "picture" not in (chunk.metadata.get("labels") or []):
+                    continue
+                refs = chunk.metadata.get("doc_item_refs") or []
+                data = next((picture_data[r] for r in refs if r in picture_data), None)
+                if data is not None:
+                    chunk._picture_data = data
+                else:
+                    logger.warning(
+                        "Document %s picture chunk %s has no recoverable bytes; "
+                        "embedding its caption as text.",
+                        doc.id,
+                        chunk.id,
+                    )
 
-        for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
+        content_fts_list = contextualize(chunks)
+        embedded_chunks = await embed_chunks(chunks, embedder, client._config)
+
+        for chunk, content_fts, embedded in zip(
+            chunks, content_fts_list, embedded_chunks
+        ):
             assert chunk.id is not None
             assert chunk.document_id is not None
+            assert embedded.embedding is not None
             pending_records.append(
                 client.store.ChunkRecord(
                     id=chunk.id,
@@ -403,7 +423,7 @@ async def _rebuild_embed_only(
                     content_fts=content_fts,
                     metadata=json.dumps(chunk.metadata),
                     order=chunk.order,
-                    vector=embedding,
+                    vector=embedded.embedding,
                 )
             )
 

tests/test_picture_in_context.py

@@ -202,6 +202,70 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
         assert after.get("#/pictures/0") == before.get("#/pictures/0")
 
 
+@pytest.mark.asyncio
+async def test_embed_only_preserves_picture_vectors(temp_db_path, monkeypatch):
+    """``rebuild --embed-only`` must re-embed picture chunks through the image
+    path. With a multimodal embedder, picture vectors must survive the rebuild
+    instead of being overwritten by a text embedding of the caption."""
+    from haiku.rag.client import RebuildMode
+    from haiku.rag.client.documents import _store_document_with_chunks
+    from haiku.rag.config import EmbeddingModelConfig, EmbeddingsConfig
+    from haiku.rag.embeddings import EmbedderWrapper, embed_chunks
+    from haiku.rag.store.models.document import Document
+    from tests.store.test_document_items import _docling_doc_with_picture
+
+    TEXT_VEC = [0.1, 0.1, 0.1, 0.1]
+    IMAGE_VEC = [0.9, 0.9, 0.9, 0.9]
+
+    class StubMultimodalEmbedder(EmbedderWrapper):
+        supports_images = True
+
+        def __init__(self):
+            super().__init__(embedder=None, vector_dim=4)
+
+        async def embed_documents(self, texts):
+            return [list(TEXT_VEC) for _ in texts]
+
+        async def embed_image(self, image):
+            return list(IMAGE_VEC)
+
+    monkeypatch.setattr(
+        "haiku.rag.store.engine.get_embedder",
+        lambda *a, **kw: StubMultimodalEmbedder(),
+    )
+
+    config = AppConfig(
+        embeddings=EmbeddingsConfig(
+            model=EmbeddingModelConfig(provider="ollama", name="stub", vector_dim=4)
+        )
+    )
+
+    docling_doc = _docling_doc_with_picture()
+
+    async def _picture_chunk_row(rag):
+        rows = await rag.chunk_repository.store.chunks_table.query().to_list()
+        picture_rows = [r for r in rows if "#/pictures/0" in (r.get("metadata") or "")]
+        assert len(picture_rows) == 1
+        return picture_rows[0]
+
+    async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
+        chunks = await rag.chunk(docling_doc)
+        embedded = await embed_chunks(chunks, rag.embedder, rag._config)
+        document = Document(content="x", uri="test://doc")
+        document.set_docling(docling_doc)
+        await _store_document_with_chunks(rag, document, embedded, docling_doc)
+
+        before = await _picture_chunk_row(rag)
+        assert list(before["vector"]) == pytest.approx(IMAGE_VEC)
+
+        async for _ in rag.rebuild_database(mode=RebuildMode.EMBED_ONLY):
+            pass
+
+        after = await _picture_chunk_row(rag)
+        assert list(after["vector"]) == pytest.approx(IMAGE_VEC)
+        assert after["id"] == before["id"]
+
+
 @pytest.mark.asyncio
 async def test_expand_context_does_not_attach_expansion_added_pictures(temp_db_path):
     """expand_context preserves picture bytes from the pre-expansion result and