"""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 = 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()