Store haiku.rag version, enabling upgrades
This commit is contained in:
parent
f4aef4eb9c
commit
0bf5606117
8 changed files with 129 additions and 19 deletions
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
3
src/haiku/rag/store/upgrades/__init__.py
Normal file
3
src/haiku/rag/store/upgrades/__init__.py
Normal 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
|
||||
26
src/haiku/rag/store/upgrades/v0_3_4.py
Normal file
26
src/haiku/rag/store/upgrades/v0_3_4.py
Normal 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])
|
||||
]
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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
15
tests/test_utils.py
Normal 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
|
||||
2
uv.lock
2
uv.lock
|
|
@ -816,7 +816,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "haiku-rag"
|
||||
version = "0.3.3"
|
||||
version = "0.3.4"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "fastmcp" },
|
||||
|
|
|
|||
Loading…
Reference in a new issue