model-eval-site/database.py

# -*- coding: utf-8 -*-
"""SQLite 存储：账号 / 评测提交 / 能力测试 / 统计"""
import os
import json
import sqlite3
import threading

import config

_lock = threading.RLock()

SCHEMA = """
CREATE TABLE IF NOT EXISTS accounts(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  name TEXT UNIQUE NOT NULL,
  remark TEXT DEFAULT '',
  created_at TEXT DEFAULT (datetime('now','localtime'))
);

CREATE TABLE IF NOT EXISTS submissions(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  account_id INTEGER NOT NULL,
  source_test_id INTEGER DEFAULT 0,
  source_site TEXT DEFAULT 'llm-speed-tester',
  created_at TEXT DEFAULT (datetime('now','localtime')),
  provider TEXT DEFAULT '',
  model TEXT DEFAULT '',
  test_name TEXT DEFAULT '',
  hardware TEXT DEFAULT '',
  samples_ok INTEGER DEFAULT 0,
  samples_total INTEGER DEFAULT 0,
  avg_ttft_ms REAL,
  avg_prefill_speed REAL,
  avg_decode_speed REAL,
  avg_stream_decode REAL,
  avg_prompt_tokens REAL,
  avg_output_tokens REAL,
  avg_total_ms REAL,
  min_ttft_ms REAL, max_ttft_ms REAL,
  min_prefill_speed REAL, max_prefill_speed REAL,
  min_decode_speed REAL, max_decode_speed REAL,
  min_total_ms REAL, max_total_ms REAL,
  context_lengths TEXT DEFAULT '[]',
  concurrency_levels TEXT DEFAULT '[]',
  by_length TEXT DEFAULT '{}',
  by_concurrency TEXT DEFAULT '{}',
  runs TEXT DEFAULT '[]',
  raw TEXT DEFAULT '{}'
);

CREATE TABLE IF NOT EXISTS capabilities(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  name TEXT NOT NULL,
  category TEXT DEFAULT '',
  icon TEXT DEFAULT '🧪',
  description TEXT DEFAULT '',
  items TEXT DEFAULT '[]',
  status INTEGER DEFAULT 1,
  sort INTEGER DEFAULT 0,
  created_at TEXT DEFAULT (datetime('now','localtime')),
  updated_at TEXT DEFAULT (datetime('now','localtime'))
);

CREATE TABLE IF NOT EXISTS speed_items(
  id INTEGER PRIMARY KEY AUTOINCREMENT,
  provider TEXT DEFAULT '',
  model TEXT DEFAULT '',
  heat INTEGER DEFAULT 0,
  intro TEXT DEFAULT '',
  count INTEGER DEFAULT 0,
  remark TEXT DEFAULT '',
  updated_at TEXT DEFAULT (datetime('now','localtime'))
);
CREATE UNIQUE INDEX IF NOT EXISTS idx_si_model ON speed_items(provider, model);

CREATE TABLE IF NOT EXISTS settings(
  key TEXT PRIMARY KEY,
  value TEXT DEFAULT ''
);

CREATE INDEX IF NOT EXISTS idx_sub_model ON submissions(model, provider);
CREATE INDEX IF NOT EXISTS idx_sub_account ON submissions(account_id);
CREATE INDEX IF NOT EXISTS idx_sub_created ON submissions(created_at);
"""


def _connect():
    os.makedirs(config.DATA_DIR, exist_ok=True)
    conn = sqlite3.connect(config.DB_PATH, check_same_thread=False, timeout=30)
    conn.row_factory = sqlite3.Row
    conn.execute("PRAGMA journal_mode=WAL")
    return conn


def _migrate(conn):
    cur = conn.execute("PRAGMA table_info(submissions)")
    cols = [r[1] for r in cur.fetchall()]
    if "hardware" not in cols:
        conn.execute("ALTER TABLE submissions ADD COLUMN hardware TEXT DEFAULT ''")


