diff --git a/pyproject.toml b/pyproject.toml index 30dbb443..3bfd9784 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "haiku.rag" -version = "0.3.3" +version = "0.3.4" description = "Retrieval Augmented Generation (RAG) with SQLite" authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }] license = { text = "MIT" } diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index a2701ddf..72f17402 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -1,11 +1,16 @@ import sqlite3 import struct +from importlib import metadata from pathlib import Path from typing import Literal import sqlite_vec +from packaging.version import parse +from haiku.rag.config import Config from haiku.rag.embeddings import get_embedder +from haiku.rag.store.upgrades import upgrades +from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int class Store: @@ -13,7 +18,7 @@ class Store: self, db_path: Path | Literal[":memory:"], skip_validation: bool = False ): self.db_path: Path | Literal[":memory:"] = db_path - self._connection = self.create_db() + self.create_or_update_db() # Validate config compatibility after connection is established if not skip_validation: @@ -21,12 +26,35 @@ class Store: settings_repo = SettingsRepository(self) settings_repo.validate_config_compatibility() + current_version = metadata.version("haiku.rag") + self.set_user_version(current_version) - def create_db(self) -> sqlite3.Connection: + def create_or_update_db(self): """Create the database and tables with sqlite-vec support for embeddings.""" + current_version = metadata.version("haiku.rag") + db = sqlite3.connect(self.db_path) db.enable_load_extension(True) sqlite_vec.load(db) + self._connection = db + existing_tables = [ + row[0] + for row in db.execute( + "SELECT name FROM sqlite_master WHERE type='table';" + ).fetchall() + ] + + # If we have a db already, perform upgrades and return + if self.db_path != ":memory:" and "documents" in existing_tables: + # Upgrade database + db_version = self.get_user_version() + for version, steps in upgrades: + if parse(current_version) >= parse(version) and parse(version) > parse( + db_version + ): + for step in steps: + step(db) + return # Create documents table db.execute(""" @@ -39,7 +67,6 @@ class Store: updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ) """) - # Create chunks table db.execute(""" CREATE TABLE IF NOT EXISTS chunks ( @@ -50,7 +77,6 @@ class Store: FOREIGN KEY (document_id) REFERENCES documents (id) ON DELETE CASCADE ) """) - # Create vector table for chunk embeddings embedder = get_embedder() db.execute(f""" @@ -59,7 +85,6 @@ class Store: embedding FLOAT[{embedder._vector_dim}] ) """) - # Create FTS5 table for full-text search db.execute(""" CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5( @@ -68,7 +93,6 @@ class Store: content_rowid='id' ) """) - # Create settings table for storing current configuration db.execute(""" CREATE TABLE IF NOT EXISTS settings ( @@ -76,23 +100,35 @@ class Store: settings TEXT NOT NULL DEFAULT '{}' ) """) - - # Create indexes for better performance - db.execute( - "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)" - ) - # Save current settings to the new database - from haiku.rag.config import Config - settings_json = Config.model_dump_json() db.execute( "INSERT OR IGNORE INTO settings (id, settings) VALUES (1, ?)", (settings_json,), ) - + # Create indexes for better performance + db.execute( + "CREATE INDEX IF NOT EXISTS idx_chunks_document_id ON chunks(document_id)" + ) db.commit() - return db + + def get_user_version(self) -> str: + """Returns the SQLite user version""" + if self._connection is None: + raise ValueError("Store connection is not available") + + cursor = self._connection.execute("PRAGMA user_version;") + version = cursor.fetchone() + return int_to_semantic_version(version[0]) + + def set_user_version(self, version: str) -> None: + """Updates the SQLite user version""" + if self._connection is None: + raise ValueError("Store connection is not available") + + self._connection.execute( + f"PRAGMA user_version = {semantic_version_to_int(version)};" + ) def recreate_embeddings_table(self) -> None: """Recreate the embeddings table with current vector dimensions.""" diff --git a/src/haiku/rag/store/upgrades/__init__.py b/src/haiku/rag/store/upgrades/__init__.py new file mode 100644 index 00000000..3954309e --- /dev/null +++ b/src/haiku/rag/store/upgrades/__init__.py @@ -0,0 +1,3 @@ +from haiku.rag.store.upgrades.v0_3_4 import upgrades as v0_3_4_upgrades + +upgrades = v0_3_4_upgrades diff --git a/src/haiku/rag/store/upgrades/v0_3_4.py b/src/haiku/rag/store/upgrades/v0_3_4.py new file mode 100644 index 00000000..9ba912a1 --- /dev/null +++ b/src/haiku/rag/store/upgrades/v0_3_4.py @@ -0,0 +1,26 @@ +from collections.abc import Callable +from sqlite3 import Connection + +from haiku.rag.config import Config + + +def add_settings_table(db: Connection) -> None: + # Create settings table for storing current configuration + db.execute(""" + CREATE TABLE settings ( + id INTEGER PRIMARY KEY DEFAULT 1, + settings TEXT NOT NULL DEFAULT '{}' + ) + """) + + settings_json = Config.model_dump_json() + db.execute( + "INSERT INTO settings (id, settings) VALUES (1, ?)", + (settings_json,), + ) + db.commit() + + +upgrades: list[tuple[str, list[Callable[[Connection], None]]]] = [ + ("0.3.4", [add_settings_table]) +] diff --git a/src/haiku/rag/utils.py b/src/haiku/rag/utils.py index 8b78f008..91ead27b 100644 --- a/src/haiku/rag/utils.py +++ b/src/haiku/rag/utils.py @@ -29,6 +29,37 @@ def get_default_data_dir() -> Path: return data_path +def semantic_version_to_int(version: str) -> int: + """ + Convert a semantic version string to an integer. + + :param version: Semantic version string + :type version: str + :return: Integer representation of semantic version + :rtype: int + """ + major, minor, patch = version.split(".") + major = int(major) << 16 + minor = int(minor) << 8 + patch = int(patch) + return major + minor + patch + + +def int_to_semantic_version(version: int) -> str: + """ + Convert an integer to a semantic version string. + + :param version: Integer representation of semantic version + :type version: int + :return: Semantic version string + :rtype: str + """ + major = version >> 16 + minor = (version >> 8) & 255 + patch = version & 255 + return f"{major}.{minor}.{patch}" + + async def is_up_to_date() -> tuple[bool, Version, Version]: """ Checks whether haiku.rag is current. diff --git a/tests/test_settings.py b/tests/test_settings.py index 6a531685..08c21b07 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -50,7 +50,6 @@ async def test_config_validation_on_db_load(): # Create store and save settings store1 = Store(db_path) - SettingsRepository(store1) store1.close() # Change config diff --git a/tests/test_utils.py b/tests/test_utils.py new file mode 100644 index 00000000..cd5b742c --- /dev/null +++ b/tests/test_utils.py @@ -0,0 +1,15 @@ +from haiku.rag.utils import int_to_semantic_version, semantic_version_to_int + + +def test_sqlite_user_version(): + version = "0.1.5" + assert semantic_version_to_int(version) == 261 + assert int_to_semantic_version(261) == version + + version = "0.0.0" + assert semantic_version_to_int(version) == 0 + assert int_to_semantic_version(0) == version + + version = "255.255.255" + assert semantic_version_to_int(version) == 16777215 + assert int_to_semantic_version(16777215) == version diff --git a/uv.lock b/uv.lock index fd6b3649..d0ac66e2 100644 --- a/uv.lock +++ b/uv.lock @@ -816,7 +816,7 @@ wheels = [ [[package]] name = "haiku-rag" -version = "0.3.3" +version = "0.3.4" source = { editable = "." } dependencies = [ { name = "fastmcp" },