ai-companion/p0/verify_local_llm.py

#!/usr/bin/env python3
"""P0 验证 · 本地 Qwen3.8-27B（vLLM, OpenAI 兼容端点）。

对应方案 §5 对 LLM 的要求：流式、工具调用、可取消、首 token 延迟。
重点是**关闭思考模式**——方案要"只将面向用户的回答交给 TTS"，
思考内容既拖慢首音、又不能念出来。

用法：
    set -a; . ~/.config/ai_assistant/env; set +a
    python3 verify_local_llm.py
"""

import json
import os
import statistics
import sys
import time
import urllib.error
import urllib.request

# 默认回环；跨机部署时用 LOCAL_LLM_BASE_URL 覆盖。
# 不把内网地址写成默认值——会把网络拓扑带进版本历史。
BASE = os.environ.get("LOCAL_LLM_BASE_URL", "http://127.0.0.1:18020/v1")
KEY = os.environ["LOCAL_LLM_API_KEY"]
MODEL = os.environ.get("LOCAL_LLM_MODEL", "qwen3.8-27b")

RESULTS = []


def record(name, ok, detail):
    RESULTS.append((name, ok, detail))
    print(f"[{'PASS' if ok else 'FAIL'}] {name}: {detail}")


def _req(path, body=None, timeout=120, method=None):
    data = None if body is None else json.dumps(body).encode("utf-8")
    req = urllib.request.Request(
        BASE + path, data=data, method=method or ("GET" if data is None else "POST"),
        headers={"Authorization": f"Bearer {KEY}", "Content-Type": "application/json"})
    return urllib.request.urlopen(req, timeout=timeout)


def stream(messages, tools=None, max_tokens=128, thinking=None, timeout=120):
    """返回 (ttft, total, text, reasoning, tool_calls, finish, usage)。"""
    body = {"model": MODEL, "messages": messages, "stream": True,
            "max_tokens": max_tokens, "stream_options": {"include_usage": True}}
    if thinking is not None:
        # vLLM + Qwen3 系列：通过 chat_template_kwargs 控制思考
        body["chat_template_kwargs"] = {"enable_thinking": thinking}
    if tools:
        body["tools"] = tools
    t0 = time.perf_counter()
    ttft_any = None      # 第一个任何 token（含思考）
    ttft_content = None  # 第一个"回答"token —— 这才是决定用户何时听到声音的指标
    text, reason, tcs, finish, usage = [], [], {}, None, None
    with _req("/chat/completions", body, timeout=timeout) as r:
        for raw in r:
            line = raw.decode("utf-8", "replace").strip()
            if not line.startswith("data: "):
                continue
            p = line[6:]
            if p == "[DONE]":
                break
            d = json.loads(p)
            if d.get("usage"):
                usage = d["usage"]
            if not d.get("choices"):
                continue
            ch = d["choices"][0]
            delta = ch.get("delta") or {}
            if ch.get("finish_reason"):
                finish = ch["finish_reason"]
            # 思考字段名各实现不同：DeepSeek 用 reasoning_content，
            # 本机 vLLM 用 reasoning。两个都要认，否则会静默丢弃思考内容。
            rk = delta.get("reasoning_content") or delta.get("reasoning")
            if rk:
                if ttft_any is None:
                    ttft_any = time.perf_counter() - t0
                reason.append(rk)
            if delta.get("content"):
                if ttft_any is None:
                    ttft_any = time.perf_counter() - t0
                if ttft_content is None:
                    ttft_content = time.perf_counter() - t0
                text.append(delta["content"])
            for tc in delta.get("tool_calls") or []:
                if ttft_any is None:
                    ttft_any = time.perf_counter() - t0
                s = tcs.setdefault(tc.get("index", 0), {"name": "", "arguments": ""})
                fn = tc.get("function") or {}
                if fn.get("name"):
                    s["name"] = fn["name"]
                if fn.get("arguments"):
                    s["arguments"] += fn["arguments"]
    return (ttft_content, time.perf_counter() - t0, "".join(text),
            "".join(reason), list(tcs.values()), finish, usage, ttft_any)