def init_db():
    with _lock:
        conn = _connect()
        try:
            conn.executescript(SCHEMA)
            _migrate(conn)
            conn.commit()
        finally:
            conn.close()


def _jload(s, default):
    try:
        return json.loads(s or "")
    except Exception:
        return default


# ───────────────────────── 账号 ─────────────────────────

def get_or_create_account(name: str, remark: str = "") -> int:
    name = (name or "").strip() or "默认账号"
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT id FROM accounts WHERE name=?", (name,)).fetchone()
            if r:
                return r["id"]
            cur = conn.execute("INSERT INTO accounts(name, remark) VALUES(?,?)",
                               (name, (remark or "").strip()))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def add_account(name: str, remark: str = "") -> int:
    name = (name or "").strip()
    if not name:
        raise ValueError("账号名称不能为空")
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT id FROM accounts WHERE name=?", (name,)).fetchone()
            if r:
                raise ValueError("账号已存在")
            cur = conn.execute("INSERT INTO accounts(name, remark) VALUES(?,?)", (name, remark))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def list_accounts():
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT a.id,a.name,a.remark,a.created_at,"
                "(SELECT COUNT(*) FROM submissions s WHERE s.account_id=a.id) AS cnt "
                "FROM accounts a ORDER BY cnt DESC, a.id").fetchall()
            return [dict(r) for r in rows]
        finally:
            conn.close()


def rename_account(aid: int, name: str, remark: str):
    with _lock:
        conn = _connect()
        try:
            conn.execute("UPDATE accounts SET name=?, remark=? WHERE id=?",
                         (name.strip(), (remark or "").strip(), aid))
            conn.commit()
        finally:
            conn.close()


def delete_account(aid: int):
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM submissions WHERE account_id=?", (aid,))
            conn.execute("DELETE FROM accounts WHERE id=?", (aid,))
            conn.commit()
        finally:
            conn.close()


# ───────────────────────── 提交 ─────────────────────────

def add_submission(account_id: int, data: dict) -> int:
    """data 是 llm-speed-tester 发来的整包（含 summary/gen/runs/config/hardware）"""
    summary = data.get("summary") or {}
    gen = data.get("gen") or {}
    runs = data.get("runs") or []
    with _lock:
        conn = _connect()
        try:
            cur = conn.execute(
                "INSERT INTO submissions(account_id,source_test_id,source_site,provider,model,test_name,"
                "hardware,samples_ok,samples_total,avg_ttft_ms,avg_prefill_speed,avg_decode_speed,avg_stream_decode,"
                "avg_prompt_tokens,avg_output_tokens,avg_total_ms,min_ttft_ms,max_ttft_ms,min_prefill_speed,"
                "max_prefill_speed,min_decode_speed,max_decode_speed,min_total_ms,max_total_ms,"
                "context_lengths,concurrency_levels,by_length,by_concurrency,runs,raw) VALUES("
                "?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)",
                (account_id,
                 int(data.get("source_test_id") or 0),
                 data.get("source_site") or "llm-speed-tester",
                 data.get("provider") or "",
                 data.get("model") or "",
                 data.get("test_name") or data.get("name") or "",
                 (data.get("hardware") or "").strip(),
                 int(summary.get("samples_ok") or 0),
                 int(summary.get("samples_total") or 0),
                 summary.get("avg_ttft_ms"), summary.get("avg_prefill_speed"),
                 summary.get("avg_decode_speed"), summary.get("avg_stream_decode"),
                 summary.get("avg_prompt_tokens"), summary.get("avg_output_tokens"),
                 summary.get("avg_total_ms"),
                 summary.get("min_ttft_ms"), summary.get("max_ttft_ms"),
                 summary.get("min_prefill_speed"), summary.get("max_prefill_speed"),
                 summary.get("min_decode_speed"), summary.get("max_decode_speed"),
                 summary.get("min_total_ms"), summary.get("max_total_ms"),
                 json.dumps(gen.get("context_lengths") or summary.get("context_lengths") or [],
                            ensure_ascii=False),
                 json.dumps(summary.get("concurrency_levels") or gen.get("concurrency_levels") or [1],
                            ensure_ascii=False),
                 json.dumps(summary.get("by_length") or {}, ensure_ascii=False),
                 json.dumps(summary.get("by_concurrency") or {}, ensure_ascii=False),
                 json.dumps(runs[:200], ensure_ascii=False),
                 json.dumps(data, ensure_ascii=False)))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def get_submission(sid: int):
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT * FROM submissions WHERE id=?", (sid,)).fetchone()
            if not r:
                return None
            d = dict(r)
            d["account"] = conn.execute("SELECT name FROM accounts WHERE id=?",
                                        (d["account_id"],)).fetchone()["name"]
            for k in ("context_lengths", "concurrency_levels", "by_length",
                      "by_concurrency", "runs", "raw"):
                d[k] = _jload(d[k], "[]" if k in ("context_lengths", "concurrency_levels", "runs") else "{}")
            return d
        finally:
            conn.close()


def list_submissions(page=1, page_size=20, q="", account_id=None, model="", provider=""):
    with _lock:
        conn = _connect()
        try:
            where = ["1=1"]
            args = []
            if q:
                where.append("(s.model LIKE ? OR s.test_name LIKE ? OR a.name LIKE ? OR s.provider LIKE ? OR s.hardware LIKE ?)")
                like = "%" + q + "%"
                args += [like, like, like, like, like]
            if account_id:
                where.append("s.account_id=?")
                args.append(account_id)
            if model:
                where.append("s.model=?")
                args.append(model)
            if provider:
                where.append("s.provider=?")
                args.append(provider)
            wsql = " AND ".join(where)
            total = conn.execute(
                "SELECT COUNT(*) c FROM submissions s LEFT JOIN accounts a ON a.id=s.account_id WHERE " + wsql,
                args).fetchone()["c"]
            rows = conn.execute(
                "SELECT s.*, a.name AS account FROM submissions s "
                "LEFT JOIN accounts a ON a.id=s.account_id WHERE " + wsql +
                " ORDER BY s.id DESC LIMIT ? OFFSET ?",
                args + [page_size, (page - 1) * page_size]).fetchall()
            out = []
            for r in rows:
                d = dict(r)
                for k in ("context_lengths", "concurrency_levels"):
                    d[k] = _jload(d[k], "[]")
                out.append(d)
            return {"total": total, "items": out, "page": page, "page_size": page_size,
                    "pages": max(1, -(-total // page_size))}
        finally:
            conn.close()


def update_submission_hardware(sid: int, hardware: str):
    with _lock:
        conn = _connect()
        try:
            conn.execute("UPDATE submissions SET hardware=? WHERE id=?", (hardware.strip(), sid))
            conn.commit()
        finally:
            conn.close()


def delete_submission(sid: int):
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM submissions WHERE id=?", (sid,))
            conn.commit()
        finally:
            conn.close()


# ───────────────────────── 汇总统计 / 排行 ─────────────────────────

def get_stats():
    with _lock:
        conn = _connect()
        try:
            model_cnt = conn.execute("SELECT COUNT(DISTINCT model) c FROM submissions WHERE model<>''").fetchone()["c"]
            sub_cnt = conn.execute("SELECT COUNT(*) c FROM submissions").fetchone()["c"]
            acc_cnt = conn.execute("SELECT COUNT(*) c FROM accounts").fetchone()["c"]
            ok_cnt = conn.execute("SELECT SUM(samples_ok) c FROM submissions").fetchone()["c"] or 0
            cap_cnt = conn.execute("SELECT COUNT(*) c FROM capabilities").fetchone()["c"]
            return {"models": model_cnt, "submissions": sub_cnt,
                    "accounts": acc_cnt, "samples_ok": ok_cnt, "capabilities": cap_cnt}
        finally:
            conn.close()


def leaderboard(sort="heat", order="desc", limit=500):
    """按 (provider, model) 聚合所有提交 → 模型速度排行（含热度/配置简介/硬件）
    默认按热度(heat)倒序；sort 支持 heat/avg_decode_speed 等"""
    sort_whitelist = {
        "heat": "heat",
        "avg_decode_speed": "avg_decode_speed", "avg_prefill_speed": "avg_prefill_speed",
        "avg_ttft_ms": "avg_ttft_ms", "best_decode": "best_decode",
        "cnt": "cnt", "last_tested": "last_tested", "avg_stream_decode": "avg_stream_decode",
    }
    srt = sort_whitelist.get(sort, "heat")
    od = "DESC" if order == "asc" and srt not in ("avg_ttft_ms", "avg_total_ms") else (
        "ASC" if order == "asc" else "DESC")
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT s.provider, s.model, "
                "COUNT(*) AS cnt, "
                "COUNT(DISTINCT s.account_id) AS accounts, "
                "AVG(s.avg_decode_speed) AS avg_decode_speed, "
                "AVG(s.avg_prefill_speed) AS avg_prefill_speed, "
                "AVG(s.avg_stream_decode) AS avg_stream_decode, "
                "AVG(s.avg_ttft_ms) AS avg_ttft_ms, "
                "AVG(s.avg_output_tokens) AS avg_output_tokens, "
                "AVG(s.avg_total_ms) AS avg_total_ms, "
                "MAX(s.avg_decode_speed) AS best_decode, "
                "MAX(s.created_at) AS last_tested, "
                "COALESCE(MAX(si.heat), 0) AS heat, "
                "MAX(si.intro) AS intro, "
                "COALESCE(MAX(si.count), 0) AS count_override, "
                "(SELECT s2.hardware FROM submissions s2 "
                " WHERE s2.provider=s.provider AND s2.model=s.model AND s2.hardware<>'' "
                " ORDER BY s2.id DESC LIMIT 1) AS hardware "
                "FROM submissions s "
                "LEFT JOIN speed_items si ON si.provider=s.provider AND si.model=s.model "
                "WHERE s.model<>'' "
                "GROUP BY s.provider, s.model ORDER BY %s %s LIMIT ?" % (srt, od),
                (limit,)).fetchall()
            return [dict(r) for r in rows]
        finally:
            conn.close()


def get_model_summary(provider, model):
    """单个模型的聚合汇总（对比页用，含硬件）"""
    with _lock:
        conn = _connect()
        try:
            r = conn.execute(
                "SELECT COUNT(*) AS cnt, COUNT(DISTINCT account_id) AS accounts, "
                "AVG(avg_decode_speed) AS avg_decode_speed, "
                "AVG(avg_prefill_speed) AS avg_prefill_speed, "
                "AVG(avg_stream_decode) AS avg_stream_decode, "
                "AVG(avg_ttft_ms) AS avg_ttft_ms, "
                "AVG(avg_output_tokens) AS avg_output_tokens, "
                "AVG(avg_total_ms) AS avg_total_ms, "
                "SUM(samples_ok) AS samples_ok, "
                "MAX(avg_decode_speed) AS best_decode, "
                "MIN(avg_decode_speed) AS min_decode_speed, "
                "MAX(created_at) AS last_tested, "
                "(SELECT s2.hardware FROM submissions s2 "
                " WHERE s2.provider=submissions.provider AND s2.model=submissions.model AND s2.hardware<>'' "
                " ORDER BY s2.id DESC LIMIT 1) AS hardware "
                "FROM submissions WHERE provider=? AND model=?",
                (provider, model)).fetchone()
            return dict(r) if r else None
        finally:
            conn.close()


def get_model_detail(provider, model):
    """单个模型的所有提交 + 按上下文长度聚合"""
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT s.*, a.name AS account FROM submissions s "
                "LEFT JOIN accounts a ON a.id=s.account_id "
                "WHERE s.model=? AND s.provider=? ORDER BY s.id DESC",
                (model, provider)).fetchall()
            subs = []
            agg_length = {}
            for r in rows:
                d = dict(r)
                d["by_length"] = _jload(d.pop("by_length"), "{}")
                d["by_concurrency"] = _jload(d.pop("by_concurrency"), "{}")
                d["context_lengths"] = _jload(d.pop("context_lengths"), "[]")
                d["concurrency_levels"] = _jload(d.pop("concurrency_levels"), "[]")
                subs.append(d)
                for L, bl in d["by_length"].items():
                    a = agg_length.setdefault(str(L), {"count": 0, "sum_decode": 0.0,
                                                       "sum_prefill": 0.0, "sum_ttft": 0.0})
                    if bl.get("avg_decode_speed") is not None:
                        a["count"] += 1
                        a["sum_decode"] += float(bl["avg_decode_speed"] or 0)
                        a["sum_prefill"] += float(bl.get("avg_prefill_speed") or 0)
                        a["sum_ttft"] += float(bl.get("avg_ttft_ms") or 0)
            length_rows = []
            for L in sorted(agg_length, key=int):
                a = agg_length[L]
                if not a["count"]:
                    continue
                length_rows.append({
                    "length": int(L), "count": a["count"],
                    "avg_decode_speed": round(a["sum_decode"] / a["count"], 2),
                    "avg_prefill_speed": round(a["sum_prefill"] / a["count"], 2),
                    "avg_ttft_ms": round(a["sum_ttft"] / a["count"], 1),
                })
            return {"provider": provider, "model": model,
                    "submissions": subs, "by_length": length_rows,
                    "sub_cnt": len(subs)}
        finally:
            conn.close()


# ───────────────────────── 模型对比 ─────────────────────────

def compare_models(model_keys):
    """model_keys: [(provider, model), ...] → 每模型聚合汇总行"""
    out = []
    for provider, model in model_keys:
        r = get_model_summary(provider, model)
        if r and r.get("cnt"):
            out.append({"provider": provider, "model": model, **r})
    return out


# ───────────────────────── 能力测试模块 ─────────────────────────

def list_capabilities(status_only=True):
    with _lock:
        conn = _connect()
        try:
            sql = "SELECT * FROM capabilities"
            args = []
            if status_only:
                sql += " WHERE status=1"
            sql += " ORDER BY sort, id"
            rows = conn.execute(sql, args).fetchall()
            out = []
            for r in rows:
                d = dict(r)
                d["items"] = _jload(d["items"], "[]")
                out.append(d)
            return out
        finally:
            conn.close()


def add_capability(data: dict) -> int:
    with _lock:
        conn = _connect()
        try:
            cur = conn.execute(
                "INSERT INTO capabilities(name,category,icon,description,items,status,sort) "
                "VALUES(?,?,?,?,?,?,?)",
                ((data.get("name") or "").strip(),
                 (data.get("category") or "").strip(),
                 (data.get("icon") or "🧪").strip(),
                 (data.get("description") or "").strip(),
                 json.dumps(data.get("items") or [], ensure_ascii=False),
                 1 if data.get("status", 1) else 0,
                 int(data.get("sort") or 0)))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def update_capability(cid: int, data: dict):
    with _lock:
        conn = _connect()
        try:
            conn.execute(
                "UPDATE capabilities SET name=?,category=?,icon=?,description=?,items=?,status=?,sort=?,"
                "updated_at=datetime('now','localtime') WHERE id=?",
                ((data.get("name") or "").strip(),
                 (data.get("category") or "").strip(),
                 (data.get("icon") or "🧪").strip(),
                 (data.get("description") or "").strip(),
                 json.dumps(data.get("items") or [], ensure_ascii=False),
                 1 if data.get("status", 1) else 0,
                 int(data.get("sort") or 0), cid))
            conn.commit()
        finally:
            conn.close()


def delete_capability(cid: int):
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM capabilities WHERE id=?", (cid,))
            conn.commit()
        finally:
            conn.close()


# ───────────────────────── 速度项（热度/配置简介/个数，后台可改） ─────────────────────────

def touch_speed_item(provider, model):
    """提交到达时确保速度项存在；不存在则以当前提交数为默认热度"""
    provider = provider or ""
    model = (model or "").strip()
    if not model:
        return None
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT id FROM speed_items WHERE provider=? AND model=?",
                             (provider, model)).fetchone()
            if r:
                return r["id"]
            cnt = conn.execute("SELECT COUNT(*) c FROM submissions WHERE provider=? AND model=?",
                               (provider, model)).fetchone()["c"]
            cur = conn.execute(
                "INSERT INTO speed_items(provider,model,heat,intro,count) VALUES(?,?,?,?,?)",
                (provider, model, cnt, "", cnt))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def list_speed_items():
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT si.*, "
                "(SELECT COUNT(*) FROM submissions s WHERE s.provider=si.provider AND s.model=si.model) AS real_cnt "
                "FROM speed_items si ORDER BY si.heat DESC, si.model").fetchall()
            return [dict(r) for r in rows]
        finally:
            conn.close()


