web_db.py

"""SQLite persistence for the PureLLM web interface."""

import json
import sqlite3
from pathlib import Path
from typing import Any, Optional


ROOT = Path(__file__).resolve().parent
DATA_DIR = ROOT / "data"
DB_PATH = DATA_DIR / "purellm.sqlite3"


def connect() -> sqlite3.Connection:
    DATA_DIR.mkdir(exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA foreign_keys = ON")
    return conn


def init_db() -> None:
    with connect() as conn:
        conn.executescript(
            """
            CREATE TABLE IF NOT EXISTS provider (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE,
                base_url TEXT,
                supported_models TEXT NOT NULL DEFAULT '[]',
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
            );

            CREATE TABLE IF NOT EXISTS official_model (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL UNIQUE,
                api_type TEXT NOT NULL DEFAULT 'openai',
                base_url TEXT,
                model_name TEXT NOT NULL,
                latest_baseline_run_id INTEGER,
                created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY(latest_baseline_run_id) REFERENCES benchmark_run(id)
            );

            CREATE TABLE IF NOT EXISTS benchmark_run (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                mode TEXT NOT NULL,
                api_type TEXT NOT NULL,
                provider_name TEXT,
                base_url TEXT,
                model TEXT NOT NULL,
                official_model_id INTEGER,
                baseline_run_id INTEGER,
                status TEXT NOT NULL,
                started_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
                completed_at TEXT,
                total_tests INTEGER NOT NULL DEFAULT 0,
                passed_tests INTEGER NOT NULL DEFAULT 0,
                overall_score REAL NOT NULL DEFAULT 0,
                authenticity_score REAL,
                category_scores TEXT NOT NULL DEFAULT '{}',
                run_json TEXT,
                comparison_json TEXT,
                error TEXT,
                FOREIGN KEY(official_model_id) REFERENCES official_model(id),
                FOREIGN KEY(baseline_run_id) REFERENCES benchmark_run(id)
            );

            CREATE TABLE IF NOT EXISTS test_result (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                run_id INTEGER NOT NULL,
                test_id TEXT NOT NULL,
                test_name TEXT NOT NULL,
                category TEXT NOT NULL,
                passed INTEGER NOT NULL,
                score REAL NOT NULL,
                details TEXT,
                elapsed_seconds REAL,
                token_count INTEGER,
                content TEXT,
                metadata TEXT,
                FOREIGN KEY(run_id) REFERENCES benchmark_run(id) ON DELETE CASCADE
            );

            CREATE INDEX IF NOT EXISTS idx_run_started_at ON benchmark_run(started_at DESC);
            CREATE INDEX IF NOT EXISTS idx_run_provider ON benchmark_run(provider_name);
            CREATE INDEX IF NOT EXISTS idx_run_model ON benchmark_run(model);
            CREATE INDEX IF NOT EXISTS idx_test_result_run ON test_result(run_id);
            """
        )

        seed_providers(conn)
        seed_official_models(conn)


def seed_providers(conn: sqlite3.Connection) -> None:
    providers_path = ROOT / "providers.json"
    if providers_path.exists():
        with providers_path.open("r", encoding="utf-8") as f:
            data = json.load(f)
        rows = [
            (p.get("name", ""), p.get("baseurl"), [])
            for group in ("official", "routers")
            for p in data.get(group, [])
            if p.get("name")
        ]
    else:
        rows = []

    for name, base_url, models in rows:
        conn.execute(
            """
            INSERT INTO provider (name, base_url, supported_models)
            VALUES (?, ?, ?)
            ON CONFLICT(name) DO UPDATE SET
                base_url = excluded.base_url,
                supported_models = excluded.supported_models
            """,
            (name, base_url, json.dumps(models)),
        )


def seed_official_models(conn: sqlite3.Connection) -> None:
    rows = [
        ("OpenAI GPT-4o", "openai", None, "gpt-4o"),
        ("OpenAI GPT-4o mini", "openai", None, "gpt-4o-mini"),
        ("Claude 3.5 Sonnet", "anthropic", None, "claude-3-5-sonnet-20241022"),
    ]
    for name, api_type, base_url, model_name in rows:
        conn.execute(
            """
            INSERT OR IGNORE INTO official_model (name, api_type, base_url, model_name)
            VALUES (?, ?, ?, ?)
            """,
            (name, api_type, base_url, model_name),
        )


def dict_row(row: sqlite3.Row) -> dict[str, Any]:
    return {k: row[k] for k in row.keys()}


def get_official_models() -> list[dict[str, Any]]:
    with connect() as conn:
        rows = conn.execute(
            """
            SELECT om.*, br.started_at AS baseline_started_at,
                   br.overall_score AS baseline_score
            FROM official_model om
            LEFT JOIN benchmark_run br ON br.id = om.latest_baseline_run_id
            ORDER BY om.name
            """
        ).fetchall()
    return [dict_row(r) for r in rows]


def get_providers() -> list[dict[str, Any]]:
    with connect() as conn:
        rows = conn.execute("SELECT * FROM provider ORDER BY name").fetchall()
    providers = []
    for row in rows:
        item = dict_row(row)
        item["supported_models"] = json.loads(item["supported_models"] or "[]")
        providers.append(item)
    return providers


def get_official_model(model_id: int) -> Optional[dict[str, Any]]:
    with connect() as conn:
        row = conn.execute(
            "SELECT * FROM official_model WHERE id = ?",
            (model_id,),
        ).fetchone()
    return dict_row(row) if row else None


def upsert_official_model(
    *, name: str, api_type: str, base_url: Optional[str], model_name: str
) -> int:
    with connect() as conn:
        row = conn.execute(
            "SELECT id FROM official_model WHERE name = ?",
            (name,),
        ).fetchone()
        if row:
            conn.execute(
                """
                UPDATE official_model
                SET api_type = ?, base_url = ?, model_name = ?
                WHERE id = ?
                """,
                (api_type, base_url, model_name, row["id"]),
            )
            return int(row["id"])
        cur = conn.execute(
            """
            INSERT INTO official_model (name, api_type, base_url, model_name)
            VALUES (?, ?, ?, ?)
            """,
            (name, api_type, base_url, model_name),
        )
        return int(cur.lastrowid)


def create_run(
    *,
    mode: str,
    api_type: str,
    model: str,
    base_url: Optional[str],
    provider_name: Optional[str] = None,
    official_model_id: Optional[int] = None,
    baseline_run_id: Optional[int] = None,
) -> int:
    with connect() as conn:
        cur = conn.execute(
            """
            INSERT INTO benchmark_run (
                mode, api_type, provider_name, base_url, model,
                official_model_id, baseline_run_id, status
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, 'running')
            """,
            (
                mode,
                api_type,
                provider_name,
                base_url,
                model,
                official_model_id,
                baseline_run_id,
            ),
        )
        return int(cur.lastrowid)


def update_run_complete(
    run_id: int,
    *,
    status: str,
    run_data: Optional[dict[str, Any]] = None,
    comparison: Optional[dict[str, Any]] = None,
    error: Optional[str] = None,
) -> None:
    run_data = run_data or {}
    summary = (comparison or {}).get("summary", {})
    with connect() as conn:
        conn.execute(
            """
            UPDATE benchmark_run
            SET status = ?, completed_at = CURRENT_TIMESTAMP,
                total_tests = ?, passed_tests = ?, overall_score = ?,
                authenticity_score = ?, category_scores = ?,
                run_json = ?, comparison_json = ?, error = ?
            WHERE id = ?
            """,
            (
                status,
                run_data.get("total_tests", 0),
                run_data.get("passed_tests", 0),
                run_data.get("overall_score", 0),
                summary.get("authenticity_score"),
                json.dumps(run_data.get("category_scores", {})),
                json.dumps(run_data),
                json.dumps(comparison) if comparison else None,
                error,
                run_id,
            ),
        )


def store_test_result(run_id: int, result: dict[str, Any]) -> None:
    with connect() as conn:
        conn.execute(
            """
            INSERT INTO test_result (
                run_id, test_id, test_name, category, passed, score, details,
                elapsed_seconds, token_count, content, metadata
            )
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                run_id,
                result.get("test_id"),
                result.get("test_name"),
                result.get("category"),
                1 if result.get("passed") else 0,
                result.get("score", 0),
                result.get("details"),
                result.get("elapsed_seconds"),
                result.get("token_count"),
                result.get("content"),
                json.dumps(result.get("metadata")) if result.get("metadata") else None,
            ),
        )


def set_latest_baseline(official_model_id: int, run_id: int) -> None:
    with connect() as conn:
        conn.execute(
            """
            UPDATE official_model
            SET latest_baseline_run_id = ?
            WHERE id = ?
            """,
            (run_id, official_model_id),
        )


def get_run(run_id: int) -> Optional[dict[str, Any]]:
    with connect() as conn:
        row = conn.execute(
            "SELECT * FROM benchmark_run WHERE id = ?",
            (run_id,),
        ).fetchone()
    if not row:
        return None
    item = dict_row(row)
    item["category_scores"] = json.loads(item["category_scores"] or "{}")
    item["run_json"] = json.loads(item["run_json"]) if item.get("run_json") else None
    item["comparison_json"] = (
        json.loads(item["comparison_json"]) if item.get("comparison_json") else None
    )
    return item


def get_baseline_data(run_id: int) -> dict[str, Any]:
    run = get_run(run_id)
    if not run or not run.get("run_json"):
        raise ValueError(f"Baseline run {run_id} has no saved run JSON")
    return {
        r["test_id"]: r
        for r in run["run_json"].get("results", [])
        if r.get("test_id")
    }


def list_runs(provider: Optional[str] = None, model: Optional[str] = None) -> list[dict[str, Any]]:
    query = """
        SELECT br.*, om.name AS official_model_name
        FROM benchmark_run br
        LEFT JOIN official_model om ON om.id = br.official_model_id
        WHERE 1 = 1
    """
    params: list[Any] = []
    if provider:
        query += " AND br.provider_name = ?"
        params.append(provider)
    if model:
        query += " AND br.model = ?"
        params.append(model)
    query += " ORDER BY br.started_at DESC"

    with connect() as conn:
        rows = conn.execute(query, params).fetchall()

    runs = []
    for row in rows:
        item = dict_row(row)
        item["category_scores"] = json.loads(item["category_scores"] or "{}")
        runs.append(item)
    return runs


def category_pass_counts(run_id: int) -> dict[str, dict[str, int]]:
    with connect() as conn:
        rows = conn.execute(
            """
            SELECT category, COUNT(*) AS total, SUM(passed) AS passed
            FROM test_result
            WHERE run_id = ?
            GROUP BY category
            ORDER BY category
            """,
            (run_id,),
        ).fetchall()
    return {
        row["category"]: {"passed": int(row["passed"] or 0), "total": int(row["total"])}
        for row in rows
    }