Spaces:
Sleeping
Sleeping
File size: 2,242 Bytes
fda8fb3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from __future__ import annotations
import sqlite3
from pathlib import Path
def connect(database_path: str) -> sqlite3.Connection:
path = Path(database_path)
path.parent.mkdir(parents=True, exist_ok=True)
connection = sqlite3.connect(path, check_same_thread=False)
connection.row_factory = sqlite3.Row
connection.execute("PRAGMA journal_mode=WAL")
connection.execute("PRAGMA foreign_keys=ON")
return connection
def initialize_schema(connection: sqlite3.Connection) -> None:
connection.executescript(
"""
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
status TEXT NOT NULL,
question TEXT NOT NULL,
model_name TEXT NOT NULL,
owner_id TEXT NOT NULL,
owner_name TEXT,
error TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS generation_results (
session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE,
answer TEXT NOT NULL,
raw_generation_text TEXT NOT NULL,
raw_trace_text TEXT NOT NULL,
normalized_trace_text TEXT NOT NULL,
sentences_json TEXT NOT NULL,
generation_metadata_json TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS analysis_results (
session_id TEXT PRIMARY KEY REFERENCES sessions(id) ON DELETE CASCADE,
result_json TEXT NOT NULL
);
"""
)
columns = {
row["name"] for row in connection.execute("PRAGMA table_info(generation_results)").fetchall()
}
if "raw_generation_text" not in columns:
connection.execute(
"ALTER TABLE generation_results ADD COLUMN raw_generation_text TEXT NOT NULL DEFAULT ''"
)
session_columns = {row["name"] for row in connection.execute("PRAGMA table_info(sessions)").fetchall()}
if "owner_id" not in session_columns:
connection.execute("ALTER TABLE sessions ADD COLUMN owner_id TEXT NOT NULL DEFAULT 'legacy-user'")
if "owner_name" not in session_columns:
connection.execute("ALTER TABLE sessions ADD COLUMN owner_name TEXT")
connection.commit()
|