vent_deep_agent/tests/test_fast_stream.py

# -*- coding: utf-8 -*-
"""
Fast 模式流式输出测试 Demo

模拟 /api/interpret/click/tun?mode=fast 的完整 LLM 调用链路：
  1. 读取 .env 配置 → 获取 deepseek-v4-flash 模型
  2. 构造与 Fast 模式完全一致的 system + user prompt
  3. 流式调用模型，逐 token 打印并统计延迟

  用法：
  cd D:\\workspace\\py\\vent_deep_agent_1.4
  python tests/test_fast_stream.py

  也可以带参数指定巷道 ID（会走 MCP 真实查询）：
  python tests/test_fast_stream.py --tun-id 581 --tun-name "二水平中部辅运大巷"
"""

import argparse
import asyncio
import json
import os
import sys
import time
from pathlib import Path

# 确保项目根目录在 sys.path 中
sys.path.insert(0, str(Path(__file__).parent.parent))

from dotenv import dotenv_values
from openai import AsyncOpenAI


# ── 与 routes.py 完全一致的 Fast 模式系统提示词 ──
FAST_TUN_SYSTEM_PROMPT = """你是一名煤矿通风安全专家。请根据提供的巷道监测数据，直接生成标准化结构化解读报告。

## 报告格式（严格按此模板输出，不要添加额外标题或代码块）

**{巷道名称}**
- 当前风速：{数值} m/s
- 规程允许范围：{min} ~ {max} m/s
- 当前风量：{数值} m³/min{，计划风量：{数值} m³/min}
- 瓦斯浓度：{数值}%（限值 {限值}%）
- CO浓度：{数值} ppm（限值 {限值} ppm）
- 温度：{数值}℃（限值 {限值}℃）
- 设备状态：{状态}
- 结论：{正常/重大风险/高风险/中风险}。{详细异常分析}
- 依据：《煤矿安全规程》相关条款
- 建议：{针对性通风整改措施}

## 风险分级标准
- 风速偏低：实测＜0.5倍下限→重大；0.5倍≤实测＜0.8倍下限→高；0.8倍≤实测＜下限→中
- 风速偏高：实测≥1.8倍上限→重大；1.4倍≤实测＜1.8倍上限→高；上限≤实测＜1.4倍上限→中
- 风量超配：实测≥1.8倍计划→重大；1.5倍≤实测＜1.8倍计划→高；计划＜实测＜1.5倍计划→中
- 风量不足：实测＜0.75倍计划→重大；0.75倍≤实测＜0.9倍计划→高；0.9倍≤实测＜计划→中

## 巷道类型编码对照
- 0-回采工作面；1-掘进工作面；2-辅运巷；3-主运巷；4-硐室；5-联络巷；6-进风井；7-回风井；8-专用回风巷

## 规则
- 所有分析必须基于提供的数据，绝不编造
- 全程使用简体中文输出
- 严格按照模板格式输出，不要添加额外标题、总结、代码块、JSON
- 禁止暴露巷道ID，统一用巷道名称代替
- 数据缺失项标注"无数据"并说明影响
"""

# ── 模拟巷道数据（与 MCP 返回格式一致） ──
MOCK_TUN_DATA = {
    "tunId": "581",
    "tunnelName": "二水平中部辅运大巷",
    "needAirVolume": 1200,
    "usingType": 2,
    "usingTypeName": "辅运巷",
    "permissibleMin": 0.25,
    "permissibleMax": 6.0,
    "permissibleVelocity": {
        "fmin": 0.25,
        "fmax": 6.0,
    },
    "devices": [
        {
            "sourceType": "wind",
            "sourceName": "二水平中部辅运大巷测风站",
            "installPos": "二水平中部辅运大巷中部",
            "airVolume": 1150.5,
            "windSpeed": 1.85,
            "warnFlag": 0,
            "netStatus": 1,
            "deviceStatusName": "在线",
            "readTime": "2026-07-15 10:30:00",
            "m3": 1150.5,
            "va": 1.85,
            "alarmDescription": "",
            "alarmDescriptions": [],
        }
    ],
    # 模拟瓦斯传感器
    "gasSensor": {
        "concentration": 0.12,
        "limit": 0.8,
        "unit": "%",
    },
    # 模拟 CO 传感器
    "coSensor": {
        "concentration": 5,
        "limit": 24,
        "unit": "ppm",
    },
    # 模拟温度传感器
    "tempSensor": {
        "value": 24.5,
        "limit": 30,
        "unit": "℃",
    },
}


