model-watchdog/probes.py

"""模型可用性探测：按 API 端点分类，每类一个探测函数。

统一返回结果 dict：
    {model, endpoint, label, ok, code, message, raw, latency_ms}

判定原则：
- HTTP 200 且响应含该端点应有的字段 -> 可用
- 其他情况（非 200 / 缺字段 / 连接异常 / WS task-failed）-> 加入告警列表，
  raw 保留完整响应体，供告警正文里的"完整响应"使用。
"""
import json
import logging
import re
import time
import uuid
from concurrent.futures import ThreadPoolExecutor

import requests

log = logging.getLogger("probe")

# websocket-client 在正常关闭连接时也会打 ERROR 日志，屏蔽掉避免污染告警日志
logging.getLogger("websocket").setLevel(logging.CRITICAL)

RAW_LIMIT = 1500          # 完整响应在告警里的最大保留长度
MAX_WORKERS = 5           # 并发探测数

# validate 模式下，命中该模式的 400 报错说明"模型在线，只是没给 prompt"
_PROMPT_ERR_RE = re.compile(r"prompt", re.IGNORECASE)


# ========== 对外入口 ==========

def probe_all(targets, api_key):
    """并发探测 targets = [(模型名, 分组配置), ...]，返回结果列表（保持入参顺序）。"""
    if not targets:
        return []
    with ThreadPoolExecutor(max_workers=min(MAX_WORKERS, len(targets))) as pool:
        results = list(pool.map(lambda t: probe_one(t[0], t[1], api_key), targets))
    for r in results:
        if r["ok"]:
            log.info("[探测] %s (%s) 正常, %d ms", r["model"], r["label"], r["latency_ms"])
        else:
            log.warning("[探测] %s (%s) 异常: %s %s", r["model"], r["label"], r["code"], r["message"])
    return results


def probe_one(model, group, api_key):
    endpoint = group["endpoint"]
    fn = _PROBES.get(endpoint)
    result = {
        "model": model,
        "endpoint": endpoint,
        "label": group.get("label", endpoint),
        "ok": False,
        "code": "",
        "message": "",
        "raw": "",
        "latency_ms": 0,
    }
    if fn is None:
        result["code"] = "ConfigError"
        result["message"] = f"未实现的端点类型: {endpoint}"
        return result

    start = time.monotonic()
    try:
        ok, code, message, raw = fn(model, group, api_key)
    except Exception as e:  # noqa: BLE001 - 网络/超时/协议异常统一转为告警
        ok, code, message, raw = False, type(e).__name__, str(e), ""
    result["latency_ms"] = int((time.monotonic() - start) * 1000)
    result.update(ok=ok, code=str(code or ""), message=str(message or ""), raw=_clip(raw))
    return result


# ========== 各端点探测实现 ==========

def _probe_chat(model, g, api_key):
    """OpenAI 兼容模式对话：只让模型回一个 ok，输出上限 max_tokens。"""
    body = {
        "model": model,
        "messages": [{"role": "user", "content": g["prompt"]}],
        "max_tokens": g["max_tokens"],
        "temperature": 0,
        "stream": False,
    }
    body.update(g.get("extra_params") or {})
    resp = requests.post(
        g["url"], headers=_headers(api_key), json=body, timeout=g["timeout_seconds"]
    )
    return _judge_http(resp, ("choices",))


def _probe_asr(model, g, api_key):
    """录音文件识别：异步提交任务，拿到 task_id 即认为模型可用（不轮询结果）。"""
    body = {
        "model": model,
        "input": {"file_urls": [g["sample_audio_url"]]},
        "parameters": {"channel_id": [0]},
    }
    resp = requests.post(
        g["url"],
        headers=_headers(api_key, async_task=True),
        json=body,
        timeout=g["timeout_seconds"],
    )
    return _judge_http(resp, ("output", "task_id"))


def _probe_text2image(model, g, api_key):
    """文生图：异步提交任务。

    probe_mode=validate（默认）：故意不传 prompt。服务端先解析模型再校验参数，
        因此"模型不存在"与"缺少 prompt"可区分开，且不会真的出图、不产生费用。
    probe_mode=submit：真实提交生成任务（会计费）。
    """
    validate = g.get("probe_mode", "validate") == "validate"
    body = {
        "model": model,
        "input": {} if validate else {"prompt": g["prompt"]},
        "parameters": {"n": 1, "size": g["size"]},
    }
    resp = requests.post(
        g["url"],
        headers=_headers(api_key, async_task=True),
        json=body,
        timeout=g["timeout_seconds"],
    )
    return _judge_async(resp, validate)


def _probe_text2video(model, g, api_key):
    """文生视频：异步提交任务，模式含义同文生图。"""
    validate = g.get("probe_mode", "validate") == "validate"
    body = {
        "model": model,
        "input": {} if validate else {"prompt": g["prompt"]},
        "parameters": {"size": g["size"], "duration": g["duration"]},
    }
    resp = requests.post(
        g["url"],
        headers=_headers(api_key, async_task=True),
        json=body,
        timeout=g["timeout_seconds"],
    )
    return _judge_async(resp, validate)


