Merge pull request #235 from ggozad/fix/rebuild-changed-vector-size

Fix embed-only rebuild with changed vector dimensions

CHANGELOG.md

@@ -7,6 +7,12 @@
   - When set, tools emit `{state_key: snapshot}` instead of bare state, enabling state merging when multiple agents share state
   - Default `None` preserves backwards compatibility (bare state emission)
 - **Page Image Generation Control**: New `generate_page_images` option in `ConversionOptions` to control PDF page image extraction
+
+### Fixed
+
+- **Embed-only rebuild with changed vector dimensions**: Fixed `haiku-rag rebuild --embed-only` failing when the configured embedding model has different dimensions than the database
+  - Store now reads stored vector dimension when opening existing databases, allowing chunks to be read regardless of current config
+  - `_rebuild_embed_only` recreates the chunks table to handle dimension changes
   - `generate_page_images: bool = True` - Enable/disable rendered page images (used by `visualize_chunk()`)
   - Works with both `docling-local` and `docling-serve` converters
   - For `docling-serve`, maps to `image_export_mode` API parameter (`embedded`/`placeholder`)

haiku_rag_slim/haiku/rag/client.py

@@ -1415,57 +1415,53 @@ class HaikuRAG:
         """Re-embed all chunks without changing chunk boundaries."""
         from haiku.rag.embeddings import contextualize
 
-        batch_size = 50
-        pending_records: list = []
-        pending_doc_ids: list[str] = []
+        # Collect all chunks with new embeddings
+        all_chunk_data: list[tuple[str, dict]] = []
 
         for doc in documents:
             assert doc.id is not None
-
-            # Get existing chunks
             chunks = await self.chunk_repository.get_by_document_id(doc.id)
             if not chunks:
-                yield doc.id
                 continue
 
-            # Generate new embeddings using contextualize for consistency
             texts = contextualize(chunks)
             embeddings = await self.chunk_repository.embedder.embed_documents(texts)
 
-            # Build updated records
             for chunk, content_fts, embedding in zip(chunks, texts, embeddings):
