""" 聊天 WebSocket 路由 处理实时对话生成 """ from fastapi import APIRouter, WebSocket, WebSocketDisconnect from typing import Dict, Any import json import asyncio try: from backend.services.chat_workflow_service import ChatWorkflowService from backend.services.chat_service import ChatService from backend.services.task_queue_manager import task_queue_manager from backend.core.config import settings except ImportError: from services.chat_workflow_service import ChatWorkflowService from services.chat_service import ChatService from services.task_queue_manager import task_queue_manager from core.config import settings router = APIRouter(prefix="/chat", tags=["chat-websocket"]) # 初始化服务 workflow_service = ChatWorkflowService() chat_service = ChatService(settings.DATA_PATH) # ✅ 全局变量：用于存储需要中断的聊天会话 interrupt_flags: Dict[str, bool] = {} @router.websocket("/{role_name}/{chat_name}/ws") async def websocket_chat_endpoint( websocket: WebSocket, role_name: str, chat_name: str ): """ WebSocket 聊天端点 接收前端发送的完整对话请求，调用工作流生成回复，支持流式输出 """ await websocket.accept() print(f"\n{'='*80}") print(f"[WebSocket] 📡 连接建立: {role_name}/{chat_name}") print(f"{'='*80}\n") chat_id = f"{role_name}/{chat_name}" try: while True: # 1. 接收前端消息 print(f"[WebSocket] ⏳ 等待接收消息...") data = await websocket.receive_text() print(f"[WebSocket] ✅ 收到消息，长度: {len(data)}") request_data = json.loads(data) # ✅ 检查是否是取消任务的请求 if request_data.get("type") == "cancel_task": task_id = request_data.get("taskId") print(f"[WebSocket] ❌ 收到取消任务请求: {task_id}") # ✅ 特殊处理：如果是 LLM 生成任务，需要中断当前流式生成 if task_id == "current_llm_generation": print(f"[WebSocket] 🛑 正在终止 LLM 流式生成...") # TODO: 实现 LLM 生成的中断逻辑 # 目前只能通过关闭连接来终止 await websocket.send_json({ "type": "task_cancelled", "taskId": task_id, "success": True, "message": "LLM 生成已终止" }) else: # 取消其他类型的任务（图像生成、动态表格等） success = await task_queue_manager.cancel_task(task_id) await websocket.send_json({ "type": "task_cancelled", "taskId": task_id, "success": success }) print(f"[WebSocket] ✅ 任务取消结果: {success}") continue print(f"\n{'-'*80}") print(f"[WebSocket] 📨 收到请求:") print(f" - Floor: {request_data.get('floor')}") print(f" - Role: {request_data.get('currentRole')}") print(f" - Chat: {request_data.get('currentChat')}") print(f" - Stream: {request_data.get('stream', False)}") print(f" - Message Length: {len(request_data.get('mes', ''))}") # ✅ 打印 API 配置信息（隐藏密钥） api_config = request_data.get('apiConfig', {}) current_profile = request_data.get('currentProfile', {}) profile_id = current_profile.get('id') if isinstance(current_profile, dict) else None print(f" - Profile ID: {profile_id or 'N/A'}") print(f" - API URL: {api_config.get('api_url', 'N/A')[:50]}..." if len(api_config.get('api_url', '')) > 50 else f" - API URL: {api_config.get('api_url', 'N/A')}") print(f" - Model: {api_config.get('model', 'N/A')}") # ✅ 始终从配置文件中读取 API Key（不信任前端传来的 Key） if profile_id: try: from .apiConfigRoute import load_profile profile = load_profile(profile_id) if profile: # 找到 mainLLM 的配置 main_llm_config = profile.get('apis', {}).get('mainLLM', {}) api_key = main_llm_config.get('apiKey', '') if api_key: # 使用明文 API Key api_config['api_key'] = api_key request_data['apiConfig'] = api_config print(f" - API Key: ✅ 已从配置文件加载") else: print(f" - API Key: ⚠️ 配置文件中未找到 Key") else: print(f" - API Key: ❌ 无法加载配置文件: {profile_id}") except Exception as e: print(f" - API Key: ❌ 加载失败: {str(e)}") import traceback traceback.print_exc() else: print(f" - API Key: ⚠️ 未提供 profileId，无法加载") print(f"{'-'*80}\n") # 2. 提取流式输出标志 stream_output = request_data.get("stream", False) if stream_output: # === 真正的流式输出模式 === print(f"[WebSocket] 🌊 进入流式处理模式") await _handle_stream_chat( websocket, role_name, chat_name, request_data, workflow_service ) else: # === 非流式输出模式 === print(f"[WebSocket] 📦 进入非流式处理模式") result = await workflow_service.process_chat_request(request_data) if result["success"]: content = result["content"] print(f"\n[WebSocket] ✨ 生成成功，内容长度: {len(content)}") # ✅ 发送激活的世界书条目信息 active_entries = result.get("activeEntries", []) print(f"[WebSocket] 📚 发送世界书激活信息: {len(active_entries)} 个条目") await websocket.send_json({ "type": "worldbook_active", "entries": active_entries }) # ✅ 发送任务ID信息 task_ids = result.get("taskIds", {}) if task_ids.get("imageWorkflow") or task_ids.get("dynamicTable"): print(f"[WebSocket] 📋 发送任务ID信息: {task_ids}") await websocket.send_json({ "type": "tasks_created", "tasks": task_ids }) # 一次性发送完整内容 print(f"[WebSocket] 📤 发送完整内容 (chunk)") await websocket.send_json({ "type": "chunk", "content": content }) print(f"[WebSocket] ✅ 发送完成信号") await websocket.send_json({ "type": "complete" }) # 保存消息 print(f"[WebSocket] 💾 保存消息到文件...") await _save_messages(role_name, chat_name, request_data, content) print(f"[WebSocket] ✅ 消息保存完成\n") else: error_msg = result["error"] print(f"[WebSocket] ❌ 处理失败: {error_msg}") await websocket.send_json({ "type": "error", "message": error_msg }) except WebSocketDisconnect: print(f"\n{'='*80}") print(f"[WebSocket] 🔌 连接断开: {role_name}/{chat_name}") print(f"{'='*80}\n") except Exception as e: print(f"\n{'='*80}") print(f"[WebSocket] ⚠️ 错误: {str(e)}") print(f"{'='*80}\n") import traceback traceback.print_exc() try: await websocket.send_json({ "type": "error", "message": f"服务器错误: {str(e)}" }) except: pass finally: try: await websocket.close() except: pass async def _handle_stream_chat( websocket: WebSocket, role_name: str, chat_name: str, request_data: Dict[str, Any], workflow_service ): """ 处理流式聊天请求 Args: websocket: WebSocket 连接 role_name: 角色名 chat_name: 聊天名 request_data: 请求数据 workflow_service: 工作流服务实例 """ try: print(f"[StreamChat] 🚀 开始流式处理") # ✅ 第1步：加载角色卡 current_role = request_data.get("currentRole") character_data = request_data.get("characterData") if character_data: try: from backend.models.internal import CharacterCard except ImportError: from models.internal import CharacterCard character = CharacterCard(**character_data) else: from backend.services.character_service import CharacterService character_service = CharacterService() character = character_service.get_character_by_name(current_role) if not character: print(f"[StreamChat] ❌ 错误: 无法加载角色 '{current_role}'") await websocket.send_json({ "type": "error", "message": f"角色 '{current_role}' 不存在" }) return print(f"[StreamChat] ✅ 已加载角色卡: {character.name}") # ✅ 第2步：激活世界书条目（在LLM调用之前） print(f"[StreamChat] 📚 正在激活世界书条目...") active_entries = await workflow_service._collect_and_activate_worldbooks( request_data, character ) # ✅ 发送激活的世界书条目信息（在LLM调用前） if active_entries: print(f"[StreamChat] 📤 发送世界书激活信息: {len(active_entries)} 个条目") # 将 Pydantic 模型转换为字典 entries_dict = [entry.model_dump() for entry in active_entries] await websocket.send_json({ "type": "worldbook_active", "entries": entries_dict }) # ✅ TODO: RAG检索（暂时为空，待实现） rag_results = [] if rag_results: print(f"[StreamChat] 🔍 发送 RAG 检索结果: {len(rag_results)} 条") await websocket.send_json({ "type": "rag_results", "results": rag_results }) # ✅ 第2步：启动并行任务（在LLM调用前创建任务ID） options = request_data.get("options", {}) task_ids = { "imageWorkflow": None, "dynamicTable": None } if options.get("imageWorkflow", False): import uuid chat_id = f"{role_name}/{chat_name}" task_ids["imageWorkflow"] = f"img_{uuid.uuid4().hex[:8]}" from backend.services.task_queue_manager import task_queue_manager, TaskType await task_queue_manager.add_task(task_ids["imageWorkflow"], TaskType.IMAGE_WORKFLOW, chat_id) if options.get("dynamicTable", False): import uuid chat_id = f"{role_name}/{chat_name}" task_ids["dynamicTable"] = f"tbl_{uuid.uuid4().hex[:8]}" from backend.services.task_queue_manager import task_queue_manager, TaskType await task_queue_manager.add_task(task_ids["dynamicTable"], TaskType.DYNAMIC_TABLE, chat_id) # ✅ 发送任务ID信息（在LLM调用前） if task_ids.get("imageWorkflow") or task_ids.get("dynamicTable"): print(f"[StreamChat] 📤 发送任务ID信息: {task_ids}") await websocket.send_json({ "type": "tasks_created", "tasks": task_ids }) # ✅ 第3步：调用LLM流式生成 chunk_count = [0] # 使用列表以便在闭包中修改 result = await workflow_service.process_chat_request_stream( request_data, on_chunk=lambda chunk: asyncio.create_task( _send_chunk_with_log(websocket, chunk, chunk_count) ) ) if result["success"]: content = result["content"] print(f"\n[StreamChat] ✨ 流式生成成功，总长度: {len(content)}") # 发送完成信号 print(f"[StreamChat] ✅ 发送完成信号") await websocket.send_json({ "type": "complete" }) # 保存消息 print(f"[StreamChat] 💾 保存消息到文件...") await _save_messages(role_name, chat_name, request_data, content) print(f"[StreamChat] ✅ 消息保存完成\n") else: error_msg = result["error"] print(f"[StreamChat] ❌ 流式处理失败: {error_msg}") await websocket.send_json({ "type": "error", "message": error_msg }) except Exception as e: print(f"\n[StreamChat] ⚠️ 错误: {str(e)}") import traceback traceback.print_exc() await websocket.send_json({ "type": "error", "message": f"流式处理失败: {str(e)}" }) async def _send_chunk_with_log(websocket: WebSocket, chunk: str, chunk_count: list): """ 发送 chunk 并记录日志 Args: websocket: WebSocket 连接 chunk: 文本片段 chunk_count: 计数器（使用列表以便在闭包中修改） """ chunk_count[0] += 1 if chunk_count[0] % 10 == 0: # 每10个chunk记录一次 print(f"[StreamChat] 📤 已发送 {chunk_count[0]} 个 chunks") await websocket.send_json({"type": "chunk", "content": chunk}) async def _save_messages( role_name: str, chat_name: str, request_data: Dict[str, Any], ai_response: str ): """ 保存用户消息和AI回复到聊天文件 Args: role_name: 角色名 chat_name: 聊天名 request_data: 前端发送的请求数据 ai_response: AI生成的回复 """ try: from datetime import datetime # ✅ 应用双 false 的正则规则（永久修改存储数据） from services.regex_service import regex_service from models.regex_rules import RegexPlacement # 获取预设名称 preset_config = request_data.get("presetConfig", {}) preset_name = preset_config.get("selectedPreset") # 计算消息深度 floor = request_data.get("floor", 0) message_depth = 0 # AI 回复是最新消息，深度为 0 # ✅ 应用 AI Output 正则规则（placement=2） # 只应用双 false 的规则（markdownOnly=false 且 promptOnly=false） processed_ai_response = regex_service.apply_rules_by_placement( text=ai_response, placement=RegexPlacement.AI_OUTPUT.value, character_name=role_name, preset_name=preset_name, message_depth=message_depth, is_for_llm=False, # ✅ 不是发送给 LLM，是保存数据 is_markdown_rendered=False # ✅ 不是 Markdown 渲染后 ) # 如果处理后的内容与原始内容不同，说明有双 false 规则被应用 if processed_ai_response != ai_response: print(f"[Regex] ✅ 已应用双 false 正则规则（永久修改存储数据）") ai_response = processed_ai_response # ✅ 检查是否是重roll模式（targetFloor 存在且不为 null） target_floor = request_data.get("floor") is_reroll = target_floor is not None if is_reroll: # ✅ 重roll模式：更新现有消息的 swipes 数组 print(f"[WebSocket] 🔄 重roll模式，更新楼层 {target_floor} 的 swipes") # 获取现有的消息 existing_message = chat_service.get_message(role_name, chat_name, target_floor) if not existing_message: print(f"[WebSocket] ⚠️ 找不到楼层 {target_floor} 的消息，创建新消息") # 如果找不到，创建新消息（兼容处理） ai_message = { "id": f"msg_{datetime.now().timestamp()}_ai", "name": request_data.get("characterName", role_name), "is_user": False, "is_system": False, "sendDate": datetime.now().isoformat(), "mes": ai_response, "chatId": f"{role_name}/{chat_name}", "floor": target_floor, "swipes": [ai_response], "swipe_id": 0 } chat_service.add_message(role_name, chat_name, ai_message) else: # ✅ 更新 swipes 数组 existing_swipes = existing_message.get("swipes", []) current_mes = existing_message.get("mes", "") # 构建新的 swipes 数组 updated_swipes = list(existing_swipes) # 复制现有swipes # 如果当前 mes 不在 swipes 中，先添加它 if current_mes and current_mes not in updated_swipes: updated_swipes.append(current_mes) print(f"[WebSocket] 📝 将当前内容添加到 swipes") # 添加新生成的内容 updated_swipes.append(ai_response) print(f"[WebSocket] 📊 Swipes 更新: {len(existing_swipes)} -> {len(updated_swipes)}") # 更新消息 update_data = { "mes": ai_response, # 显示最新内容 "swipes": updated_swipes, # 更新 swipes 数组 "swipe_id": len(updated_swipes) - 1 # 自动切换到新版本 } chat_service.update_message(role_name, chat_name, target_floor, update_data) print(f"[WebSocket] ✅ 楼层 {target_floor} 已更新，swipes 数量: {len(updated_swipes)}") else: # ✅ 正常模式：创建新的用户消息和AI消息 print(f"[WebSocket] ➕ 正常模式，创建新消息") # 1. 保存用户消息 user_message = { "id": f"msg_{datetime.now().timestamp()}_user", "name": request_data.get("userName", "User"), "is_user": True, "is_system": False, "sendDate": datetime.now().isoformat(), "mes": request_data.get("mes", ""), "chatId": f"{role_name}/{chat_name}", "floor": request_data.get("floor", 0) } chat_service.add_message(role_name, chat_name, user_message) # 2. 保存AI回复 ai_message = { "id": f"msg_{datetime.now().timestamp()}_ai", "name": request_data.get("characterName", role_name), "is_user": False, "is_system": False, "sendDate": datetime.now().isoformat(), "mes": ai_response, "chatId": f"{role_name}/{chat_name}", "floor": request_data.get("floor", 0) + 1 } chat_service.add_message(role_name, chat_name, ai_message) print(f"[WebSocket] ✅ 新消息已保存: {role_name}/{chat_name}") except Exception as e: print(f"[WebSocket] 保存消息失败: {e}") import traceback traceback.print_exc() # 不抛出异常，避免影响主流程