55 lines
1.7 KiB
Python
55 lines
1.7 KiB
Python
from __future__ import annotations
|
|
|
|
import sqlite3
|
|
|
|
from backend.database.migrations.runner import Migration
|
|
|
|
|
|
def upgrade(connection: sqlite3.Connection) -> None:
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE llm_models (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
display_name TEXT NOT NULL,
|
|
display_name_key TEXT NOT NULL UNIQUE,
|
|
base_url TEXT NOT NULL,
|
|
model_identifier TEXT NOT NULL,
|
|
encrypted_api_key TEXT NOT NULL,
|
|
created_at TEXT NOT NULL,
|
|
updated_at TEXT NOT NULL,
|
|
updated_by INTEGER NOT NULL REFERENCES users(id) ON DELETE RESTRICT
|
|
)
|
|
"""
|
|
)
|
|
connection.execute(
|
|
"""
|
|
CREATE TABLE llm_configuration (
|
|
id INTEGER PRIMARY KEY CHECK (id = 1),
|
|
primary_model_id INTEGER REFERENCES llm_models(id) ON DELETE RESTRICT,
|
|
fallback_model_id INTEGER REFERENCES llm_models(id) ON DELETE RESTRICT,
|
|
updated_at TEXT,
|
|
updated_by INTEGER REFERENCES users(id) ON DELETE RESTRICT,
|
|
CHECK (
|
|
fallback_model_id IS NULL
|
|
OR primary_model_id IS NULL
|
|
OR fallback_model_id <> primary_model_id
|
|
)
|
|
)
|
|
"""
|
|
)
|
|
connection.execute("INSERT INTO llm_configuration (id) VALUES (1)")
|
|
|
|
|
|
def downgrade(connection: sqlite3.Connection) -> None:
|
|
connection.execute("DROP TABLE llm_configuration")
|
|
connection.execute("DROP TABLE llm_models")
|
|
|
|
|
|
MIGRATION = Migration(
|
|
version=2,
|
|
name="create_model_pool",
|
|
signature="model-pool:v1:models-primary-fallback-encrypted-key",
|
|
upgrade=upgrade,
|
|
downgrade=downgrade,
|
|
)
|