vent_deep_agent/api/fast_generators.py

# -*- coding: utf-8 -*-
"""
Fast 模式 SSE 流式生成器 —— 跳过 Agent 与 Skills，直接查询数据 + LLM 生成报告
"""

import json
import time
import asyncio
import traceback
from typing import AsyncGenerator

from api.prompts import FAST_TUN_SYSTEM_PROMPT, FAST_DEVICE_SYSTEM_PROMPT
from api.chat_model import get_chat_model
from api.sse_core import _clean_response
from tools.vent_tools import query_tun_data_by_id, query_device_data_by_id
from db.chat_store import save_message


async def fast_tun_sse_generator(
    tun_id: str,
    tun_name: str,
    thread_id: str,
    session_id: str,
    save_to_db: bool = True,
) -> AsyncGenerator[str, None]:
    """Fast 模式巷道解读 SSE 生成器。

    跳过 Agent 与 Skills，直接调用 query_tun_data_by_id 获取数据，
    将数据 + 报告模板发给 LLM，流式输出 SSE。

    Args:
        save_to_db: 是否将消息保存到数据库。点选解读等非对话场景应设为 False。
    """
    if save_to_db:
        save_message(session_id, "user", f"快速解读: {tun_name}")

    full_response = ""

    try:
        # 1. 立即通知前端，不阻塞（让用户第一时间看到反馈）
        yield f"data: {json.dumps({'type': 'thinking', 'source': 'fast', 'node': 'fetching_data'}, ensure_ascii=False)}\n\n"

        # 2. 获取缓存的模型实例（首次调用 ~1s，后续缓存命中 ~0s）
        model = get_chat_model()

        # 3. 直接调用数据查询工具
        data = await query_tun_data_by_id(tun_id)

        yield f"data: {json.dumps({'type': 'executing', 'source': 'fast', 'tools': ['query_tun_data_by_id']}, ensure_ascii=False)}\n\n"

        # 4. 构建消息并调用 LLM 流式生成报告
        messages = [
            {"role": "system", "content": FAST_TUN_SYSTEM_PROMPT},
            {"role": "user", "content": f"请根据以下巷道监测数据生成标准化解读报告，禁止给出任何建议：\n\n巷道名称：{tun_name}\n\n数据：\n{data}"},
        ]

        # 5. 通知前端 LLM 正在生成（消除"卡住"的感知）
        yield f"data: {json.dumps({'type': 'generating', 'source': 'fast'}, ensure_ascii=False)}\n\n"

        # 带心跳的超时迭代：避免 LLM TTFT 过长导致前端"卡住"
        first_token = False
        heartbeat_start = time.time()
        astream_iter = model.astream(messages)

        while True:
            try:
                if first_token:
                    chunk = await astream_iter.__anext__()
                else:
                    # 首 token 等待，3 秒超时发送心跳
                    chunk = await asyncio.wait_for(astream_iter.__anext__(), timeout=3.0)
                    if not first_token:
                        first_token = True
            except asyncio.TimeoutError:
                elapsed = int(time.time() - heartbeat_start)
                yield f"data: {json.dumps({'type': 'progress', 'message': f'模型正在生成报告（已等待 {elapsed} 秒）...'}, ensure_ascii=False)}\n\n"
                continue
            except StopAsyncIteration:
                break

            # 处理推理内容（DeepSeek-R1 等推理模型会先输出 reasoning_content）
            reasoning = ""
            if hasattr(chunk, "additional_kwargs") and isinstance(chunk.additional_kwargs, dict):
                reasoning = chunk.additional_kwargs.get("reasoning_content", "")
            if reasoning:
                yield f"data: {json.dumps({'type': 'reasoning', 'content': reasoning}, ensure_ascii=False)}\n\n"
                continue  # 推理内容不保存到 full_response

            content = chunk.content if hasattr(chunk, "content") else ""
            if content and isinstance(content, str):
                full_response += content
                yield f"data: {json.dumps({'type': 'token', 'source': 'fast', 'content': content}, ensure_ascii=False)}\n\n"

        # 6. 保存助手回复
        if save_to_db and full_response.strip():
            cleaned = _clean_response(full_response)
            save_message(session_id, "assistant", cleaned)

        yield f"data: {json.dumps({'type': 'done', 'thread_id': thread_id, 'session_id': session_id}, ensure_ascii=False)}\n\n"

    except Exception as e:
        traceback.print_exc()
        yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"


