286 lines
9.8 KiB
Python
286 lines
9.8 KiB
Python
#!/usr/bin/env python3
|
||
"""独立模型连通性检测脚本（中文日志输出）。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import os
|
||
import ssl
|
||
import sys
|
||
import time
|
||
import urllib.error
|
||
import urllib.request
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
|
||
|
||
def load_env_file(env_path: Path) -> None:
|
||
if not env_path.exists():
|
||
return
|
||
for raw_line in env_path.read_text(encoding="utf-8").splitlines():
|
||
line = raw_line.strip()
|
||
if not line or line.startswith("#") or "=" not in line:
|
||
continue
|
||
key, value = line.split("=", 1)
|
||
key = key.strip()
|
||
value = value.strip()
|
||
if len(value) >= 2 and ((value[0] == value[-1] == '"') or (value[0] == value[-1] == "'")):
|
||
value = value[1:-1]
|
||
os.environ.setdefault(key, value)
|
||
|
||
|
||
def provider_runtime(provider: str) -> dict[str, str]:
|
||
p = provider.strip().lower()
|
||
if p in {"glm", "zhipu"}:
|
||
return {
|
||
"provider": p,
|
||
"base_url": str(os.getenv("GLM_BASE_URL", "")).strip(),
|
||
"api_key": str(os.getenv("GLM_API_KEY", "")).strip(),
|
||
}
|
||
if p == "openai":
|
||
return {
|
||
"provider": p,
|
||
"base_url": str(os.getenv("OPENAI_BASE_URL", "https://api.openai.com/v1")).strip(),
|
||
"api_key": str(os.getenv("OPENAI_API_KEY", "")).strip(),
|
||
}
|
||
if p == "deepseek":
|
||
return {
|
||
"provider": p,
|
||
"base_url": str(os.getenv("DEEPSEEK_BASE_URL", "https://api.deepseek.com/v1")).strip(),
|
||
"api_key": str(os.getenv("DEEPSEEK_API_KEY", "")).strip(),
|
||
}
|
||
if p == "ollama":
|
||
return {
|
||
"provider": p,
|
||
"base_url": str(os.getenv("OLLAMA_BASE_URL", os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434"))).strip(),
|
||
"api_key": "",
|
||
}
|
||
return {"provider": p, "base_url": "", "api_key": ""}
|
||
|
||
|
||
def extract_json_detail(text: str) -> dict[str, str]:
|
||
detail = {
|
||
"content": "",
|
||
"reasoning_content": "",
|
||
"error_message": "",
|
||
"body_preview": text[:180].replace("\n", " "),
|
||
}
|
||
try:
|
||
data = json.loads(text)
|
||
except Exception:
|
||
return detail
|
||
|
||
if isinstance(data, dict):
|
||
choices = data.get("choices") or []
|
||
if isinstance(choices, list) and choices:
|
||
first = choices[0] if isinstance(choices[0], dict) else {}
|
||
msg = first.get("message") if isinstance(first, dict) else {}
|
||
if isinstance(msg, dict):
|
||
detail["content"] = str(msg.get("content") or "").strip()
|
||
detail["reasoning_content"] = str(msg.get("reasoning_content") or "").strip()
|
||
|
||
if not detail["error_message"]:
|
||
err = data.get("error")
|
||
if isinstance(err, dict):
|
||
detail["error_message"] = str(err.get("message") or "").strip()
|
||
elif err:
|
||
detail["error_message"] = str(err).strip()
|
||
|
||
if not detail["error_message"]:
|
||
if data.get("message"):
|
||
detail["error_message"] = str(data.get("message")).strip()
|
||
elif data.get("reason"):
|
||
detail["error_message"] = str(data.get("reason")).strip()
|
||
return detail
|
||
|
||
|
||
def probe_model(
|
||
level: str,
|
||
provider: str,
|
||
model: str,
|
||
timeout: float,
|
||
insecure: bool,
|
||
) -> dict[str, object]:
|
||
runtime = provider_runtime(provider)
|
||
base_url = runtime["base_url"].rstrip("/")
|
||
api_key = runtime["api_key"]
|
||
started = time.perf_counter()
|
||
|
||
if not base_url:
|
||
elapsed = time.perf_counter() - started
|
||
return {
|
||
"level": level,
|
||
"provider": provider,
|
||
"model": model,
|
||
"ok": False,
|
||
"elapsed": elapsed,
|
||
"reason": f"provider={provider} 缺少 base_url 配置",
|
||
"http_status": 0,
|
||
}
|
||
|
||
url = f"{base_url}/chat/completions"
|
||
payload = {
|
||
"model": model,
|
||
"messages": [
|
||
{"role": "system", "content": "你是连通性检测助手，只回复OK。"},
|
||
{"role": "user", "content": "ping"},
|
||
],
|
||
"temperature": 0,
|
||
"max_tokens": 16,
|
||
}
|
||
req = urllib.request.Request(
|
||
url=url,
|
||
data=json.dumps(payload).encode("utf-8"),
|
||
headers={"Content-Type": "application/json"},
|
||
method="POST",
|
||
)
|
||
if api_key:
|
||
req.add_header("Authorization", f"Bearer {api_key}")
|
||
|
||
status_code = 0
|
||
body = ""
|
||
error_text = ""
|
||
ssl_context = ssl._create_unverified_context() if insecure else None
|
||
try:
|
||
with urllib.request.urlopen(req, timeout=timeout, context=ssl_context) as resp:
|
||
status_code = int(resp.getcode() or 0)
|
||
body = resp.read().decode("utf-8", errors="replace")
|
||
except urllib.error.HTTPError as exc:
|
||
status_code = int(exc.code or 0)
|
||
body = exc.read().decode("utf-8", errors="replace")
|
||
error_text = str(exc)
|
||
except Exception as exc: # pragma: no cover
|
||
error_text = str(exc)
|
||
|
||
elapsed = time.perf_counter() - started
|
||
detail = extract_json_detail(body)
|
||
content = detail["content"]
|
||
reasoning = detail["reasoning_content"]
|
||
has_text = bool(content)
|
||
has_reasoning = bool(reasoning)
|
||
http_ok = 200 <= status_code < 300
|
||
|
||
if http_ok and (has_text or has_reasoning):
|
||
note = "文本返回" if has_text else "仅推理内容"
|
||
return {
|
||
"level": level,
|
||
"provider": provider,
|
||
"model": model,
|
||
"ok": True,
|
||
"elapsed": elapsed,
|
||
"reason": note,
|
||
"http_status": status_code,
|
||
}
|
||
|
||
reason_parts: list[str] = []
|
||
if status_code:
|
||
reason_parts.append(f"HTTP {status_code}")
|
||
if detail["error_message"]:
|
||
reason_parts.append(str(detail["error_message"]))
|
||
elif error_text:
|
||
reason_parts.append(error_text)
|
||
elif http_ok:
|
||
reason_parts.append("响应为空（content/reasoning_content 均为空）")
|
||
else:
|
||
reason_parts.append(detail["body_preview"] or "请求失败")
|
||
|
||
return {
|
||
"level": level,
|
||
"provider": provider,
|
||
"model": model,
|
||
"ok": False,
|
||
"elapsed": elapsed,
|
||
"reason": " | ".join(reason_parts),
|
||
"http_status": status_code,
|
||
}
|
||
|
||
|
||
def render_result_line(item: dict[str, object]) -> str:
|
||
status_text = "✓ 正常" if bool(item["ok"]) else "✗ 失败"
|
||
elapsed_text = f"{float(item['elapsed']):.2f}s"
|
||
return (
|
||
f"[{item['level']}] {item['model']} ({item['provider']}) "
|
||
f"{status_text} ({elapsed_text})"
|
||
)
|
||
|
||
|
||
def main() -> int:
|
||
parser = argparse.ArgumentParser(description="PMFlow 模型连通性独立检测")
|
||
parser.add_argument(
|
||
"--env-file",
|
||
default=str(Path(__file__).resolve().parents[1] / ".env"),
|
||
help="环境变量文件路径（默认根目录 .env）",
|
||
)
|
||
parser.add_argument(
|
||
"--timeout",
|
||
type=float,
|
||
default=float(os.getenv("PMFLOW_MODEL_PROBE_TIMEOUT", "20")),
|
||
help="单次请求超时时间（秒）",
|
||
)
|
||
parser.add_argument(
|
||
"--insecure",
|
||
action="store_true",
|
||
default=str(os.getenv("PMFLOW_MODEL_PROBE_INSECURE", "1")).strip().lower() in {"1", "true", "yes"},
|
||
help="关闭 TLS 证书校验（默认开启，可用 PMFLOW_MODEL_PROBE_INSECURE=0 关闭）",
|
||
)
|
||
parser.add_argument(
|
||
"--strict-tls",
|
||
action="store_true",
|
||
default=False,
|
||
help="强制开启 TLS 证书校验（优先级高于 --insecure）",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
load_env_file(Path(args.env_file))
|
||
|
||
insecure = bool(args.insecure and not args.strict_tls)
|
||
|
||
low_provider = str(os.getenv("LLM_ROUTING_LOW_PROVIDER", os.getenv("LLM_PROVIDER", "glm")) or "").strip() or "glm"
|
||
low_model = str(os.getenv("LLM_ROUTING_LOW_MODEL", os.getenv("GLM_MODEL", "GLM-4.7")) or "").strip() or "GLM-4.7"
|
||
high_provider = str(os.getenv("LLM_ROUTING_HIGH_PROVIDER", os.getenv("LLM_PROVIDER", "openai")) or "").strip() or "openai"
|
||
high_model = str(os.getenv("LLM_ROUTING_HIGH_MODEL", os.getenv("OPENAI_MODEL", "GLM-5")) or "").strip() or "GLM-5"
|
||
|
||
targets: list[tuple[str, str, str]] = [("LOW", low_provider, low_model), ("HIGH", high_provider, high_model)]
|
||
deduped: list[tuple[str, str, str]] = []
|
||
seen: set[tuple[str, str]] = set()
|
||
for level, provider, model in targets:
|
||
key = (provider.lower(), model)
|
||
if key in seen:
|
||
continue
|
||
seen.add(key)
|
||
deduped.append((level, provider, model))
|
||
|
||
started_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
|
||
print("┌────────────────────────────────────────┐")
|
||
print("│ PMFlow 模型连通性独立检测 │")
|
||
print(f"│ 启动时间: {started_at} │")
|
||
print(f"│ TLS校验: {'关闭' if insecure else '开启'} │")
|
||
print("├────────────────────────────────────────┤")
|
||
|
||
results: list[dict[str, object]] = []
|
||
for level, provider, model in deduped:
|
||
result = probe_model(
|
||
level=level,
|
||
provider=provider,
|
||
model=model,
|
||
timeout=args.timeout,
|
||
insecure=insecure,
|
||
)
|
||
results.append(result)
|
||
print(f"│ {render_result_line(result)}")
|
||
print(f"│ 说明: {result['reason']}")
|
||
|
||
print("└────────────────────────────────────────┘")
|
||
|
||
failed = [item for item in results if not bool(item["ok"])]
|
||
if failed:
|
||
print(f"[结果] 检测完成：{len(failed)} 个模型异常。")
|
||
return 1
|
||
print("[结果] 检测完成：全部模型连通正常。")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|