Fix rebuild --embed-only corrupting picture embeddings
This commit is contained in:
parent
6d95fbe74a
commit
651b22ddcf
3 changed files with 98 additions and 10 deletions
|
|
@ -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.
|
- `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
|
## [0.52.0] - 2026-06-01
|
||||||
|
|
||||||
### Added
|
### Added
|
||||||
|
|
|
||||||
|
|
@ -357,10 +357,10 @@ async def _rebuild_embed_only(
|
||||||
treats it as a partial phase 1, which is harmless because phase 2 has
|
treats it as a partial phase 1, which is harmless because phase 2 has
|
||||||
already finished writing the new chunks table.
|
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
|
db = client.store.db
|
||||||
batch_size = client._config.embeddings.batch_size
|
embedder = client.chunk_repository.embedder
|
||||||
|
|
||||||
if not resume_from_staging:
|
if not resume_from_staging:
|
||||||
# Phase 1: copy chunks into staging, then mark it complete. After the
|
# Phase 1: copy chunks into staging, then mark it complete. After the
|
||||||
|
|
@ -384,17 +384,37 @@ async def _rebuild_embed_only(
|
||||||
if not chunks:
|
if not chunks:
|
||||||
continue
|
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.id is not None
|
||||||
assert chunk.document_id is not None
|
assert chunk.document_id is not None
|
||||||
|
assert embedded.embedding is not None
|
||||||
pending_records.append(
|
pending_records.append(
|
||||||
client.store.ChunkRecord(
|
client.store.ChunkRecord(
|
||||||
id=chunk.id,
|
id=chunk.id,
|
||||||
|
|
@ -403,7 +423,7 @@ async def _rebuild_embed_only(
|
||||||
content_fts=content_fts,
|
content_fts=content_fts,
|
||||||
metadata=json.dumps(chunk.metadata),
|
metadata=json.dumps(chunk.metadata),
|
||||||
order=chunk.order,
|
order=chunk.order,
|
||||||
vector=embedding,
|
vector=embedded.embedding,
|
||||||
)
|
)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -202,6 +202,70 @@ async def test_rechunk_preserves_picture_data(temp_db_path):
|
||||||
assert after.get("#/pictures/0") == before.get("#/pictures/0")
|
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
|
@pytest.mark.asyncio
|
||||||
async def test_expand_context_does_not_attach_expansion_added_pictures(temp_db_path):
|
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
|
"""expand_context preserves picture bytes from the pre-expansion result and
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue