295 lines
9.2 KiB
Python
295 lines
9.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""SQLite 存储：提供商配置 / 测试记录 / 每次采样指标 / 全量日志"""
|
|
import os
|
|
import json
|
|
import sqlite3
|
|
import threading
|
|
import time
|
|
|
|
import config
|
|
|
|
_lock = threading.RLock() # RLock：允许 get_logs_after 内嵌套调用 get_last_run
|
|
|
|
SCHEMA = """
|
|
CREATE TABLE IF NOT EXISTS configs(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT DEFAULT '',
|
|
provider TEXT DEFAULT 'openai',
|
|
base_url TEXT DEFAULT '',
|
|
api_key TEXT DEFAULT '',
|
|
model TEXT DEFAULT '',
|
|
temperature REAL DEFAULT 0.7,
|
|
created_at TEXT DEFAULT (datetime('now','localtime'))
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS tests(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
created_at TEXT DEFAULT (datetime('now','localtime')),
|
|
status TEXT DEFAULT 'running',
|
|
provider TEXT DEFAULT '',
|
|
model TEXT DEFAULT '',
|
|
config_json TEXT DEFAULT '{}',
|
|
gen_cfg_json TEXT DEFAULT '{}',
|
|
summary_json TEXT DEFAULT '{}',
|
|
error TEXT DEFAULT '',
|
|
started_at REAL DEFAULT 0
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS test_runs(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
test_id INTEGER NOT NULL,
|
|
run_index INTEGER DEFAULT 0,
|
|
metrics_json TEXT DEFAULT '{}',
|
|
error TEXT DEFAULT ''
|
|
);
|
|
|
|
CREATE TABLE IF NOT EXISTS logs(
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
test_id INTEGER NOT NULL,
|
|
level TEXT DEFAULT 'INFO',
|
|
msg TEXT DEFAULT '',
|
|
rel REAL DEFAULT 0,
|
|
ts TEXT DEFAULT (datetime('now','localtime'))
|
|
);
|
|
|
|
CREATE INDEX IF NOT EXISTS idx_runs_test ON test_runs(test_id);
|
|
CREATE INDEX IF NOT EXISTS idx_logs_test ON logs(test_id);
|
|
"""
|
|
|
|
|
|
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 init_db():
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
conn.executescript(SCHEMA)
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ───────────────────────── 提供商配置 ─────────────────────────
|
|
|
|
def add_config(cfg: dict) -> int:
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
cur = conn.execute(
|
|
"INSERT INTO configs(name,provider,base_url,api_key,model,temperature) VALUES(?,?,?,?,?,?)",
|
|
(cfg.get("name", ""), cfg.get("provider", "openai"),
|
|
cfg.get("base_url", ""), cfg.get("api_key", ""),
|
|
cfg.get("model", ""), float(cfg.get("temperature", 0.7))))
|
|
conn.commit()
|
|
return cur.lastrowid
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def list_configs():
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
rows = conn.execute("SELECT id,name,provider,base_url,model,temperature,"
|
|
"CASE WHEN api_key<>'' THEN 1 ELSE 0 END AS has_key "
|
|
"FROM configs ORDER BY id DESC").fetchall()
|
|
return [dict(r) for r in rows]
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_config(cid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
r = conn.execute("SELECT * FROM configs WHERE id=?", (cid,)).fetchone()
|
|
return dict(r) if r else None
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def delete_config(cid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
conn.execute("DELETE FROM configs WHERE id=?", (cid,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ───────────────────────── 测试记录 ─────────────────────────
|
|
|
|
def create_test(cfg: dict, gen: dict) -> int:
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
cur = conn.execute(
|
|
"INSERT INTO tests(status,provider,model,config_json,gen_cfg_json,started_at) "
|
|
"VALUES('running',?,?,?,?,?)",
|
|
(cfg.get("provider", "openai"), cfg.get("model", ""),
|
|
json.dumps(cfg, ensure_ascii=False), json.dumps(gen, ensure_ascii=False),
|
|
time.time()))
|
|
conn.commit()
|
|
return cur.lastrowid
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def update_status(tid: int, status: str, summary: dict = None, error: str = ""):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
conn.execute("UPDATE tests SET status=?, summary_json=?, error=? WHERE id=?",
|
|
(status,
|
|
json.dumps(summary or {}, ensure_ascii=False),
|
|
error or "", tid))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_test(tid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
r = conn.execute("SELECT * FROM tests WHERE id=?", (tid,)).fetchone()
|
|
if not r:
|
|
return None
|
|
d = dict(r)
|
|
d["config"] = json.loads(d.pop("config_json") or "{}")
|
|
d["gen"] = json.loads(d.pop("gen_cfg_json") or "{}")
|
|
d["summary"] = json.loads(d.pop("summary_json") or "{}")
|
|
return d
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def list_tests(limit=100):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT id,created_at,status,provider,model,summary_json,error "
|
|
"FROM tests ORDER BY id DESC LIMIT ?", (limit,)).fetchall()
|
|
out = []
|
|
for r in rows:
|
|
d = dict(r)
|
|
d["summary"] = json.loads(d.pop("summary_json") or "{}")
|
|
out.append(d)
|
|
return out
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def delete_test(tid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
conn.execute("DELETE FROM tests WHERE id=?", (tid,))
|
|
conn.execute("DELETE FROM test_runs WHERE test_id=?", (tid,))
|
|
conn.execute("DELETE FROM logs WHERE test_id=?", (tid,))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ───────────────────────── 采样指标 ─────────────────────────
|
|
|
|
def add_run(tid: int, run_index: int, metrics: dict, error: str = ""):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
conn.execute(
|
|
"INSERT INTO test_runs(test_id,run_index,metrics_json,error) VALUES(?,?,?,?)",
|
|
(tid, run_index, json.dumps(metrics, ensure_ascii=False), error))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_runs(tid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT run_index,metrics_json,error FROM test_runs "
|
|
"WHERE test_id=? ORDER BY run_index", (tid,)).fetchall()
|
|
out = []
|
|
for r in rows:
|
|
d = dict(r)
|
|
d["metrics"] = json.loads(d.pop("metrics_json") or "{}")
|
|
out.append(d)
|
|
return out
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_last_run(tid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
r = conn.execute(
|
|
"SELECT metrics_json FROM test_runs WHERE test_id=? "
|
|
"ORDER BY run_index DESC LIMIT 1", (tid,)).fetchone()
|
|
return json.loads(r["metrics_json"] or "{}") if r else None
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
# ───────────────────────── 日志 ─────────────────────────
|
|
|
|
def add_log(tid: int, level: str, msg: str, rel: float = None):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
if rel is None:
|
|
r = conn.execute("SELECT started_at FROM tests WHERE id=?", (tid,)).fetchone()
|
|
rel = (time.time() - (r["started_at"] or time.time())) if r else 0.0
|
|
conn.execute("INSERT INTO logs(test_id,level,msg,rel) VALUES(?,?,?,?)",
|
|
(tid, level, msg, round(rel, 3)))
|
|
conn.commit()
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_logs_after(tid: int, after_id: int = 0):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
t = conn.execute("SELECT status,summary_json,error FROM tests WHERE id=?", (tid,)).fetchone()
|
|
if not t:
|
|
return None
|
|
rows = conn.execute(
|
|
"SELECT id,level,msg,rel FROM logs WHERE test_id=? AND id>? ORDER BY id",
|
|
(tid, after_id)).fetchall()
|
|
logs = [dict(r) for r in rows]
|
|
last = logs[-1]["id"] if logs else after_id
|
|
return {
|
|
"status": t["status"],
|
|
"error": t["error"],
|
|
"summary": json.loads(t["summary_json"] or "{}"),
|
|
"last_run": get_last_run(tid),
|
|
"logs": logs,
|
|
"after": last,
|
|
}
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def get_logs(tid: int):
|
|
with _lock:
|
|
conn = _connect()
|
|
try:
|
|
rows = conn.execute(
|
|
"SELECT id,level,msg,rel,ts FROM logs WHERE test_id=? ORDER BY id", (tid,)).fetchall()
|
|
return [dict(r) for r in rows]
|
|
finally:
|
|
conn.close()
|