-                pending_records.append(
-                    self.store.ChunkRecord(
-                        id=chunk.id,  # type: ignore[arg-type]
-                        document_id=chunk.document_id,  # type: ignore[arg-type]
-                        content=chunk.content,
-                        content_fts=content_fts,
-                        metadata=json.dumps(chunk.metadata),
-                        order=chunk.order,
-                        vector=embedding,
+                all_chunk_data.append(
+                    (
+                        doc.id,
+                        {
+                            "id": chunk.id,
+                            "document_id": chunk.document_id,
+                            "content": chunk.content,
+                            "content_fts": content_fts,
+                            "metadata": json.dumps(chunk.metadata),
+                            "order": chunk.order,
+                            "vector": embedding,
+                        },
                     )
                 )
 
-            pending_doc_ids.append(doc.id)
+        # Recreate chunks table (handles dimension changes)
+        self.store.recreate_embeddings_table()
 
-            # Flush batch when size reached
-            if len(pending_doc_ids) >= batch_size:
-                if pending_records:
-                    self.store.chunks_table.merge_insert(
-                        "id"
-                    ).when_matched_update_all().execute(pending_records)
-                for doc_id in pending_doc_ids:
-                    yield doc_id
-                pending_records = []
-                pending_doc_ids = []
+        # Insert all chunks
+        if all_chunk_data:
+            records = [self.store.ChunkRecord(**data) for _, data in all_chunk_data]
+            self.store.chunks_table.add(records)
 
-        # Flush remaining
-        if pending_records:
-            self.store.chunks_table.merge_insert(
-                "id"
-            ).when_matched_update_all().execute(pending_records)
-        for doc_id in pending_doc_ids:
-            yield doc_id
+        # Yield all processed doc IDs
+        yielded_docs: set[str] = set()
+        for doc_id, _ in all_chunk_data:
+            if doc_id not in yielded_docs:
+                yielded_docs.add(doc_id)
+                yield doc_id
+
+        # Yield docs with no chunks
+        for doc in documents:
+            if doc.id and doc.id not in yielded_docs:
+                yield doc.id
 
     async def _flush_rebuild_batch(
         self, documents: list[Document], chunks: list[Chunk]

haiku_rag_slim/haiku/rag/store/engine.py

@@ -90,12 +90,8 @@ class Store:
         self._before = before
         # Time-travel mode is always read-only
         self._read_only = read_only or (before is not None)
-        self.embedder = get_embedder(config=self._config)
         self._vacuum_lock = asyncio.Lock()
 
-        # Create the ChunkRecord model with the correct vector dimension
-        self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
-
         # Check if database exists (for local filesystem only)
         is_new_db = False
         if not self._has_cloud_config():
@@ -113,6 +109,19 @@ class Store:
         # Connect to LanceDB
         self.db = self._connect_to_lancedb(db_path)
 
+        # For existing databases, read stored vector dimension to create ChunkRecord
+        # that can read existing chunks. For new databases, use config's dimension.
+        stored_vector_dim = None
+        if not is_new_db:
+            stored_vector_dim = self._get_stored_vector_dim()
+
+        # Create embedder with config's dimension (for generating new embeddings)
+        self.embedder = get_embedder(config=self._config)
+
+        # Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB)
+        chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
+        self.ChunkRecord = create_chunk_model(chunk_vector_dim)
+
         # Initialize tables (creates them if they don't exist)
         self._init_tables()
 
@@ -137,6 +146,35 @@ class Store:
         """Whether the store is in read-only mode."""
         return self._read_only
 
+    def _get_stored_vector_dim(self) -> int | None:
+        """Read the stored vector dimension from the settings table.
+
+        Returns:
+            The stored vector dimension, or None if not found.
+        """
+        try:
+            existing_tables = self.db.table_names()
+            if "settings" not in existing_tables:
+                return None
+
+            settings_table = self.db.open_table("settings")
+            rows = (
+                settings_table.search()
+                .where("id = 'settings'")
+                .limit(1)
+                .to_arrow()
+                .to_pylist()
+            )
+            if not rows or not rows[0].get("settings"):
+                return None
+
+            settings = json.loads(rows[0]["settings"])
+            embeddings = settings.get("embeddings", {})
+            model = embeddings.get("model", {})
+            return model.get("vector_dim")
+        except Exception:
+            return None
+
     def _assert_writable(self) -> None:
         """Raise ReadOnlyError if the store is in read-only mode."""
         if self._read_only:

tests/test_rebuild.py

@@ -106,6 +106,116 @@ async def test_rebuild_embed_only_skips_unchanged(qa_corpus: Dataset, temp_db_pa
             assert embeddings_before[chunk_id] == embeddings_after[chunk_id]
 
 
+@pytest.mark.vcr()
+async def test_rebuild_embed_only_with_changed_vector_dim(
+    qa_corpus: Dataset, temp_db_path
+):
+    """Test embed-only rebuild when vector dimension changes.
+
+    This tests the scenario where a database was created with one embedding model
+    (e.g., qwen3-embedding:8b with 4096 dims) and rebuild is run with a different
+    model (e.g., qwen3-embedding:4b with 2560 dims).
+
+    The Store should use the stored vector_dim for reading existing chunks,
+    then rebuild should handle changing to the new dimension.
+    """
+    import json
+
+    import lancedb
+    from lancedb.pydantic import LanceModel, Vector
+    from pydantic import Field
+
+    # Step 1: Create a database with normal 2560-dim embeddings
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        doc = await client.create_document(content=qa_corpus["document_extracted"][0])
+        assert doc.id is not None
+
+        chunks_before = await client.chunk_repository.get_by_document_id(doc.id)
+        assert len(chunks_before) > 0
+        chunk_data = [
+            {
+                "id": c.id,
+                "document_id": c.document_id,
+                "content": c.content,
+                "content_fts": c.content,
+                "metadata": json.dumps(c.metadata),
+                "order": c.order,
+            }
+            for c in chunks_before
+        ]
+
+    # Step 2: Manually recreate chunks table with 4096-dim vectors (simulating old DB)
+    db = lancedb.connect(temp_db_path)
+
+    class ChunkRecord4096(LanceModel):
+        id: str
+        document_id: str
+        content: str
+        content_fts: str = Field(default="")
+        metadata: str = Field(default="{}")
+        order: int = Field(default=0)
+        vector: Vector(4096) = Field(default_factory=lambda: [0.0] * 4096)  # type: ignore
+
+    db.drop_table("chunks")
+    chunks_table = db.create_table("chunks", schema=ChunkRecord4096)
+
+    # Insert chunks with 4096-dim fake vectors
+    records_4096 = [
+        ChunkRecord4096(
+            id=c["id"],
+            document_id=c["document_id"],
+            content=c["content"],
+            content_fts=c["content_fts"],
+            metadata=c["metadata"],
+            order=c["order"],
+            vector=[0.1] * 4096,
+        )
+        for c in chunk_data
+    ]
+    chunks_table.add(records_4096)
+
+    # Update settings to reflect the 4096-dim model used
+    settings_table = db.open_table("settings")
+    rows = (
+        settings_table.search().where("id = 'settings'").limit(1).to_arrow().to_pylist()
+    )
+    settings = json.loads(rows[0]["settings"])
+    settings["embeddings"]["model"]["vector_dim"] = 4096
+    settings["embeddings"]["model"]["name"] = "qwen3-embedding:8b"
+    settings_table.update(
+        where="id = 'settings'", values={"settings": json.dumps(settings)}
+    )
+
+    # Step 3: Open with skip_validation (different config) and run embed-only rebuild
+    # This should work: Store should use stored vector_dim for reading,
+    # then rebuild should migrate to new dimension
+    async with HaikuRAG(temp_db_path, skip_validation=True) as client:
+        processed_ids = [
+            doc_id
+            async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY)
+        ]
+        assert doc.id in processed_ids
+
+        # Verify chunks now have 2560-dim embeddings (from current config's model)
+        chunks_after = await client.chunk_repository.get_by_document_id(doc.id)
+        assert len(chunks_after) == len(chunks_before)
+
+        # Check that embeddings in DB are now 2560-dim
+        raw_chunks = (
+            client.store.chunks_table.search()
+            .where(f"document_id = '{doc.id}'")
+            .to_arrow()
+            .to_pylist()
+        )
+        for raw_chunk in raw_chunks:
+            assert len(raw_chunk["vector"]) == 2560
+
+        # Chunk IDs should be preserved
+        chunk_ids_before = {c.id for c in chunks_before}
+        chunk_ids_after = {c.id for c in chunks_after}
+        assert chunk_ids_before == chunk_ids_after
+
+
 @pytest.mark.vcr()
 async def test_rebuild_rechunk(qa_corpus: Dataset, temp_db_path):
     """Test rechunk rebuild: re-chunks from content without accessing source files."""