SillyTavern_replica/backend/api/routes/chatWsRoute.py

"""
聊天 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
):
    """
    处理流式聊天请求 – engine callbacks emit worldbook_active / tasks_created / chunk.
    """
    try:
        print(f"[StreamChat] 🚀 开始流式处理")

        chunk_count = [0]

        async def on_worldbook_active(entries):
            if entries:
                print(f"[StreamChat] 📤 发送世界书激活信息: {len(entries)} 个条目")
                await websocket.send_json({
                    "type": "worldbook_active",
                    "entries": entries,
                })

        async def on_tasks_created(task_ids):
            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,
                })

        async def on_chunk(chunk):
            chunk_count[0] += 1
            if chunk_count[0] % 10 == 0:
                print(f"[StreamChat] 📤 已发送 {chunk_count[0]} 个 chunks")
            await websocket.send_json({"type": "chunk", "content": chunk})

        result = await workflow_service.process_chat_request_stream(
            request_data,
            on_chunk=on_chunk,
            on_worldbook_active=on_worldbook_active,
            on_tasks_created=on_tasks_created,
        )

        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 _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()
        # 不抛出异常，避免影响主流程