ai_site/scripts/probe_dashscope.py

"""Probe DashScope / 通义千问 connectivity and latency."""
from __future__ import annotations

import base64
import os
import socket
import ssl
import sys
import time
from pathlib import Path

import httpx

ROOT = Path(__file__).resolve().parents[1]


def load_env() -> None:
    for envp in (ROOT / ".env", ROOT / "ai-service" / ".env"):
        if not envp.exists():
            continue
        for line in envp.read_text(encoding="utf-8", errors="replace").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            k, v = line.split("=", 1)
            k, v = k.strip(), v.strip().strip('"').strip("'")
            if k and k not in os.environ:
                os.environ[k] = v


def post(base: str, key: str, payload: dict, timeout: float, label: str) -> tuple[float, int | None]:
    t0 = time.time()
    try:
        with httpx.Client(timeout=timeout) as c:
            r = c.post(
                f"{base}/chat/completions",
                headers={"Authorization": f"Bearer {key}", "Content-Type": "application/json"},
                json=payload,
            )
        dt = time.time() - t0
        print(f"[{label}] status={r.status_code} time={dt:.1f}s bytes_in={len(r.content)}")
        if r.status_code >= 400:
            print(f"[{label}] body {r.text[:500].replace(chr(10), ' ')}")
        else:
            j = r.json()
            txt = j["choices"][0]["message"]["content"]
            if isinstance(txt, list):
                txt = "".join((x.get("text") if isinstance(x, dict) else str(x)) for x in txt)
            print(f"[{label}] ok content_len={len(str(txt))} preview={str(txt)[:100]!r}")
        return dt, r.status_code
    except Exception as e:
        dt = time.time() - t0
        print(f"[{label}] FAIL after {dt:.1f}s: {type(e).__name__}: {e}")
        return dt, None


def main() -> int:
    load_env()
    key = (os.getenv("DASHSCOPE_API_KEY") or "").strip()
    base = (os.getenv("DASHSCOPE_BASE_URL") or "https://dashscope.aliyuncs.com/compatible-mode/v1").rstrip("/")
    model = (os.getenv("VISION_MODEL") or "qwen3.6-plus").strip()
    print(f"base={base}")
    print(f"model={model}")
    print(f"key_set={bool(key)} key_len={len(key)}")
    if not key:
        print("NO DASHSCOPE_API_KEY")
        return 1

    host = "dashscope.aliyuncs.com"
    t0 = time.time()
    try:
        ips = socket.getaddrinfo(host, 443, type=socket.SOCK_STREAM)
        print(f"DNS ok {time.time() - t0:.2f}s -> {sorted({x[4][0] for x in ips})[:4]}")
    except Exception as e:
        print(f"DNS FAIL {type(e).__name__}: {e}")

    t0 = time.time()
    try:
        ctx = ssl.create_default_context()
        with socket.create_connection((host, 443), timeout=10) as sock:
            with ctx.wrap_socket(sock, server_hostname=host) as ssock:
                print(f"TLS ok {time.time() - t0:.2f}s protocol={ssock.version()}")
    except Exception as e:
        print(f"TLS FAIL {type(e).__name__}: {e}")

    # 1) text-only
    post(
        base,
        key,
        {
            "model": model,
            "temperature": 0,
            "messages": [{"role": "user", "content": "只回复：通义可达"}],
            "max_tokens": 32,
        },
        30.0,
        "text-only",
    )

    # 2) tiny image
    png_b64 = (
        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
    )
    post(
        base,
        key,
        {
            "model": model,
            "temperature": 0,
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": "这张图是什么颜色？一句话。"},
                        {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{png_b64}"}},
                    ],
                }
            ],
            "max_tokens": 64,
        },
        60.0,
        "tiny-image",
    )

    # 3) real refs
    refs = list((ROOT / "test" / "refs").glob("*.png"))[:2]
    if not refs:
        refs = list((ROOT / "test").rglob("board_*.png"))[:2]
    print("refs", [f"{p.name}:{p.stat().st_size}" for p in refs])
    if refs:
        content: list[dict] = [{"type": "text", "text": "用一句话描述这些界面截图里最显眼的标题。"}]
        total = 0
        for p in refs:
            raw = p.read_bytes()
            total += len(raw)
            b64 = base64.b64encode(raw).decode("ascii")
            content.append({"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}})
        print(f"ref payload images_bytes={total} b64_approx={total * 4 // 3}")
        post(
            base,
            key,
            {
                "model": model,
                "temperature": 0,
                "messages": [{"role": "user", "content": content}],
                "max_tokens": 128,
            },
            150.0,
            "real-refs",
        )

        # long prompt + 2 images like fidelity scoring (heavier)
        long_text = ("请逐区对照并打分。" + "区" * 800)[:2000]
        content2 = [{"type": "text", "text": long_text}] + content[1:]
        post(
            base,
            key,
            {
                "model": model,
                "temperature": 0,
                "messages": [{"role": "user", "content": content2}],
                "max_tokens": 512,
            },
            150.0,
            "fidelity-like",
        )

    # 4) alternate model text
    if model != "qwen-vl-plus":
        post(
            base,
            key,
            {
                "model": "qwen-vl-plus",
                "temperature": 0,
                "messages": [{"role": "user", "content": "只回复：vl-plus可达"}],
                "max_tokens": 32,
            },
            30.0,
            "alt-qwen-vl-plus-text",
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())