def load_env():
    """读取 .env 配置（与 _get_chat_model 逻辑一致）"""
    env_file = Path(__file__).parent.parent / ".env"
    _cfg = dotenv_values(str(env_file))

    model_str = _cfg.get("SUMMARY_MODEL", "") or _cfg.get("DEEPAGENT_MODEL", "deepseek-v4-flash")
    base_url = _cfg.get("OPENAI_BASE_URL", "https://api.deepseek.com").strip().strip('"').strip("'")
    api_key = _cfg.get("OPENAI_API_KEY", "").strip().strip('"').strip("'")

    return model_str, base_url, api_key


def build_prompt(tun_name: str, data: dict) -> list[dict]:
    """构造与 routes.py _fast_tun_sse_generator 完全一致的 messages"""
    data_str = json.dumps(data, ensure_ascii=False, indent=2)
    return [
        {"role": "system", "content": FAST_TUN_SYSTEM_PROMPT},
        {"role": "user", "content": f"请根据以下巷道监测数据生成标准化解读报告：\n\n巷道名称：{tun_name}\n\n数据：\n{data_str}"},
    ]


async def stream_model(client: AsyncOpenAI, model: str, messages: list[dict]):
    """流式调用模型，逐 token 输出并计时"""
    t0 = time.perf_counter()
    first_token_time = None
    token_count = 0
    full_text = ""

    print("─" * 60)
    print(f"模型: {model}")
    print(f"System prompt 长度: {len(messages[0]['content'])} 字符")
    print(f"User prompt 长度:   {len(messages[1]['content'])} 字符")
    print("─" * 60)
    print("\n[LLM 流式输出开始]\n")

    stream = await client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0,
        stream=True,
        stream_options={"include_usage": True},
    )

    async for chunk in stream:
        if chunk.choices and chunk.choices[0].delta.content:
            content = chunk.choices[0].delta.content
            if first_token_time is None:
                first_token_time = time.perf_counter()
                ttft = first_token_time - t0
                print(f"\n{'─' * 40}")
                print(f"⏱ 首 token 延迟 (TTFT): {ttft:.2f}s")
                print(f"{'─' * 40}\n")

            token_count += 1
            full_text += content
            # 实时逐 token 打印（不换行）
            print(content, end="", flush=True)

    t_total = time.perf_counter() - t0
    ttft = (first_token_time - t0) if first_token_time else t_total

    print("\n")
    print("─" * 60)
    print(f"📊 统计:")
    print(f"   首 token 延迟: {ttft:.2f}s")
    print(f"   总耗时:        {t_total:.2f}s")
    print(f"   总 token 数:   {token_count}")
    print(f"   平均速度:      {token_count / max(t_total - ttft, 0.01):.1f} tok/s")
    print(f"   输出长度:      {len(full_text)} 字符")
    print("─" * 60)

    return full_text


async def main():
    parser = argparse.ArgumentParser(description="Fast 模式流式输出测试")
    parser.add_argument("--tun-id", default="581", help="巷道 ID（默认 581，使用 mock 数据）")
    parser.add_argument("--tun-name", default="二水平中部辅运大巷", help="巷道名称")
    parser.add_argument("--real-data", action="store_true", help="通过 MCP 获取真实数据（需要 MCP 服务可用）")
    parser.add_argument("--model", default=None, help="覆盖模型名称（默认从 .env 读取 SUMMARY_MODEL）")
    args = parser.parse_args()

    # 1. 加载环境配置
    model_str, base_url, api_key = load_env()
    if args.model:
        model_str = args.model

    print(f"🔧 API Base: {base_url}")
    print(f"🔧 模型:     {model_str}")

    if not api_key:
        print("❌ 未找到 OPENAI_API_KEY，请检查 .env 文件")
        return

    # 2. 获取数据
    if args.real_data:
        print("🌐 通过 MCP 获取真实数据...")
        t_mcp_start = time.perf_counter()
        from tools.vent_tools import query_tun_data_by_id
        raw = await query_tun_data_by_id(args.tun_id)
        data = json.loads(raw) if isinstance(raw, str) else raw
        t_mcp = time.perf_counter() - t_mcp_start
        print(f"   MCP 耗时: {t_mcp:.2f}s")
    else:
        print("📦 使用 Mock 数据（跳过 MCP 调用）")
        data = MOCK_TUN_DATA

    # 3. 构造提示词
    messages = build_prompt(args.tun_name, data)

    # 4. 初始化 OpenAI 客户端
    client = AsyncOpenAI(
        api_key=api_key,
        base_url=base_url + "/v1" if not base_url.endswith("/v1") else base_url,
        timeout=180.0,
        max_retries=2,
    )

    # 5. 流式调用
    full_text = await stream_model(client, model_str, messages)

    # 6. 可选：保存结果
    output_file = Path(__file__).parent / "fast_test_output.txt"
    output_file.write_text(full_text, encoding="utf-8")
    print(f"\n💾 完整输出已保存到: {output_file}")


if __name__ == "__main__":
    asyncio.run(main())