llm-speed-tester/app.py

# -*- 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)