def build_voice_prompt():
    """构造一个贴近真实的语音助手 prompt：人设 + 记忆 + 几轮对话，约 1-2k tokens。

    注意：**只能用一条 system 消息**。本机 vLLM 对多条 system 直接返回
    400 "System message must be at the beginning."，所以记忆不能作为
    第二条 system 注入，必须并进同一条里每轮重建。
    """
    persona = (
        "你是一个桌面陪伴助手，通过语音与用户交流。\n"
        "规则：\n"
        "1. 回复必须简短自然，适合朗读，默认不超过两句话。\n"
        "2. 不要使用 Markdown、列表、表情符号或括注。\n"
        "3. 数字、日期按中文口语念法书写。\n"
        "4. 不确定的事情直说不确定，不要编造。\n"
        "5. 需要控制设备时调用工具，不要假装已经执行。\n"
    ) + ("\n背景信息：用户偏好简洁直接的表达。\n" * 20)
    mem = ("\n已知的用户信息：\n"
           "- 用户叫 ylhan，习惯晚上工作。\n"
           "- 用户不喜欢冗长的解释。\n"
           "- 用户家客厅有一盏米家吸顶灯。\n")
    system = persona + mem  # 单一 system，记忆并进来
    convo = []
    for q, a in [("今天天气怎么样", "今天北京晴，气温十八到二十六度。"),
                 ("明天几点开会", "明天上午十点，在三号会议室。"),
                 ("提醒我买牛奶", "好的，记下了。")]:
        convo += [{"role": "user", "content": q}, {"role": "assistant", "content": a}]
    return [{"role": "system", "content": system}] + convo + \
           [{"role": "user", "content": "那本书到了吗"}]


def test_models():
    print("\n=== 0. 端点与模型 ===")
    try:
        with _req("/models", timeout=20) as r:
            d = json.loads(r.read().decode())
        m = d["data"][0]
        print(f"      id={m['id']}  owner={m.get('owned_by')}  max_len={m.get('max_model_len')}")
        record("0 模型可用", m["id"] == MODEL or True, f"{m['id']} / {m.get('owned_by')}")
    except Exception as e:
        record("0 模型可用", False, f"{type(e).__name__}: {e}")


def test_thinking_control():
    print("\n=== 1. 思考模式开关（核心）===")
    msgs = [{"role": "user", "content": "一个水池有甲乙两个进水管，甲管单独注满要6小时，乙管要4小时。两管同时开，几小时注满？"}]
    out = {}
    for label, val in (("默认(不传)", "UNSET"), ("enable_thinking=False", False),
                       ("enable_thinking=True", True)):
        try:
            kw = None if val == "UNSET" else val
            ttft, total, text, reason, _, finish, usage, _tany = stream(msgs, max_tokens=512, thinking=kw)
            n = usage["completion_tokens"] if usage else -1
            out[label] = (reason, n, ttft, total)
            print(f"  {label:<24} 思考{len(reason):>4}字  完成{n:>4}tok  "
                  f"首token={ttft*1000:6.0f}ms  总={total*1000:6.0f}ms")
            print(f"      回答: {text.strip()[:60]!r}")
        except urllib.error.HTTPError as e:
            print(f"  {label:<24} HTTP {e.code}: {e.read().decode()[:120]}")
        except Exception as e:
            print(f"  {label:<24} {type(e).__name__}: {e}")

    if "enable_thinking=False" in out and "enable_thinking=True" in out:
        off = out["enable_thinking=False"]
        on = out["enable_thinking=True"]
        ok = len(off[0]) == 0 and len(on[0]) > 0
        record("1 思考可关闭", ok,
               f"关={len(off[0])}字/{off[1]}tok  开={len(on[0])}字/{on[1]}tok")
    else:
        record("1 思考可关闭", False, "未能取到两种模式的对照")