def add_speed_item(provider, model, heat=0, intro="", count=0, remark=""):
    provider = provider or ""
    model = (model or "").strip()
    if not model:
        raise ValueError("模型名称不能为空")
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT id FROM speed_items WHERE provider=? AND model=?",
                             (provider, model)).fetchone()
            if r:
                raise ValueError("该速度项已存在")
            cur = conn.execute(
                "INSERT INTO speed_items(provider,model,heat,intro,count,remark) VALUES(?,?,?,?,?,?)",
                (provider, model, int(heat or 0), intro or "", int(count or 0), remark or ""))
            conn.commit()
            return cur.lastrowid
        finally:
            conn.close()


def update_speed_item(sid: int, data: dict):
    with _lock:
        conn = _connect()
        try:
            conn.execute(
                "UPDATE speed_items SET provider=?,model=?,heat=?,intro=?,count=?,remark=?,"
                "updated_at=datetime('now','localtime') WHERE id=?",
                ((data.get("provider") or "").strip(),
                 (data.get("model") or "").strip(),
                 int(data.get("heat") or 0),
                 (data.get("intro") or "").strip(),
                 int(data.get("count") or 0),
                 (data.get("remark") or "").strip(),
                 sid))
            conn.commit()
        finally:
            conn.close()


def delete_speed_item(sid: int):
    with _lock:
        conn = _connect()
        try:
            conn.execute("DELETE FROM speed_items WHERE id=?", (sid,))
            conn.commit()
        finally:
            conn.close()


def sync_speed_items():
    """从提交自动同步速度项：为所有有提交的模型建速度项（默认热度=提交数）"""
    added = 0
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute(
                "SELECT provider, model, COUNT(*) c FROM submissions WHERE model<>'' "
                "GROUP BY provider, model").fetchall()
            for r in rows:
                r0 = conn.execute("SELECT id FROM speed_items WHERE provider=? AND model=?",
                                  (r["provider"], r["model"])).fetchone()
                if r0:
                    continue
                conn.execute(
                    "INSERT INTO speed_items(provider,model,heat,intro,count) VALUES(?,?,?,?,?)",
                    (r["provider"], r["model"], r["c"], "", r["c"]))
                added += 1
            conn.commit()
            return added
        finally:
            conn.close()


# ───────────────────────── 设置（水印/默认条数等） ─────────────────────────

def get_setting(key, default=""):
    with _lock:
        conn = _connect()
        try:
            r = conn.execute("SELECT value FROM settings WHERE key=?", (key,)).fetchone()
            return r["value"] if r else default
        finally:
            conn.close()


def set_setting(key, value):
    with _lock:
        conn = _connect()
        try:
            conn.execute("INSERT OR REPLACE INTO settings(key,value) VALUES(?,?)",
                         (key, str(value or "")))
            conn.commit()
        finally:
            conn.close()


def get_all_settings():
    with _lock:
        conn = _connect()
        try:
            rows = conn.execute("SELECT key,value FROM settings").fetchall()
            return {r["key"]: r["value"] for r in rows}
        finally:
            conn.close()