Store haiku.rag version, enabling upgrades

This commit is contained in:
Yiorgis Gozadinos 2025-07-16 15:32:24 +03:00
parent f4aef4eb9c
commit 0bf5606117
No known key found for this signature in database
8 changed files with 129 additions and 19 deletions

View file

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

View file

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

View file

@ -0,0 +1,3 @@
from haiku.rag.store.upgrades.v0_3_4 import upgrades as v0_3_4_upgrades
upgrades = v0_3_4_upgrades

View file

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

View file

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

View file

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

15
tests/test_utils.py Normal file
View file

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

View file

@ -816,7 +816,7 @@ wheels = [
[[package]]
name = "haiku-rag"
version = "0.3.3"
version = "0.3.4"
source = { editable = "." }
dependencies = [
{ name = "fastmcp" },