Prevent rebuild from throwing and recreate the embeddings table

This commit is contained in:
Yiorgis Gozadinos 2025-07-12 21:21:14 +03:00
parent 4444ec50da
commit 62a519b094
No known key found for this signature in database
4 changed files with 52 additions and 13 deletions

View file

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

View file

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

View file

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

View file

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