完成请求推送，但组装mes还有问题

2026-05-04 00:33:29 +08:00
parent 7fc9e10c99
commit 2050a30a52
93 changed files with 9499 additions and 2719 deletions
--- a/backend/api/routes/chatWsRoute.py
+++ b/backend/api/routes/chatWsRoute.py
@@ -0,0 +1,420 @@
+"""
+聊天 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. 接收前端消息
+            data = await websocket.receive_text()
+            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)} 个条目")
+            await websocket.send_json({
+                "type": "worldbook_active",
+                "entries": active_entries
+            })
+        
+        # ✅ 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
+        
+        # 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}")
+        # 不抛出异常，避免影响主流程