From 651b22ddcf2131c4bbd785b7d8ae98307b72086b Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 3 Jun 2026 11:46:43 +0300 Subject: [PATCH 1/2] Fix rebuild --embed-only corrupting picture embeddings --- CHANGELOG.md | 4 ++ haiku_rag_slim/haiku/rag/client/rebuild.py | 40 ++++++++++---- tests/test_picture_in_context.py | 64 ++++++++++++++++++++++ 3 files changed, 98 insertions(+), 10 deletions(-) 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 From a65757807e311ae9fdd5dd8d5e2293b2bc82b09a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 3 Jun 2026 12:01:43 +0300 Subject: [PATCH 2/2] Use cached HuggingFace models offline in test job The Qwen tokenizer and cross-encoder pre-downloads call the HF metadata API to revalidate even on a cache hit; a 429 there propagates instead of falling back to the cached files, failing CI on HF throttling. Skip the pre-download steps when the cache is restored and run pytest with HF_HUB_OFFLINE/TRANSFORMERS_OFFLINE on a hit, so cached models are used without any network revalidation; allow online on a miss so a fresh cache key still populates. Rename the cache key so the snapshot re-populates with every test model (the old key predated the cross-encoder step and never cached it). Co-Authored-By: Claude Opus 4.8 (1M context) --- .github/workflows/test.yml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 31d7dd7e..5422bb65 100644 --- a/.github/workflows/test.yml +++ b/.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