diff --git a/CHANGELOG.md b/CHANGELOG.md index b9aa8f36..ba10337a 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 6144dae4..f44b6b38 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/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) - embeddings: list[list[float]] = [] - for i in range(0, len(texts), batch_size): - batch_embeddings = await client.chunk_repository.embedder.embed_documents( - texts[i : i + batch_size] + # Re-attach picture bytes stripped by the staging copy so picture + # chunks route through embed_image rather than being text-embedded. + # Bytes live in document_items (embed-only never touches that table). + if embedder.supports_images: + 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, ) ) diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 431e75e8..173e1c7a 100644 --- a/tests/test_picture_in_context.py +++ b/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