Fix rebuild --embed-only corrupting picture embeddings

This commit is contained in:
Yiorgis Gozadinos 2026-06-03 11:46:43 +03:00
parent 6d95fbe74a
commit 651b22ddcf
No known key found for this signature in database
3 changed files with 98 additions and 10 deletions

View file

@ -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

View file

@ -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,
)
)

View file

@ -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