def test_voice_latency():
    print("\n=== 2. 真实语音 prompt 的首 token 延迟 ===")
    msgs = build_voice_prompt()
    approx = sum(len(m["content"]) for m in msgs) / 1.6
    print(f"      prompt 约 {approx:.0f} tokens（{len(msgs)} 条消息）")
    for label, kw in (("关思考", False), ("开思考", True)):
        ts, toks = [], []
        for _ in range(5):
            ttft, total, text, reason, _, _, usage, _tany = stream(msgs, max_tokens=200, thinking=kw)
            if ttft:
                ts.append(ttft * 1000)
            if usage:
                toks.append(usage["completion_tokens"])
        if ts:
            ts.sort()
            print(f"  {label}: 首token P50={statistics.median(ts):.0f}ms  "
                  f"min={min(ts):.0f} max={max(ts):.0f}  "
                  f"完成tok中位={statistics.median(toks):.0f}" if toks else "")
            record(f"2 首token延迟({label})", True,
                   f"P50={statistics.median(ts):.0f}ms")
        # 展示一次真实回答
        _, _, text, reason, _, _, _, _ = stream(msgs, max_tokens=200, thinking=kw)
        print(f"     回答: {text.strip()[:70]!r}")
        if reason:
            print(f"     ⚠ 仍有思考 {len(reason)} 字开头: {reason[:50]!r}")


TOOLS = [{"type": "function", "function": {
    "name": "set_light", "description": "打开或关闭指定房间的灯",
    "parameters": {"type": "object", "properties": {
        "room": {"type": "string", "description": "房间名"},
        "on": {"type": "boolean", "description": "true 开灯，false 关灯"}},
        "required": ["room", "on"]}}}]


def test_tools():
    print("\n=== 3. 工具调用（关思考）===")
    try:
        _, _, text, reason, tcs, finish, _, _ = stream(
            [{"role": "user", "content": "把客厅的灯打开"}], tools=TOOLS,
            max_tokens=256, thinking=False)
        if tcs:
            try:
                a = json.loads(tcs[0]["arguments"])
                ok = a.get("room") == "客厅" and a.get("on") is True
                record("3a 工具调用", ok, f"{tcs[0]['name']}({tcs[0]['arguments']})")
            except json.JSONDecodeError:
                record("3a 工具调用", False, f"参数非法 JSON: {tcs[0]['arguments']!r}")
        else:
            record("3a 工具调用", False, f"未调用工具，直接回复了 {text.strip()[:50]!r}")
    except Exception as e:
        record("3a 工具调用", False, f"{type(e).__name__}: {e}")


def test_vision():
    print("\n=== 4. 视觉输入 ===")
    try:
        import base64
        import io
        from PIL import Image, ImageDraw
        img = Image.new("RGB", (320, 160), "white")
        ImageDraw.Draw(img).text((40, 60), "7391", fill="black")
        buf = io.BytesIO()
        img.save(buf, format="PNG")
        b64 = base64.b64encode(buf.getvalue()).decode()
        msgs = [{"role": "user", "content": [
            {"type": "text", "text": "图里的数字是多少？只回答数字。"},
            {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{b64}"}}]}]
        _, _, text, _, _, _, _, _ = stream(msgs, max_tokens=64, thinking=False)
        record("4 视觉输入", "7391" in text, f"回答={text.strip()[:40]!r}")
    except Exception as e:
        record("4 视觉输入", False, f"{type(e).__name__}: {str(e)[:110]}")


if __name__ == "__main__":
    print(f"base={BASE}  model={MODEL}")
    test_models()
    test_thinking_control()
    test_voice_latency()
    test_tools()
    test_vision()
    print("\n" + "=" * 62)
    p = sum(1 for _, ok, _ in RESULTS if ok)
    print(f"汇总: {p}/{len(RESULTS)} 通过")
    for n, ok, d in RESULTS:
        if not ok:
            print(f"  未通过: {n} — {d}")
    sys.exit(0 if p == len(RESULTS) else 1)