def _probe_tts(model, g, api_key):
    """语音合成（WebSocket）：建连 -> run-task -> 等 task-started。

    不发送待合成文本、不发 finish-task，收到 task-started 立即断开，
    因此只验证模型可用性，几乎不产生合成用量。
    """
    try:
        import websocket  # websocket-client
    except ImportError:
        return False, "DependencyMissing", "缺少依赖 websocket-client，无法探测 wss 端点", ""

    timeout = g["timeout_seconds"]
    task_id = uuid.uuid4().hex
    run_task = {
        "header": {"action": "run-task", "task_id": task_id, "streaming": "duplex"},
        "payload": {
            "task_group": "audio",
            "task": "tts",
            "function": "SpeechSynthesizer",
            "model": model,
            "parameters": {
                "text_type": "PlainText",
                "voice": g["voice"],
                "format": g["audio_format"],
                "sample_rate": g["sample_rate"],
            },
            "input": {},
        },
    }

    ws = websocket.create_connection(
        g["url"],
        header=[f"Authorization: bearer {api_key}", "X-DashScope-DataInspection: enable"],
        timeout=timeout,
    )
    try:
        ws.send(json.dumps(run_task))
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            frame = ws.recv()
            if isinstance(frame, (bytes, bytearray)):   # 音频帧，探测阶段不应出现
                continue
            try:
                msg = json.loads(frame)
            except ValueError:
                continue
            event = (msg.get("header") or {}).get("event", "")
            if event == "task-started":
                return True, "", "", frame
            if event in ("task-failed", "task-finished"):
                header = msg.get("header") or {}
                code = header.get("error_code") or event
                message = header.get("error_message") or f"服务端返回 {event}"
                return False, code, message, frame
        return False, "Timeout", f"{timeout}s 内未收到 task-started 事件", ""
    finally:
        try:
            ws.close()
        except Exception:  # noqa: BLE001
            pass


_PROBES = {
    "chat": _probe_chat,
    "tts": _probe_tts,
    "asr": _probe_asr,
    "text2image": _probe_text2image,
    "text2video": _probe_text2video,
}


# ========== 工具函数 ==========

def _headers(api_key, async_task=False):
    h = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
    if async_task:
        h["X-DashScope-Async"] = "enable"
    return h


def _judge_async(resp, validate):
    """判定异步生成端点（文生图/文生视频）的提交响应。

    validate 模式下，仅因缺少 prompt 被拒说明模型解析已通过 -> 判为可用；
    "Model not exist" / 鉴权失败 / 5xx 等一律判为不可用。
    """
    if not validate:
        return _judge_http(resp, ("output", "task_id"))

    raw = resp.text or ""
    try:
        data = resp.json()
    except ValueError:
        data = None

    if resp.status_code == 200 and _has_path(data, ("output", "task_id")):
        return True, "", "", raw

    code, message = _extract_error(data)
    if resp.status_code == 400 and _PROMPT_ERR_RE.search(message or ""):
        return True, "", "", raw
    return False, code or resp.status_code, message or f"HTTP {resp.status_code}", raw


def _judge_http(resp, required_path):
    """判定 HTTP 响应。required_path 为响应中必须存在的字段路径。"""
    raw = resp.text or ""
    try:
        data = resp.json()
    except ValueError:
        data = None

    if resp.status_code != 200:
        code, message = _extract_error(data)
        return False, code or resp.status_code, message or f"HTTP {resp.status_code}", raw

    # DashScope 原生端点即使 HTTP 200 也可能在 body 里带 code 表示失败
    if isinstance(data, dict) and data.get("code"):
        code, message = _extract_error(data)
        return False, code, message, raw

    if not _has_path(data, required_path):
        return (
            False,
            "BadResponse",
            "HTTP 200 但响应缺少字段 " + ".".join(required_path),
            raw,
        )
    return True, "", "", raw


def _extract_error(data):
    """兼容两种错误结构：OpenAI 兼容模式 {"error": {...}} 与原生 {"code","message"}。"""
    if not isinstance(data, dict):
        return "", ""
    err = data.get("error")
    if isinstance(err, dict):
        return str(err.get("code") or err.get("type") or ""), str(err.get("message") or "")
    if isinstance(err, str) and err:
        return "", err
    return str(data.get("code") or ""), str(data.get("message") or "")


def _has_path(data, path):
    cur = data
    for seg in path:
        if not isinstance(cur, dict) or seg not in cur:
            return False
        cur = cur[seg]
    return cur not in (None, "", [], {})


def _clip(text):
    text = (text or "").strip()
    if len(text) <= RAW_LIMIT:
        return text
    return text[:RAW_LIMIT] + f"…（已截断，原长 {len(text)}）"