Prevent rebuild from throwing and recreate the embeddings table
This commit is contained in:
parent
4444ec50da
commit
62a519b094
4 changed files with 52 additions and 13 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue