Files
model-eval-site/database.py
T

683 lines
26 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
# -*- 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 '',
framework 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 ''")
cur = conn.execute("PRAGMA table_info(speed_items)")
scols = [r[1] for r in cur.fetchall()]
if "framework" not in scols:
conn.execute("ALTER TABLE speed_items ADD COLUMN framework 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, "
"MAX(si.framework) AS framework, "
"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,framework,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,framework,count,remark) VALUES(?,?,?,?,?,?,?)",
(provider, model, int(heat or 0), intro or "", (data.get("framework") or "").strip(),
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=?,framework=?,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(),
(data.get("framework") 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,framework,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()