513 lines
21 KiB
Python
513 lines
21 KiB
Python
"""
|
||
聊天 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()
|
||
# 不抛出异常，避免影响主流程
|