381 lines
14 KiB
Python
381 lines
14 KiB
Python
"""HTTP 服务器：提供 API + 托管前端静态文件
|
|
|
|
启动后访问 http://localhost:8000 即可打开测试界面。
|
|
"""
|
|
import json
|
|
import logging
|
|
import mimetypes
|
|
import os
|
|
import threading
|
|
from dataclasses import asdict, dataclass, field
|
|
from datetime import datetime
|
|
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
|
from urllib.parse import urlparse
|
|
|
|
from . import config
|
|
from .client_base import ClientError
|
|
from .engine import TEST_CASES, run_all_tests
|
|
from .ollama_client import OllamaClient, OllamaError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class RunState:
|
|
"""测试运行状态（线程安全：所有读写通过 RunManager 加锁）"""
|
|
status: str = "idle" # idle | running | done | stopped | failed
|
|
run_id: str = ""
|
|
model: str = ""
|
|
repeats: int = 0
|
|
skip_heavy: bool = False
|
|
cases: list = field(default_factory=list)
|
|
current_case_index: int = 0
|
|
total_cases: int = 0
|
|
current_case_name: str = ""
|
|
message: str = "准备中"
|
|
progress: float = 0.0
|
|
started_at: str = ""
|
|
finished_at: str = ""
|
|
summary: dict = field(default_factory=dict)
|
|
|
|
def update(self, values=None, **kwargs) -> None:
|
|
"""批量更新字段（引擎通过 status.update(kw) 或 status.update(**kw) 汇报进度）"""
|
|
updates = dict(values or {})
|
|
updates.update(kwargs)
|
|
for key, value in updates.items():
|
|
if hasattr(self, key):
|
|
setattr(self, key, value)
|
|
|
|
def to_dict(self) -> dict:
|
|
return asdict(self)
|
|
|
|
|
|
class RunManager:
|
|
"""管理测试运行状态（单实例、同一时间只允许一次运行）"""
|
|
|
|
def __init__(self):
|
|
self._lock = threading.Lock()
|
|
self._stop_event = threading.Event()
|
|
self._state = RunState()
|
|
self._thread = None
|
|
|
|
@property
|
|
def is_running(self) -> bool:
|
|
return self._state.status == "running"
|
|
|
|
def start(self, cases: list, repeats: int, model: str, skip_heavy: bool) -> str:
|
|
with self._lock:
|
|
if self.is_running:
|
|
raise RuntimeError("已有测试正在运行，请等待完成或先停止")
|
|
self._stop_event.clear()
|
|
run_id = datetime.now().strftime("%Y%m%d-%H%M%S")
|
|
self._state = RunState(
|
|
status="running",
|
|
run_id=run_id,
|
|
model=model,
|
|
repeats=repeats,
|
|
skip_heavy=skip_heavy,
|
|
cases=list(cases),
|
|
total_cases=len(cases),
|
|
started_at=datetime.now().isoformat(),
|
|
)
|
|
self._thread = threading.Thread(
|
|
target=self._run, args=(cases, repeats, model, skip_heavy), daemon=True
|
|
)
|
|
self._thread.start()
|
|
return run_id
|
|
|
|
def stop(self) -> None:
|
|
"""请求停止（置位事件，引擎在用例之间检查）"""
|
|
self._stop_event.set()
|
|
|
|
def status(self) -> dict:
|
|
"""返回当前运行状态（对外只读拷贝）"""
|
|
with self._lock:
|
|
return self._state.to_dict()
|
|
|
|
def _run(self, cases, repeats, model, skip_heavy):
|
|
try:
|
|
report = run_all_tests(
|
|
cases=cases,
|
|
repeats=repeats,
|
|
model=model,
|
|
skip_heavy=skip_heavy,
|
|
status=self._state,
|
|
stop_event=self._stop_event,
|
|
)
|
|
stopped = report["summary"].get("stopped", False)
|
|
with self._lock:
|
|
self._state.status = "stopped" if stopped else "done"
|
|
self._state.progress = 100
|
|
self._state.message = "测试已停止" if stopped else "测试完成"
|
|
self._state.finished_at = datetime.now().isoformat()
|
|
self._state.summary = report["summary"]
|
|
except Exception as e:
|
|
logger.exception("测试引擎异常")
|
|
with self._lock:
|
|
self._state.status = "failed"
|
|
self._state.message = f"测试失败: {e}"
|
|
self._state.finished_at = datetime.now().isoformat()
|
|
|
|
|
|
manager = RunManager()
|
|
|
|
|
|
class Handler(BaseHTTPRequestHandler):
|
|
"""API + 静态文件处理"""
|
|
|
|
# ---------- HTTP 基础 ----------
|
|
|
|
def _send_json(self, obj, code=200):
|
|
body = json.dumps(obj, ensure_ascii=False).encode("utf-8")
|
|
self.send_response(code)
|
|
self.send_header("Content-Type", "application/json; charset=utf-8")
|
|
self.send_header("Content-Length", str(len(body)))
|
|
self.end_headers()
|
|
self.wfile.write(body)
|
|
|
|
def _read_json(self) -> dict:
|
|
try:
|
|
length = int(self.headers.get("Content-Length", 0) or 0)
|
|
except ValueError:
|
|
length = 0
|
|
if length <= 0:
|
|
return {}
|
|
try:
|
|
return json.loads(self.rfile.read(length).decode("utf-8"))
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
|
|
def log_message(self, format, *args):
|
|
print(f"[API] {args[0]}")
|
|
|
|
# ---------- GET ----------
|
|
|
|
def do_GET(self):
|
|
path = urlparse(self.path).path
|
|
if path == "/api/health":
|
|
self._send_json({"status": "ok"})
|
|
elif path == "/api/config":
|
|
self._handle_config()
|
|
elif path == "/api/models":
|
|
self._handle_models()
|
|
elif path == "/api/cases":
|
|
cases = {
|
|
cid: {
|
|
"name": info["name"],
|
|
"difficulty": info["difficulty"],
|
|
"estimated_seconds": info["estimated_seconds"],
|
|
}
|
|
for cid, info in TEST_CASES.items()
|
|
}
|
|
self._send_json(cases)
|
|
elif path == "/api/status":
|
|
self._send_json(manager.status())
|
|
elif path == "/api/results/latest.json":
|
|
self._handle_latest()
|
|
elif path == "/api/results/list":
|
|
self._handle_results_list()
|
|
elif path.startswith("/api/results/"):
|
|
self._handle_result_file(path)
|
|
elif path == "/report.html":
|
|
self._serve_file(os.path.join(config.REPORT_DIR, "report.html"))
|
|
else:
|
|
self._serve_static(path)
|
|
|
|
# ---------- POST ----------
|
|
|
|
def do_POST(self):
|
|
path = urlparse(self.path).path
|
|
body = self._read_json()
|
|
if path == "/api/run":
|
|
self._handle_run(body)
|
|
elif path == "/api/stop":
|
|
manager.stop()
|
|
self._send_json({"ok": True, "message": "已请求停止"})
|
|
else:
|
|
self._send_json({"error": "Not found"}, 404)
|
|
|
|
# ---------- API 处理器 ----------
|
|
|
|
def _handle_config(self):
|
|
self._send_json({
|
|
"backend": config.LLM_BACKEND,
|
|
"ollama_base_url": config.OLLAMA_BASE_URL,
|
|
"openai_base_url": config.OPENAI_BASE_URL,
|
|
"default_model": config.DEFAULT_MODEL,
|
|
"openai_default_model": config.OPENAI_DEFAULT_MODEL,
|
|
"concurrency": config.DEFAULT_CONCURRENCY,
|
|
})
|
|
|
|
def _handle_models(self):
|
|
try:
|
|
if config.LLM_BACKEND == "openai":
|
|
from .openai_client import OpenAIClient
|
|
models = OpenAIClient().list_models()
|
|
else:
|
|
models = OllamaClient().list_models()
|
|
self._send_json({"models": models, "error": None})
|
|
except ClientError as e:
|
|
self._send_json({"models": [], "error": str(e)})
|
|
|
|
def _handle_run(self, body):
|
|
cases = body.get("cases") or []
|
|
repeats = int(body.get("repeats", 3) or 3)
|
|
model = body.get("model") or config.DEFAULT_MODEL
|
|
skip_heavy = bool(body.get("skip_heavy", False))
|
|
|
|
# 校验用例
|
|
unknown = [c for c in cases if c not in TEST_CASES]
|
|
if not cases:
|
|
self._send_json({"error": "未选择测试用例"}, 400)
|
|
return
|
|
if unknown:
|
|
self._send_json({"error": f"未知用例: {', '.join(unknown)}"}, 400)
|
|
return
|
|
if repeats < 1 or repeats > 50:
|
|
self._send_json({"error": "重复次数需在 1-50 之间"}, 400)
|
|
return
|
|
|
|
try:
|
|
run_id = manager.start(cases, repeats, model, skip_heavy)
|
|
self._send_json({"run_id": run_id})
|
|
except RuntimeError as e:
|
|
self._send_json({"error": str(e)}, 409)
|
|
|
|
def _handle_latest(self):
|
|
"""返回最新测试结果，文件损坏时优雅降级"""
|
|
latest_path = os.path.join(config.RESULTS_DIR, "latest.json")
|
|
if not os.path.exists(latest_path):
|
|
self._send_json({"error": "暂无测试结果，请先运行测试"}, 404)
|
|
return
|
|
try:
|
|
with open(latest_path, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
self._send_json(data)
|
|
except (json.JSONDecodeError, OSError) as e:
|
|
logger.error("读取 latest.json 失败: %s", e)
|
|
self._send_json({"error": "结果文件损坏，请先运行新测试"}, 500)
|
|
|
|
def _handle_results_list(self):
|
|
"""列出所有 run_*.json 结果文件"""
|
|
import glob as glob_mod
|
|
pattern = os.path.join(config.RESULTS_DIR, "run_*.json")
|
|
files = sorted(glob_mod.glob(pattern))
|
|
result = []
|
|
for f in files:
|
|
name = os.path.basename(f)
|
|
try:
|
|
with open(f, "r", encoding="utf-8") as fh:
|
|
data = json.load(fh)
|
|
summary = data.get("summary", {})
|
|
result.append({
|
|
"filename": name,
|
|
"timestamp": data.get("timestamp", ""),
|
|
"model": data.get("environment", {}).get("model", ""),
|
|
"total_cases": summary.get("total_cases", 0),
|
|
"passed": summary.get("passed", 0),
|
|
"failed": summary.get("failed", 0),
|
|
"elapsed_seconds": summary.get("total_elapsed_seconds", 0),
|
|
})
|
|
except (json.JSONDecodeError, OSError):
|
|
result.append({"filename": name, "error": "无法读取"})
|
|
self._send_json({"results": result})
|
|
|
|
def _handle_result_file(self, path):
|
|
"""返回单个历史结果文件 /api/results/<filename>"""
|
|
filename = path.rsplit("/", 1)[-1]
|
|
if not filename.endswith(".json") or "/" in filename or "\\" in filename:
|
|
self._send_json({"error": "Forbidden"}, 403)
|
|
return
|
|
full = os.path.join(config.RESULTS_DIR, filename)
|
|
if not os.path.isfile(full):
|
|
self._send_json({"error": "Not found"}, 404)
|
|
return
|
|
try:
|
|
with open(full, "r", encoding="utf-8") as f:
|
|
data = json.load(f)
|
|
self._send_json(data)
|
|
except (json.JSONDecodeError, OSError):
|
|
self._send_json({"error": "文件损坏"}, 500)
|
|
|
|
# ---------- 静态文件 ----------
|
|
|
|
def _serve_static(self, path):
|
|
# 只允许从 frontend/ 目录读取
|
|
if path == "/" or path == "/index.html":
|
|
rel = "index.html"
|
|
else:
|
|
rel = path.lstrip("/")
|
|
|
|
full = os.path.normpath(os.path.join(config.FRONTEND_DIR, rel))
|
|
frontend_root = os.path.normpath(config.FRONTEND_DIR)
|
|
|
|
# 防止路径遍历攻击
|
|
if not os.path.abspath(full).startswith(os.path.abspath(frontend_root + os.sep)):
|
|
self._send_json({"error": "Forbidden"}, 403)
|
|
return
|
|
|
|
# 防止符号链接指向外部目录
|
|
# 逐段检查路径中每个组件是否安全
|
|
parts = os.path.relpath(full, frontend_root).split(os.sep)
|
|
current = frontend_root
|
|
for part in parts:
|
|
current = os.path.join(current, part)
|
|
if os.path.islink(current):
|
|
link_target = os.path.realpath(current)
|
|
if not link_target.startswith(os.path.realpath(frontend_root) + os.sep):
|
|
self._send_json({"error": "Forbidden"}, 403)
|
|
return
|
|
if not os.path.exists(current):
|
|
break
|
|
|
|
self._serve_file(full)
|
|
|
|
def _serve_file(self, full_path):
|
|
if not os.path.isfile(full_path):
|
|
self._send_json({"error": "Not found"}, 404)
|
|
return
|
|
content_type = mimetypes.guess_type(full_path)[0] or "application/octet-stream"
|
|
with open(full_path, "rb") as f:
|
|
data = f.read()
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", content_type)
|
|
self.send_header("Content-Length", str(len(data)))
|
|
self.end_headers()
|
|
self.wfile.write(data)
|
|
|
|
|
|
def main():
|
|
# 配置根日志
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
|
datefmt="%H:%M:%S",
|
|
)
|
|
|
|
server = ThreadingHTTPServer((config.SERVER_HOST, config.SERVER_PORT), Handler)
|
|
backend = config.LLM_BACKEND
|
|
print("=" * 60)
|
|
print(f"大模型速度测试助手 v2.1 ({'OpenAI 兼容' if backend == 'openai' else 'Ollama'})")
|
|
print("=" * 60)
|
|
if backend == "openai":
|
|
print(f"服务地址 : {config.OPENAI_BASE_URL}")
|
|
print(f"默认模型 : {config.OPENAI_DEFAULT_MODEL}")
|
|
print(f"并发测试 : {config.DEFAULT_CONCURRENCY} 并发")
|
|
else:
|
|
print(f"Ollama 地址 : {config.OLLAMA_BASE_URL}")
|
|
print(f"默认模型 : {config.DEFAULT_MODEL}")
|
|
print(f"前端界面 : http://localhost:{config.SERVER_PORT}")
|
|
print("按 Ctrl+C 停止服务器")
|
|
print("=" * 60)
|
|
try:
|
|
server.serve_forever()
|
|
except KeyboardInterrupt:
|
|
print("\n服务器已停止")
|
|
server.server_close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|