# -*- 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"