From 62a519b094cb5425e9aa558a139f461999db24df Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Sat, 12 Jul 2025 21:21:14 +0300
Subject: [PATCH] Prevent rebuild from throwing and recreate the embeddings
 table

---
 src/haiku/rag/app.py          |  2 +-
 src/haiku/rag/client.py       | 17 +++++++++++------
 src/haiku/rag/store/engine.py | 30 ++++++++++++++++++++++++++----
 tests/test_settings.py        | 16 ++++++++++++++--
 4 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py
index dedfce4c..0e2d1822 100644
--- a/src/haiku/rag/app.py
+++ b/src/haiku/rag/app.py
@@ -74,7 +74,7 @@ class HaikuRAGApp:
                 self.console.print(f"[red]Error: {e}[/red]")
 
     async def rebuild(self):
-        async with HaikuRAG(db_path=self.db_path) as client:
+        async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
             try:
                 documents = await client.list_documents()
                 total_docs = len(documents)
diff --git a/src/haiku/rag/client.py b/src/haiku/rag/client.py
index 74478654..72e5b0a6 100644
--- a/src/haiku/rag/client.py
+++ b/src/haiku/rag/client.py
@@ -24,12 +24,13 @@ class HaikuRAG:
         self,
         db_path: Path | Literal[":memory:"] = Config.DEFAULT_DATA_DIR
         / "haiku.rag.sqlite",
+        skip_validation: bool = False,
     ):
         """Initialize the RAG client with a database path."""
         if isinstance(db_path, Path):
             if not db_path.parent.exists():
                 Path.mkdir(db_path.parent, parents=True)
-        self.store = Store(db_path)
+        self.store = Store(db_path, skip_validation=skip_validation)
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
 
@@ -277,12 +278,16 @@ class HaikuRAG:
         Yields:
             int: The ID of the document currently being processed
         """
-        documents = await self.list_documents()
-
-        if not documents:
-            return
-
         await self.chunk_repository.delete_all()
+        self.store.recreate_embeddings_table()
+
+        # Update settings to current config
+        from haiku.rag.store.repositories.settings import SettingsRepository
+
+        settings_repo = SettingsRepository(self.store)
+        settings_repo.save()
+
+        documents = await self.list_documents()
 
         for doc in documents:
             if doc.id is not None:
diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py
index af7b27c4..a2701ddf 100644
--- a/src/haiku/rag/store/engine.py
+++ b/src/haiku/rag/store/engine.py
@@ -9,15 +9,18 @@ from haiku.rag.embeddings import get_embedder
 
 
 class Store:
-    def __init__(self, db_path: Path | Literal[":memory:"]):
+    def __init__(
+        self, db_path: Path | Literal[":memory:"], skip_validation: bool = False
+    ):
         self.db_path: Path | Literal[":memory:"] = db_path
         self._connection = self.create_db()
 
         # Validate config compatibility after connection is established
-        from haiku.rag.store.repositories.settings import SettingsRepository
+        if not skip_validation:
+            from haiku.rag.store.repositories.settings import SettingsRepository
 
-        settings_repo = SettingsRepository(self)
-        settings_repo.validate_config_compatibility()
+            settings_repo = SettingsRepository(self)
+            settings_repo.validate_config_compatibility()
 
     def create_db(self) -> sqlite3.Connection:
         """Create the database and tables with sqlite-vec support for embeddings."""
@@ -91,6 +94,25 @@ class Store:
         db.commit()
         return db
 
+    def recreate_embeddings_table(self) -> None:
+        """Recreate the embeddings table with current vector dimensions."""
+        if self._connection is None:
+            raise ValueError("Store connection is not available")
+
+        # Drop existing embeddings table
+        self._connection.execute("DROP TABLE IF EXISTS chunk_embeddings")
+
+        # Recreate with current dimensions
+        embedder = get_embedder()
+        self._connection.execute(f"""
+            CREATE VIRTUAL TABLE chunk_embeddings USING vec0(
+                chunk_id INTEGER PRIMARY KEY,
+                embedding FLOAT[{embedder._vector_dim}]
+            )
+        """)
+
+        self._connection.commit()
+
     @staticmethod
     def serialize_embedding(embedding: list[float]) -> bytes:
         """Serialize a list of floats to bytes for sqlite-vec storage."""
diff --git a/tests/test_settings.py b/tests/test_settings.py
index 8f6f49d7..0db16cc8 100644
--- a/tests/test_settings.py
+++ b/tests/test_settings.py
@@ -3,6 +3,7 @@ from pathlib import Path
 
 import pytest
 
+from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.store.engine import Store
 from haiku.rag.store.repositories.settings import (
@@ -41,7 +42,7 @@ def test_settings_save_and_retrieve():
     store.close()
 
 
-def test_config_validation_on_db_load():
+async def test_config_validation_on_db_load():
     """Test that config validation fails when loading db with mismatched settings."""
     # Create a temporary database file
     with tempfile.NamedTemporaryFile(suffix=".sqlite", delete=False) as tmp:
@@ -64,7 +65,18 @@ def test_config_validation_on_db_load():
         assert "CHUNK_SIZE" in str(exc_info.value)
         assert "Consider rebuilding" in str(exc_info.value)
 
-        # Restore original config
+        # Rebuild
+        async with HaikuRAG(db_path=db_path, skip_validation=True) as client:
+            async for _ in client.rebuild_database():
+                pass  # Process all documents
+
+        # Verify we can now load the database without exception (settings were updated)
+        store2 = Store(db_path)
+        settings_repo2 = SettingsRepository(store2)
+        db_settings = settings_repo2.get()
+        assert db_settings["CHUNK_SIZE"] == 999
+        store2.close()
+
         Config.CHUNK_SIZE = original_chunk_size
 
     finally: