144 lines
4.3 KiB
Python
144 lines
4.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""LLM 速度测试台 - Flask 主应用"""
|
|
import json
|
|
|
|
from flask import Flask, jsonify, request, send_from_directory
|
|
|
|
import config
|
|
import database as db
|
|
from llm_providers import DEFAULT_URLS, ProviderError, call_stream
|
|
from tester import TestRunner
|
|
|
|
app = Flask(__name__, static_folder="static", static_url_path="")
|
|
app.json.ensure_ascii = False
|
|
|
|
db.init_db()
|
|
|
|
RUNNERS = {} # test_id -> TestRunner
|
|
|
|
|
|
@app.route("/")
|
|
def index():
|
|
return send_from_directory(app.static_folder, "index.html")
|
|
|
|
|
|
@app.route("/api/health")
|
|
def health():
|
|
running = [tid for tid, r in RUNNERS.items() if r.is_alive()]
|
|
return jsonify({"ok": True, "port": config.PORT, "running_tests": running})
|
|
|
|
|
|
def _fill_defaults(cfg):
|
|
p = cfg.get("provider", "openai")
|
|
if not cfg.get("base_url"):
|
|
cfg["base_url"] = DEFAULT_URLS.get(p, "")
|
|
if cfg.get("temperature") is None:
|
|
cfg["temperature"] = 0.7
|
|
return cfg
|
|
|
|
|
|
# ───────────────────────── 提供商配置 ─────────────────────────
|
|
|
|
@app.route("/api/configs", methods=["GET"])
|
|
def list_configs():
|
|
return jsonify(db.list_configs())
|
|
|
|
|
|
@app.route("/api/configs", methods=["POST"])
|
|
def add_config():
|
|
cfg = request.get_json(force=True) or {}
|
|
if not cfg.get("name"):
|
|
return jsonify({"ok": False, "error": "请填写配置名称"}), 400
|
|
cid = db.add_config(cfg)
|
|
return jsonify({"ok": True, "id": cid})
|
|
|
|
|
|
@app.route("/api/configs/<int:cid>", methods=["GET"])
|
|
def get_one_config(cid):
|
|
c = db.get_config(cid)
|
|
if not c:
|
|
return jsonify({"ok": False, "error": "配置不存在"}), 404
|
|
return jsonify(c)
|
|
|
|
|
|
@app.route("/api/configs/<int:cid>", methods=["DELETE"])
|
|
def del_config(cid):
|
|
db.delete_config(cid)
|
|
return jsonify({"ok": True})
|
|
|
|
|
|
@app.route("/api/configs/test", methods=["POST"])
|
|
def test_config():
|
|
cfg = _fill_defaults(request.get_json(force=True) or {})
|
|
if not cfg.get("api_key"):
|
|
return jsonify({"ok": False, "error": "请填写 API Key"}), 400
|
|
try:
|
|
m = call_stream(cfg, "你好，请只回复：OK", {"max_tokens": 16, "avoid_cache": False})
|
|
return jsonify({"ok": True, "total_ms": m["total_ms"], "metrics": m})
|
|
except ProviderError as e:
|
|
return jsonify({"ok": False, "error": str(e)})
|
|
except Exception as e:
|
|
return jsonify({"ok": False, "error": str(e)})
|
|
|
|
|
|
# ───────────────────────── 测试 ─────────────────────────
|
|
|
|
@app.route("/api/tests", methods=["POST"])
|
|
def start_test():
|
|
body = request.get_json(force=True) or {}
|
|
cfg = _fill_defaults(body.get("config") or {})
|
|
gen = body.get("gen") or {}
|
|
if not cfg.get("api_key"):
|
|
return jsonify({"ok": False, "error": "请填写 API Key"}), 400
|
|
if not cfg.get("model"):
|
|
return jsonify({"ok": False, "error": "请填写模型名称"}), 400
|
|
tid = db.create_test(cfg, gen)
|
|
runner = TestRunner(tid, cfg, gen)
|
|
RUNNERS[tid] = runner
|
|
runner.start()
|
|
return jsonify({"ok": True, "id": tid})
|
|
|
|
|
|
@app.route("/api/tests", methods=["GET"])
|
|
def list_tests():
|
|
return jsonify(db.list_tests())
|
|
|
|
|
|
@app.route("/api/tests/<int:tid>", methods=["GET"])
|
|
def get_test(tid):
|
|
t = db.get_test(tid)
|
|
if not t:
|
|
return jsonify({"ok": False, "error": "测试不存在"}), 404
|
|
t["runs"] = db.get_runs(tid)
|
|
t["logs"] = db.get_logs(tid)
|
|
return jsonify(t)
|
|
|
|
|
|
@app.route("/api/tests/<int:tid>/logs", methods=["GET"])
|
|
def get_logs(tid):
|
|
after = int(request.args.get("after", 0))
|
|
data = db.get_logs_after(tid, after)
|
|
if data is None:
|
|
return jsonify({"ok": False, "error": "测试不存在"}), 404
|
|
return jsonify(data)
|
|
|
|
|
|
@app.route("/api/tests/<int:tid>/cancel", methods=["POST"])
|
|
def cancel_test(tid):
|
|
r = RUNNERS.get(tid)
|
|
if r and r.is_alive():
|
|
r.request_cancel()
|
|
return jsonify({"ok": True, "msg": "正在停止..."})
|
|
return jsonify({"ok": False, "msg": "测试未在运行"})
|
|
|
|
|
|
@app.route("/api/tests/<int:tid>", methods=["DELETE"])
|
|
def del_test(tid):
|
|
db.delete_test(tid)
|
|
RUNNERS.pop(tid, None)
|
|
return jsonify({"ok": True})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
app.run(host=config.HOST, port=config.PORT, threaded=True, debug=False)
|