async def fast_device_sse_generator(
    device_id: str,
    device_name: str,
    device_type: str,
    thread_id: str,
    session_id: str,
    save_to_db: bool = True,
) -> AsyncGenerator[str, None]:
    """Fast 模式设备解读 SSE 生成器。

    跳过 Agent 与 Skills，直接调用 query_device_data_by_id 获取数据，
    将数据 + 报告模板发给 LLM，流式输出 SSE。

    Args:
        save_to_db: 是否将消息保存到数据库。点选解读等非对话场景应设为 False。
    """
    if save_to_db:
        save_message(session_id, "user", f"快速解读: {device_name}")

    full_response = ""

    try:
        # 1. 立即通知前端，不阻塞（让用户第一时间看到反馈）
        yield f"data: {json.dumps({'type': 'thinking', 'source': 'fast', 'node': 'fetching_data'}, ensure_ascii=False)}\n\n"

        # 2. 获取缓存的模型实例（首次调用 ~1s，后续缓存命中 ~0s）
        model = get_chat_model()

        # 3. 直接调用数据查询工具
        data = await query_device_data_by_id(device_id)

        yield f"data: {json.dumps({'type': 'executing', 'source': 'fast', 'tools': ['query_device_data_by_id']}, ensure_ascii=False)}\n\n"

        # 4. 构建消息并调用 LLM 流式生成报告
        messages = [
            {"role": "system", "content": FAST_DEVICE_SYSTEM_PROMPT},
            {"role": "user", "content": f"请根据以下设备监测数据生成标准化解读报告,禁止给出任何建议：\n\n设备名称：{device_name}\n设备类型：{device_type}\n\n数据：\n{data}"},
        ]

        # 5. 通知前端 LLM 正在生成（消除"卡住"的感知）
        yield f"data: {json.dumps({'type': 'generating', 'source': 'fast'}, ensure_ascii=False)}\n\n"

        # 带心跳的超时迭代：避免 LLM TTFT 过长导致前端"卡住"
        first_token = False
        heartbeat_start = time.time()
        astream_iter = model.astream(messages)

        while True:
            try:
                if first_token:
                    chunk = await astream_iter.__anext__()
                else:
                    # 首 token 等待，3 秒超时发送心跳
                    chunk = await asyncio.wait_for(astream_iter.__anext__(), timeout=3.0)
                    if not first_token:
                        first_token = True
            except asyncio.TimeoutError:
                elapsed = int(time.time() - heartbeat_start)
                yield f"data: {json.dumps({'type': 'progress', 'message': f'模型正在生成报告（已等待 {elapsed} 秒）...'}, ensure_ascii=False)}\n\n"
                continue
            except StopAsyncIteration:
                break

            # 处理推理内容（DeepSeek-R1 等推理模型会先输出 reasoning_content）
            reasoning = ""
            if hasattr(chunk, "additional_kwargs") and isinstance(chunk.additional_kwargs, dict):
                reasoning = chunk.additional_kwargs.get("reasoning_content", "")
            if reasoning:
                yield f"data: {json.dumps({'type': 'reasoning', 'content': reasoning}, ensure_ascii=False)}\n\n"
                continue  # 推理内容不保存到 full_response

            content = chunk.content if hasattr(chunk, "content") else ""
            if content and isinstance(content, str):
                full_response += content
                yield f"data: {json.dumps({'type': 'token', 'source': 'fast', 'content': content}, ensure_ascii=False)}\n\n"

        # 6. 保存助手回复
        if save_to_db and full_response.strip():
            cleaned = _clean_response(full_response)
            save_message(session_id, "assistant", cleaned)

        yield f"data: {json.dumps({'type': 'done', 'thread_id': thread_id, 'session_id': session_id}, ensure_ascii=False)}\n\n"

    except Exception as e:
        traceback.print_exc()
        yield f"data: {json.dumps({'type': 'error', 'message': str(e)}, ensure_ascii=False)}\n\n"