zwitschi
712c556032
- Updated `refresh_models_cache` to include output modalities in the models cache. - Added `get_model_output_modalities` function to retrieve output modalities for a specific model. - Modified tests to cover new functionality for output modalities. - Updated OpenRouter video generation functions to support audio generation and improved error handling. - Enhanced dashboard to display generated images and videos. - Refactored frontend templates to accommodate new data structures for generated content. - Adjusted tests to validate changes in model handling and dashboard rendering. Co-authored-by: Copilot <copilot@github.com>
140 lines
4.6 KiB
Python
140 lines
4.6 KiB
Python
"""DuckDB singleton connection with asyncio write lock and schema migrations."""
|
|
import asyncio
|
|
import os
|
|
import duckdb
|
|
|
|
_conn: duckdb.DuckDBPyConnection | None = None
|
|
_write_lock = asyncio.Lock()
|
|
|
|
|
|
def get_db_path() -> str:
|
|
return os.getenv("DB_PATH", "data/app.db")
|
|
|
|
|
|
def init_db(path: str | None = None) -> duckdb.DuckDBPyConnection:
|
|
"""Open (or reuse) the DuckDB connection and run schema migrations."""
|
|
global _conn
|
|
if _conn is not None:
|
|
return _conn
|
|
db_path = path or get_db_path()
|
|
if db_path != ":memory:":
|
|
os.makedirs(os.path.dirname(db_path), exist_ok=True)
|
|
_conn = duckdb.connect(db_path)
|
|
_run_migrations(_conn)
|
|
return _conn
|
|
|
|
|
|
def get_conn() -> duckdb.DuckDBPyConnection:
|
|
"""Return the active connection; raises if not yet initialised."""
|
|
if _conn is None:
|
|
raise RuntimeError("Database not initialised. Call init_db() first.")
|
|
return _conn
|
|
|
|
|
|
def close_db() -> None:
|
|
"""Close the connection (called on app shutdown)."""
|
|
global _conn
|
|
if _conn is not None:
|
|
_conn.close()
|
|
_conn = None
|
|
|
|
|
|
def get_write_lock() -> asyncio.Lock:
|
|
"""Return the asyncio lock that serialises write operations."""
|
|
return _write_lock
|
|
|
|
|
|
def _run_migrations(conn: duckdb.DuckDBPyConnection) -> None:
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id UUID DEFAULT uuid() PRIMARY KEY,
|
|
email VARCHAR NOT NULL UNIQUE,
|
|
password_hash VARCHAR NOT NULL,
|
|
role VARCHAR DEFAULT 'user',
|
|
created_at TIMESTAMP DEFAULT now(),
|
|
updated_at TIMESTAMP DEFAULT now()
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS refresh_tokens (
|
|
jti UUID DEFAULT uuid() PRIMARY KEY,
|
|
user_id UUID NOT NULL,
|
|
issued_at TIMESTAMP DEFAULT now(),
|
|
expires_at TIMESTAMP NOT NULL,
|
|
revoked BOOLEAN DEFAULT false
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS uploaded_images (
|
|
id UUID DEFAULT uuid() PRIMARY KEY,
|
|
user_id UUID NOT NULL,
|
|
filename VARCHAR NOT NULL,
|
|
content_type VARCHAR NOT NULL,
|
|
file_path VARCHAR NOT NULL,
|
|
size_bytes BIGINT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT now()
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS models_cache (
|
|
id UUID DEFAULT uuid() PRIMARY KEY,
|
|
model_id VARCHAR NOT NULL UNIQUE,
|
|
name VARCHAR NOT NULL,
|
|
modality VARCHAR NOT NULL,
|
|
context_length BIGINT,
|
|
pricing JSON,
|
|
fetched_at TIMESTAMP NOT NULL
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS generated_images (
|
|
id UUID DEFAULT uuid() PRIMARY KEY,
|
|
user_id UUID NOT NULL,
|
|
model_id VARCHAR NOT NULL,
|
|
prompt VARCHAR NOT NULL,
|
|
image_data VARCHAR NOT NULL,
|
|
created_at TIMESTAMP DEFAULT now()
|
|
)
|
|
""")
|
|
conn.execute("""
|
|
CREATE TABLE IF NOT EXISTS generated_videos (
|
|
id UUID DEFAULT uuid() PRIMARY KEY,
|
|
user_id UUID NOT NULL,
|
|
job_id VARCHAR NOT NULL,
|
|
model_id VARCHAR NOT NULL,
|
|
prompt VARCHAR NOT NULL,
|
|
polling_url VARCHAR,
|
|
status VARCHAR NOT NULL DEFAULT 'pending',
|
|
video_url VARCHAR,
|
|
created_at TIMESTAMP DEFAULT now(),
|
|
updated_at TIMESTAMP DEFAULT now()
|
|
)
|
|
""")
|
|
# Migration: add output_modalities column if absent (stores JSON array string)
|
|
conn.execute("""
|
|
ALTER TABLE models_cache ADD COLUMN IF NOT EXISTS output_modalities VARCHAR
|
|
""")
|
|
_seed_admin(conn)
|
|
|
|
|
|
def _seed_admin(conn: duckdb.DuckDBPyConnection) -> None:
|
|
"""Insert the default admin user if it doesn't already exist."""
|
|
from passlib.context import CryptContext
|
|
_pwd = CryptContext(schemes=["bcrypt"], deprecated="auto")
|
|
|
|
email = os.getenv("ADMIN_EMAIL", "ai@allucanget.biz")
|
|
password = os.getenv("ADMIN_PASSWORD", "admin123")
|
|
|
|
existing = conn.execute(
|
|
"SELECT id FROM users WHERE email = ?", [email]
|
|
).fetchone()
|
|
if existing is None:
|
|
password_hash = _pwd.hash(password)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO users (email, password_hash, role)
|
|
VALUES (?, ?, 'admin')
|
|
""",
|
|
[email, password_hash],
|
|
)
|