1168 lines
43 KiB
Python
1168 lines
43 KiB
Python
"""
|
|
WebSocket Streaming API
|
|
|
|
提供实时执行进度更新的 WebSocket 端点
|
|
"""
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, Query
|
|
from typing import Dict, Set, Optional, Any
|
|
import json
|
|
import asyncio
|
|
from datetime import datetime
|
|
from pathlib import Path
|
|
|
|
from app.config import settings
|
|
from app.utils.logger import get_logger
|
|
from app.core.agent_runtime.director_agent import DirectorAgent
|
|
from app.db.repositories import message_repo
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
# ============================================
|
|
# WebSocket 连接管理
|
|
# ============================================
|
|
|
|
class ConnectionManager:
|
|
"""WebSocket 连接管理器"""
|
|
|
|
def __init__(self):
|
|
# 项目ID -> WebSocket连接集合
|
|
self.project_connections: Dict[str, Set[WebSocket]] = {}
|
|
# 批次ID -> WebSocket连接集合
|
|
self.batch_connections: Dict[str, Set[WebSocket]] = {}
|
|
# 项目ID -> Agent实例
|
|
self.project_agents: Dict[str, DirectorAgent] = {}
|
|
|
|
async def connect_to_project(self, websocket: WebSocket, project_id: str):
|
|
"""连接到项目执行流"""
|
|
await websocket.accept()
|
|
if project_id not in self.project_connections:
|
|
self.project_connections[project_id] = set()
|
|
self.project_connections[project_id].add(websocket)
|
|
logger.info(f"WebSocket 已连接到项目: {project_id}")
|
|
|
|
async def connect_to_batch(self, websocket: WebSocket, batch_id: str):
|
|
"""连接到批次执行流"""
|
|
await websocket.accept()
|
|
if batch_id not in self.batch_connections:
|
|
self.batch_connections[batch_id] = set()
|
|
self.batch_connections[batch_id].add(websocket)
|
|
logger.info(f"WebSocket 已连接到批次: {batch_id}")
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
"""断开连接"""
|
|
# 从所有项目连接中移除
|
|
for project_id, connections in self.project_connections.items():
|
|
if websocket in connections:
|
|
connections.remove(websocket)
|
|
logger.info(f"WebSocket 已从项目断开: {project_id}")
|
|
if not connections:
|
|
del self.project_connections[project_id]
|
|
# 清理 Agent 实例
|
|
if project_id in self.project_agents:
|
|
del self.project_agents[project_id]
|
|
|
|
# 从所有批次连接中移除
|
|
for batch_id, connections in self.batch_connections.items():
|
|
if websocket in connections:
|
|
connections.remove(websocket)
|
|
logger.info(f"WebSocket 已从批次断开: {batch_id}")
|
|
if not connections:
|
|
del self.batch_connections[batch_id]
|
|
|
|
async def send_to_project(
|
|
self,
|
|
project_id: str,
|
|
message: Dict[str, Any],
|
|
exclude: Optional[WebSocket] = None
|
|
):
|
|
"""向项目的所有连接发送消息"""
|
|
if project_id in self.project_connections:
|
|
disconnected = set()
|
|
for connection in self.project_connections[project_id]:
|
|
if connection != exclude:
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"发送消息失败: {str(e)}")
|
|
disconnected.add(connection)
|
|
|
|
# 清理断开的连接
|
|
for connection in disconnected:
|
|
self.disconnect(connection)
|
|
|
|
async def send_to_batch(
|
|
self,
|
|
batch_id: str,
|
|
message: Dict[str, Any],
|
|
exclude: Optional[WebSocket] = None
|
|
):
|
|
"""向批次的所有连接发送消息"""
|
|
if batch_id in self.batch_connections:
|
|
disconnected = set()
|
|
for connection in self.batch_connections[batch_id]:
|
|
if connection != exclude:
|
|
try:
|
|
await connection.send_json(message)
|
|
except Exception as e:
|
|
logger.error(f"发送消息失败: {str(e)}")
|
|
disconnected.add(connection)
|
|
|
|
# 清理断开的连接
|
|
for connection in disconnected:
|
|
self.disconnect(connection)
|
|
|
|
def get_agent(self, project_id: str, working_dir: Path) -> DirectorAgent:
|
|
"""获取或创建 Agent 实例"""
|
|
if project_id not in self.project_agents:
|
|
# 确保工作目录存在
|
|
working_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 检查配置的模型类型
|
|
model_name = settings.zai_model
|
|
enable_thinking = True
|
|
|
|
# 如果是 GLM 模型，禁用 thinking 模式（不支持）
|
|
if "glm" in model_name.lower():
|
|
enable_thinking = False
|
|
|
|
# 不在这里加载项目上下文，而是在 WebSocket 消息处理时加载
|
|
# 因为 get_agent 是同步方法，而 project_repo.get 是异步的
|
|
self.project_agents[project_id] = DirectorAgent(
|
|
working_directory=working_dir,
|
|
enable_thinking=enable_thinking,
|
|
model=model_name,
|
|
project_context=None # 稍后在消息处理中更新
|
|
)
|
|
return self.project_agents[project_id]
|
|
|
|
def get_project_connections_count(self, project_id: str) -> int:
|
|
"""获取项目的连接数"""
|
|
return len(self.project_connections.get(project_id, set()))
|
|
|
|
def get_batch_connections_count(self, batch_id: str) -> int:
|
|
"""获取批次的连接数"""
|
|
return len(self.batch_connections.get(batch_id, set()))
|
|
|
|
|
|
# 全局连接管理器
|
|
manager = ConnectionManager()
|
|
|
|
|
|
# ============================================
|
|
# WebSocket 端点
|
|
# ============================================
|
|
|
|
@router.websocket("/ws/projects/{project_id}/execute")
|
|
async def websocket_project_execution(
|
|
websocket: WebSocket,
|
|
project_id: str
|
|
):
|
|
"""
|
|
项目执行 WebSocket 端点
|
|
"""
|
|
await manager.connect_to_project(websocket, project_id)
|
|
|
|
# 准备工作目录 (假设在 projects/{id})
|
|
# 注意：这里需要根据实际配置调整路径
|
|
project_dir = Path(f"d:/platform/creative_studio/workspace/projects/{project_id}")
|
|
|
|
try:
|
|
# 发送连接确认
|
|
await websocket.send_json({
|
|
"type": "connected",
|
|
"data": {
|
|
"project_id": project_id,
|
|
"message": "已连接到项目执行流",
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
})
|
|
|
|
# 加载并发送历史消息
|
|
history = await message_repo.get_history(project_id)
|
|
if history:
|
|
await websocket.send_json({
|
|
"type": "history",
|
|
"messages": history
|
|
})
|
|
|
|
# 保持连接并接收客户端消息
|
|
while True:
|
|
try:
|
|
# 接收客户端消息
|
|
data = await websocket.receive_text()
|
|
message = json.loads(data)
|
|
|
|
# 处理客户端消息
|
|
await _handle_client_message(websocket, project_id, message, project_dir)
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info(f"WebSocket 客户端主动断开: {project_id}")
|
|
break
|
|
except json.JSONDecodeError:
|
|
await websocket.send_json({
|
|
"type": "error",
|
|
"data": {
|
|
"message": "无效的 JSON 格式"
|
|
}
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"处理 WebSocket 消息错误: {str(e)}")
|
|
await websocket.send_json({
|
|
"type": "error",
|
|
"data": {
|
|
"message": f"处理消息错误: {str(e)}"
|
|
}
|
|
})
|
|
|
|
except Exception as e:
|
|
logger.error(f"WebSocket 连接错误: {str(e)}")
|
|
|
|
finally:
|
|
manager.disconnect(websocket)
|
|
|
|
|
|
@router.websocket("/ws/batches/{batch_id}/execute")
|
|
async def websocket_batch_execution(
|
|
websocket: WebSocket,
|
|
batch_id: str
|
|
):
|
|
"""
|
|
批量执行 WebSocket 端点
|
|
"""
|
|
await manager.connect_to_batch(websocket, batch_id)
|
|
|
|
try:
|
|
# 发送连接确认
|
|
await websocket.send_json({
|
|
"type": "connected",
|
|
"data": {
|
|
"batch_id": batch_id,
|
|
"message": "已连接到批量执行流",
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
})
|
|
|
|
# 保持连接
|
|
while True:
|
|
try:
|
|
data = await websocket.receive_text()
|
|
# 批量执行目前不接受客户端控制消息，仅广播
|
|
# 但为了保持连接活性，可以处理 ping
|
|
message = json.loads(data)
|
|
if message.get("type") == "ping":
|
|
await websocket.send_json({
|
|
"type": "pong",
|
|
"data": {"timestamp": datetime.now().isoformat()}
|
|
})
|
|
|
|
except WebSocketDisconnect:
|
|
logger.info(f"WebSocket 客户端主动断开: {batch_id}")
|
|
break
|
|
except Exception as e:
|
|
logger.error(f"处理 WebSocket 消息错误: {str(e)}")
|
|
|
|
finally:
|
|
manager.disconnect(websocket)
|
|
|
|
|
|
# ============================================
|
|
# 消息处理
|
|
# ============================================
|
|
|
|
async def _ensure_agent_context(
|
|
agent: DirectorAgent,
|
|
project_id: str,
|
|
active_episode_number: Optional[int] = None,
|
|
active_episode_title: Optional[str] = None
|
|
):
|
|
"""确保 Agent 加载了正确的项目上下文"""
|
|
# 如果已经加载了该项目的上下文，且剧集信息一致，则跳过
|
|
if (agent.context and
|
|
agent.context.project_id == project_id and
|
|
agent.context.active_episode_number == active_episode_number):
|
|
return
|
|
|
|
try:
|
|
from app.db.repositories import project_repo
|
|
from app.core.agent_runtime.context import SkillAgentContext
|
|
from app.core.skills.skill_manager import skill_manager
|
|
|
|
project = await project_repo.get(project_id)
|
|
if not project:
|
|
logger.warning(f"Project {project_id} not found in repository")
|
|
return
|
|
|
|
# 1. 加载用户技能
|
|
user_skills = []
|
|
default_task_skills = getattr(project, 'defaultTaskSkills', [])
|
|
if isinstance(default_task_skills, list):
|
|
for task_config in default_task_skills:
|
|
task_skills = getattr(task_config, 'skills', [])
|
|
if isinstance(task_skills, list):
|
|
for skill_config in task_skills:
|
|
try:
|
|
skill_id = getattr(skill_config, 'skill_id', None) or (skill_config.get('skill_id') if isinstance(skill_config, dict) else None)
|
|
if skill_id:
|
|
skill = skill_manager.get_skill_by_id(skill_id)
|
|
if skill:
|
|
user_skills.append({
|
|
'id': skill.id,
|
|
'name': skill.name,
|
|
'behavior': skill.behavior_guide or skill.description or ''
|
|
})
|
|
except Exception as e:
|
|
logger.warning(f"Failed to load skill: {e}")
|
|
|
|
# 2. 整合角色信息
|
|
characters_text = ""
|
|
if project.globalContext:
|
|
if project.globalContext.styleGuide:
|
|
characters_text += project.globalContext.styleGuide + "\n\n"
|
|
|
|
if project.globalContext.characterProfiles:
|
|
# 确保 characterProfiles 是字典类型
|
|
profiles = project.globalContext.characterProfiles
|
|
if isinstance(profiles, dict):
|
|
characters_text += "### Detailed Character Profiles:\n"
|
|
for name, profile in profiles.items():
|
|
characters_text += f"- **{name}**: {getattr(profile, 'description', '')}\n"
|
|
if hasattr(profile, 'personality') and profile.personality:
|
|
characters_text += f" Personality: {profile.personality}\n"
|
|
else:
|
|
# characterProfiles 是字符串或其他类型，直接添加
|
|
characters_text += str(profiles) + "\n\n"
|
|
|
|
# 3. 获取所有剧集
|
|
from app.db.repositories import episode_repo
|
|
episodes_data = await episode_repo.list_by_project(project_id)
|
|
episodes = []
|
|
for ep in episodes_data:
|
|
episodes.append({
|
|
"number": ep.number,
|
|
"title": ep.title,
|
|
"status": ep.status
|
|
})
|
|
|
|
# 4. 创建项目上下文
|
|
project_context = SkillAgentContext(
|
|
skill_loader=agent.context.skill_loader,
|
|
working_directory=agent.context.working_directory,
|
|
project_id=project.id,
|
|
project_name=project.name,
|
|
project_genre=getattr(project, 'genre', '古风'),
|
|
total_episodes=project.totalEpisodes,
|
|
world_setting=project.globalContext.worldSetting if project.globalContext else None,
|
|
characters=characters_text.strip() or None,
|
|
overall_outline=project.globalContext.overallOutline if project.globalContext else None,
|
|
creation_mode='script' if (project.globalContext and project.globalContext.uploadedScript) else 'inspiration',
|
|
source_content=(project.globalContext.uploadedScript if project.globalContext and project.globalContext.uploadedScript
|
|
else project.globalContext.inspiration if project.globalContext else None),
|
|
user_skills=user_skills,
|
|
active_episode_number=active_episode_number,
|
|
active_episode_title=active_episode_title,
|
|
episodes=episodes
|
|
)
|
|
|
|
# 5. 更新 Agent
|
|
agent.context = project_context
|
|
agent.system_prompt = agent._build_system_prompt()
|
|
if hasattr(agent, 'refresh_agent'):
|
|
agent.refresh_agent()
|
|
|
|
msg = f"Successfully injected project context for {project_id}: {project.name}"
|
|
if active_episode_number:
|
|
msg += f" (Focusing on Episode {active_episode_number})"
|
|
logger.info(msg)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Failed to load project context for {project_id}: {e}", exc_info=True)
|
|
|
|
|
|
async def _handle_client_message(
|
|
websocket: WebSocket,
|
|
project_id: str,
|
|
message: Dict[str, Any],
|
|
project_dir: Path
|
|
):
|
|
"""
|
|
处理客户端发送的消息
|
|
"""
|
|
message_type = message.get("type")
|
|
logger.info(f"收到消息: type={message_type}, full_message={message}")
|
|
|
|
if message_type == "ping":
|
|
# 心跳响应
|
|
await websocket.send_json({
|
|
"type": "pong",
|
|
"data": {
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
})
|
|
|
|
elif message_type == "focus_episode":
|
|
# 客户端请求切换关注的剧集
|
|
episode_number = message.get("episodeNumber")
|
|
episode_title = message.get("episodeTitle")
|
|
|
|
agent = manager.get_agent(project_id, project_dir)
|
|
await _ensure_agent_context(agent, project_id, episode_number, episode_title)
|
|
|
|
await websocket.send_json({
|
|
"type": "focus_confirmed",
|
|
"data": {
|
|
"episodeNumber": episode_number,
|
|
"episodeTitle": episode_title
|
|
}
|
|
})
|
|
|
|
elif message_type == "update_episode_title":
|
|
# 用户手动编辑画布标题
|
|
episode_number = message.get("episodeNumber")
|
|
new_title = message.get("title")
|
|
|
|
if episode_number and new_title:
|
|
try:
|
|
from app.db.repositories import episode_repo
|
|
episodes = await episode_repo.list_by_project(project_id)
|
|
episode = next((ep for ep in episodes if ep.number == episode_number), None)
|
|
if episode:
|
|
episode.title = new_title
|
|
await episode_repo.update(episode)
|
|
logger.info(f"Manual title update for EP{episode_number}: {new_title}")
|
|
|
|
# 广播更新给所有连接
|
|
await manager.send_to_project(project_id, {
|
|
"type": "episode_updated",
|
|
"data": {
|
|
"number": episode_number,
|
|
"title": new_title,
|
|
"status": episode.status
|
|
}
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Failed to update episode title: {e}")
|
|
|
|
elif message_type == "inbox_action":
|
|
# 用户在 Inbox 中点击批准或拒绝
|
|
action = message.get("action")
|
|
item_id = message.get("itemId")
|
|
|
|
# 将操作转换为自然语言反馈给 Agent
|
|
feedback = f"User {action}ed inbox item {item_id}."
|
|
|
|
agent = manager.get_agent(project_id, project_dir)
|
|
# 确保上下文已加载
|
|
await _ensure_agent_context(agent, project_id)
|
|
|
|
try:
|
|
for event in agent.stream_events(feedback, thread_id=project_id):
|
|
# 同样的事件处理逻辑
|
|
if event.get("type") == "tool_call":
|
|
await _handle_tool_call(project_id, event)
|
|
await manager.send_to_project(project_id, event)
|
|
except Exception as e:
|
|
await manager.send_to_project(project_id, {
|
|
"type": "error",
|
|
"data": {"message": str(e)}
|
|
})
|
|
|
|
elif message_type == "chat_message":
|
|
# 用户发送聊天消息 -> 触发 Agent 执行
|
|
content = message.get("content", "")
|
|
episode_number = message.get("episodeNumber")
|
|
episode_title = message.get("episodeTitle")
|
|
|
|
if not content:
|
|
return
|
|
|
|
# 保存用户消息
|
|
await message_repo.add_message(project_id, "user", content)
|
|
|
|
# 获取 Agent
|
|
agent = manager.get_agent(project_id, project_dir)
|
|
|
|
# 确保上下文已加载（包含当前剧集信息）
|
|
await _ensure_agent_context(agent, project_id, episode_number, episode_title)
|
|
|
|
# 异步运行 Agent 并将事件流推送到前端
|
|
full_response = ""
|
|
try:
|
|
for event in agent.stream_events(content, thread_id=project_id):
|
|
# 检查特殊工具调用并转换格式
|
|
if event.get("type") == "tool_call":
|
|
await _handle_tool_call(project_id, event)
|
|
|
|
# 收集 Agent 回复内容
|
|
if event.get("type") == "text":
|
|
full_response += event.get("content", "")
|
|
|
|
await manager.send_to_project(project_id, event)
|
|
|
|
# 保存 Agent 回复
|
|
if full_response:
|
|
await message_repo.add_message(project_id, "agent", full_response)
|
|
|
|
except Exception as e:
|
|
await manager.send_to_project(project_id, {
|
|
"type": "error",
|
|
"data": {"message": str(e)}
|
|
})
|
|
|
|
elif message_type == "get_status":
|
|
# 请求状态
|
|
# 这里可以返回 Agent 的状态，或者之前的 executor 状态
|
|
pass
|
|
|
|
else:
|
|
await websocket.send_json({
|
|
"type": "error",
|
|
"data": {
|
|
"message": f"未知消息类型: {message_type}"
|
|
}
|
|
})
|
|
|
|
|
|
async def _sync_context_states(
|
|
project_id: str,
|
|
episode_number: int,
|
|
memory: Any
|
|
):
|
|
"""同步上下文状态到前端"""
|
|
try:
|
|
# 提取角色状态作为上下文状态
|
|
context_states = []
|
|
|
|
# 添加时间状态
|
|
context_states.append({
|
|
"type": "time",
|
|
"value": f"EP{episode_number} 完成后"
|
|
})
|
|
|
|
# 添加角色状态
|
|
character_states = getattr(memory, 'characterStates', {})
|
|
if isinstance(character_states, dict):
|
|
for char_name, states in character_states.items():
|
|
if states and isinstance(states, list) and len(states) > 0:
|
|
latest_state = states[-1]
|
|
if isinstance(latest_state, dict):
|
|
state_value = latest_state.get('state', f"{char_name}状态")
|
|
else:
|
|
state_value = getattr(latest_state, 'state', f"{char_name}状态")
|
|
|
|
context_states.append({
|
|
"type": "character",
|
|
"value": f"{char_name}: {state_value}",
|
|
"character": char_name,
|
|
"state": state_value
|
|
})
|
|
|
|
# 添加待收线数量
|
|
pending_threads = getattr(memory, 'pendingThreads', [])
|
|
if pending_threads:
|
|
context_states.append({
|
|
"type": "pending_threads",
|
|
"value": f"待收线: {len(pending_threads)} 条"
|
|
})
|
|
|
|
# 广播上下文更新
|
|
await manager.send_to_project(project_id, {
|
|
"type": "context_update",
|
|
"states": context_states,
|
|
"episode_number": episode_number
|
|
})
|
|
|
|
logger.info(f"已同步上下文状态到项目 {project_id}, {len(context_states)} 个状态")
|
|
|
|
except Exception as e:
|
|
logger.error(f"同步上下文状态失败: {str(e)}")
|
|
|
|
|
|
async def _handle_tool_call(project_id: str, event: Dict[str, Any]):
|
|
"""
|
|
处理工具调用，转换为特定的 WebSocket 消息
|
|
|
|
这个函数在 agent.stream_events() 中被调用，当检测到 director 工具调用时，
|
|
会将其转换为前端可以理解的 WebSocket 事件格式。
|
|
"""
|
|
name = event.get("name")
|
|
args = event.get("args", {})
|
|
|
|
# Director 工具处理
|
|
if name == "update_plan":
|
|
await manager.send_to_project(project_id, {
|
|
"type": "plan_update",
|
|
"plan": args.get("steps", []),
|
|
"status": args.get("status", "planning"),
|
|
"current_step_index": args.get("current_step_index", 0)
|
|
})
|
|
|
|
elif name == "add_inbox_task":
|
|
await manager.send_to_project(project_id, {
|
|
"type": "review_request",
|
|
"id": f"task_{args.get('title', 'unknown')}_{int(datetime.now().timestamp())}",
|
|
"title": args.get("title"),
|
|
"description": args.get("description"),
|
|
"options": args.get("options", ["Approve", "Reject"]),
|
|
"timestamp": int(datetime.now().timestamp() * 1000)
|
|
})
|
|
|
|
elif name == "add_annotation":
|
|
await manager.send_to_project(project_id, {
|
|
"type": "annotation_add",
|
|
"annotation": {
|
|
"content": args.get("content"),
|
|
"type": args.get("annotation_type", "review"),
|
|
"suggestion": args.get("suggestion", ""),
|
|
"timestamp": int(datetime.now().timestamp() * 1000)
|
|
}
|
|
})
|
|
|
|
elif name == "update_context":
|
|
# 解析 data (可能是 JSON string 或 dict)
|
|
data = args.get("data")
|
|
context_type = args.get("context_type", "state")
|
|
|
|
try:
|
|
if isinstance(data, str):
|
|
data = json.loads(data)
|
|
|
|
# 转换为前端期望的 activeStates 格式
|
|
if isinstance(data, dict):
|
|
# 将字典转换为 [{type, value}, ...] 格式
|
|
states = [{"type": k, "value": v} for k, v in data.items()]
|
|
elif isinstance(data, list):
|
|
# 已经是列表格式
|
|
states = data
|
|
else:
|
|
states = [{"type": context_type, "value": str(data)}]
|
|
|
|
await manager.send_to_project(project_id, {
|
|
"type": "context_update",
|
|
"states": states
|
|
})
|
|
except Exception as e:
|
|
logger.warning(f"Failed to process context update: {e}")
|
|
|
|
elif name == "write_to_canvas":
|
|
# 新的 write_to_canvas 工具
|
|
content = args.get("content", "")
|
|
if content:
|
|
await manager.send_to_project(project_id, {
|
|
"type": "canvas_update",
|
|
"content": content
|
|
})
|
|
|
|
elif name == "write_file":
|
|
# 如果写入的是当前画布文件，也更新画布
|
|
# 这里简化：只要有内容就更新画布
|
|
content = args.get("content")
|
|
if content:
|
|
await manager.send_to_project(project_id, {
|
|
"type": "canvas_update",
|
|
"content": content
|
|
})
|
|
|
|
elif name == "update_memory":
|
|
# 处理记忆库更新
|
|
memory_type = args.get("memory_type", "timeline")
|
|
data = args.get("data", {})
|
|
|
|
# 根据记忆类型格式化数据
|
|
memory_data = {
|
|
"type": memory_type,
|
|
"memory_type": memory_type,
|
|
"timestamp": int(datetime.now().timestamp() * 1000)
|
|
}
|
|
|
|
# 添加具体信息
|
|
if memory_type == "timeline":
|
|
memory_data.update({
|
|
"title": data.get("event", "时间线事件"),
|
|
"description": data.get("description", "")
|
|
})
|
|
elif memory_type == "character_state":
|
|
memory_data.update({
|
|
"title": f"{data.get('character', '角色')}状态变化",
|
|
"description": data.get("state", ""),
|
|
"character": data.get("character", "")
|
|
})
|
|
elif memory_type == "pending_thread":
|
|
memory_data.update({
|
|
"title": "待收线问题",
|
|
"description": data.get("description", "")
|
|
})
|
|
elif memory_type == "foreshadowing":
|
|
memory_data.update({
|
|
"title": "伏笔",
|
|
"description": data.get("description", "")
|
|
})
|
|
|
|
await manager.send_to_project(project_id, {
|
|
"type": "memory_update",
|
|
"data": memory_data
|
|
})
|
|
|
|
elif name == "save_episode":
|
|
# 处理剧集保存并持久化
|
|
episode_number = args.get("episode_number")
|
|
title = args.get("title")
|
|
content = args.get("content")
|
|
outline = args.get("outline")
|
|
|
|
if episode_number:
|
|
try:
|
|
from app.db.repositories import episode_repo, project_repo
|
|
from app.core.memory.memory_manager import get_memory_manager
|
|
episodes = await episode_repo.list_by_project(project_id)
|
|
episode = next((ep for ep in episodes if ep.number == episode_number), None)
|
|
|
|
if episode:
|
|
if title: episode.title = title
|
|
if content: episode.content = content
|
|
if outline: episode.outline = outline
|
|
episode.status = "completed"
|
|
await episode_repo.update(episode)
|
|
logger.info(f"Persisted episode {episode_number} via save_episode")
|
|
|
|
# 自动更新记忆库
|
|
try:
|
|
project = await project_repo.get(project_id)
|
|
if project and episode.content:
|
|
memory_manager = get_memory_manager()
|
|
await memory_manager.update_memory_from_episode(project, episode)
|
|
logger.info(f"Updated memory after saving episode {episode_number}")
|
|
|
|
# 同步上下文状态到前端
|
|
await _sync_context_states(project_id, episode_number, project.memory)
|
|
except Exception as memory_error:
|
|
logger.warning(f"Failed to update memory for episode {episode_number}: {memory_error}")
|
|
|
|
# 广播更新
|
|
await manager.send_to_project(project_id, {
|
|
"type": "episode_updated",
|
|
"data": {
|
|
"number": episode_number,
|
|
"title": episode.title,
|
|
"status": episode.status
|
|
}
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Failed to persist episode via save_episode: {e}")
|
|
|
|
await manager.send_to_project(project_id, {
|
|
"type": "episode_saved",
|
|
"episode_number": episode_number,
|
|
"title": title
|
|
})
|
|
|
|
elif name == "update_episode":
|
|
# 处理剧集更新并持久化
|
|
episode_number = args.get("episode_number")
|
|
title = args.get("title")
|
|
content = args.get("content")
|
|
status = args.get("status")
|
|
outline = args.get("outline")
|
|
|
|
if episode_number:
|
|
try:
|
|
from app.db.repositories import episode_repo, project_repo
|
|
from app.core.memory.memory_manager import get_memory_manager
|
|
episodes = await episode_repo.list_by_project(project_id)
|
|
episode = next((ep for ep in episodes if ep.number == episode_number), None)
|
|
|
|
if episode:
|
|
if title is not None: episode.title = title
|
|
if content is not None: episode.content = content
|
|
if status is not None: episode.status = status
|
|
if outline is not None: episode.outline = outline
|
|
|
|
await episode_repo.update(episode)
|
|
logger.info(f"Updated episode {episode_number} via update_episode")
|
|
|
|
# 如果有新内容，自动更新记忆库
|
|
if content and content.strip():
|
|
try:
|
|
project = await project_repo.get(project_id)
|
|
if project:
|
|
memory_manager = get_memory_manager()
|
|
await memory_manager.update_memory_from_episode(project, episode)
|
|
logger.info(f"Updated memory after updating episode {episode_number}")
|
|
|
|
# 同步上下文状态到前端
|
|
await _sync_context_states(project_id, episode_number, project.memory)
|
|
except Exception as memory_error:
|
|
logger.warning(f"Failed to update memory for episode {episode_number}: {memory_error}")
|
|
|
|
# 广播更新
|
|
await manager.send_to_project(project_id, {
|
|
"type": "episode_updated",
|
|
"data": {
|
|
"number": episode_number,
|
|
"title": episode.title,
|
|
"status": episode.status
|
|
}
|
|
})
|
|
except Exception as e:
|
|
logger.error(f"Failed to update episode: {e}")
|
|
|
|
elif name == "focus_episode":
|
|
# 处理剧集焦点切换
|
|
episode_number = args.get("episode_number")
|
|
title = args.get("title")
|
|
|
|
await manager.send_to_project(project_id, {
|
|
"type": "focus_update",
|
|
"episodeNumber": episode_number,
|
|
"episodeTitle": title
|
|
})
|
|
|
|
elif name == "create_episode":
|
|
# 处理剧集创作请求
|
|
episode_number = args.get("episode_number")
|
|
analyze_previous_memory = args.get("analyze_previous_memory", True)
|
|
|
|
if episode_number:
|
|
# 使用 ensure_future 确保任务在后台执行，即使 WebSocket 断开也能继续
|
|
asyncio.ensure_future(_execute_episode_creation(
|
|
project_id, episode_number, analyze_previous_memory
|
|
))
|
|
logger.info(f"已启动后台创作任务: EP{episode_number}")
|
|
|
|
# ============================================
|
|
# 辅助函数 - 用于从其他模块发送消息
|
|
# ============================================
|
|
|
|
|
|
# 全局后台任务跟踪
|
|
_background_tasks: Dict[str, asyncio.Task] = {}
|
|
|
|
|
|
async def _execute_episode_creation(
|
|
project_id: str,
|
|
episode_number: int,
|
|
analyze_previous_memory: bool
|
|
):
|
|
"""
|
|
异步执行剧集创作
|
|
|
|
这个函数在后台执行，不会阻塞 WebSocket 连接。
|
|
即使 WebSocket 断开，任务也会继续执行并保存到数据库。
|
|
|
|
它会：
|
|
1. 分析上一集的记忆（如果需要）
|
|
2. 执行剧集创作
|
|
3. 通过 WebSocket 发送进度更新（如果连接存在）
|
|
4. 将创作内容发送到画布（如果连接存在）
|
|
5. 无论如何都保存到数据库
|
|
"""
|
|
task_key = f"{project_id}_{episode_number}"
|
|
|
|
try:
|
|
from app.db.repositories import project_repo, episode_repo
|
|
from app.core.agents.series_creation_agent import get_series_agent
|
|
from app.core.memory.memory_manager import get_memory_manager
|
|
|
|
# 注册到后台任务字典
|
|
current_task = asyncio.current_task()
|
|
if current_task:
|
|
_background_tasks[task_key] = current_task
|
|
|
|
logger.info(f"开始后台创作任务: {task_key}")
|
|
|
|
# 辅助函数：安全发送 WebSocket 消息（忽略连接错误）
|
|
async def safe_send(message_type: str, data: dict = None):
|
|
try:
|
|
if data is None:
|
|
data = {}
|
|
await manager.send_to_project(project_id, {
|
|
"type": message_type,
|
|
"data": data
|
|
})
|
|
except Exception as e:
|
|
# WebSocket 可能已断开，忽略错误继续执行
|
|
logger.debug(f"WebSocket 发送失败（可能已断开）: {e}")
|
|
|
|
# 获取项目
|
|
project = await project_repo.get(project_id)
|
|
if not project:
|
|
await safe_send("error", {"message": f"项目不存在: {project_id}"})
|
|
return
|
|
|
|
# 更新计划状态 - 开始
|
|
await safe_send("plan_update", {
|
|
"plan": [
|
|
f"分析 EP{episode_number - 1 if episode_number > 1 else 'N/A'} 的记忆系统" if analyze_previous_memory and episode_number > 1 else "跳过记忆分析（首集）",
|
|
f"生成 EP{episode_number} 大纲",
|
|
f"创作 EP{episode_number} 对话内容",
|
|
f"执行质量审核",
|
|
f"更新记忆系统"
|
|
],
|
|
"status": "planning",
|
|
"current_step_index": 0
|
|
})
|
|
|
|
# 步骤 1: 分析上一集记忆（如果需要）
|
|
if analyze_previous_memory and episode_number > 1:
|
|
await safe_send("plan_update", {
|
|
"plan": [
|
|
f"分析 EP{episode_number - 1} 的记忆系统",
|
|
f"生成 EP{episode_number} 大纲",
|
|
f"创作 EP{episode_number} 对话内容",
|
|
f"执行质量审核",
|
|
f"更新记忆系统"
|
|
],
|
|
"status": "planning",
|
|
"current_step_index": 0
|
|
})
|
|
|
|
# 获取上一集内容
|
|
prev_episodes = await episode_repo.list_by_project(project_id)
|
|
prev_episode = next((ep for ep in prev_episodes if ep.number == episode_number - 1), None)
|
|
|
|
if prev_episode and prev_episode.content:
|
|
await safe_send("text", {"content": f"\n\n--- 正在分析 EP{episode_number - 1} 的记忆系统 ---\n"})
|
|
|
|
# 使用 MemoryManager 更新记忆
|
|
try:
|
|
memory_manager = get_memory_manager()
|
|
await memory_manager.update_memory_from_episode(project, prev_episode)
|
|
logger.info(f"EP{episode_number - 1} 记忆已分析并注入到 EP{episode_number}")
|
|
except Exception as e:
|
|
logger.warning(f"分析 EP{episode_number - 1} 记忆失败: {e}")
|
|
|
|
# 步骤 2-5: 执行剧集创作
|
|
agent = get_series_agent()
|
|
|
|
await safe_send("plan_update", {
|
|
"plan": [
|
|
f"分析 EP{episode_number - 1} 的记忆系统" if analyze_previous_memory and episode_number > 1 else "跳过记忆分析",
|
|
f"生成 EP{episode_number} 大纲",
|
|
f"创作 EP{episode_number} 对话内容",
|
|
f"执行质量审核",
|
|
f"更新记忆系统"
|
|
],
|
|
"status": "writing",
|
|
"current_step_index": 1
|
|
})
|
|
|
|
await safe_send("text", {"content": f"\n\n--- 开始创作 EP{episode_number} ---\n"})
|
|
|
|
# 执行创作
|
|
episode = await agent.execute_episode(
|
|
project=project,
|
|
episode_number=episode_number,
|
|
title=f"第{episode_number}集"
|
|
)
|
|
|
|
# 检查是否创作成功（错误处理）
|
|
if episode.status == "needs-review" and not episode.content:
|
|
# 创作失败，没有内容
|
|
await safe_send("error", {
|
|
"message": f"EP{episode_number} 创作失败",
|
|
"episode_number": episode_number
|
|
})
|
|
await safe_send("text", {"content": f"\n\n❌ EP{episode_number} 创作失败。请检查错误日志并重试。\n"})
|
|
logger.error(f"EP{episode_number} 创作失败，无内容生成")
|
|
return
|
|
|
|
# 创作成功，保存到数据库（无论 WebSocket 是否连接）
|
|
existing_episodes = await episode_repo.list_by_project(project_id)
|
|
episode_record = next((ep for ep in existing_episodes if ep.number == episode_number), None)
|
|
|
|
if episode_record:
|
|
episode.id = episode_record.id
|
|
episode.projectId = project_id
|
|
await episode_repo.update(episode)
|
|
logger.info(f"更新现有剧集记录: {episode.id}")
|
|
else:
|
|
episode.projectId = project_id
|
|
await episode_repo.create(episode)
|
|
logger.info(f"创建新剧集记录: {episode.id}")
|
|
|
|
# 发送内容到画布
|
|
if episode.content:
|
|
await safe_send("canvas_update", {"content": episode.content})
|
|
|
|
# 更新记忆
|
|
await project_repo.update(project_id, {
|
|
"memory": project.memory.dict()
|
|
})
|
|
|
|
# 完成消息
|
|
await safe_send("plan_update", {
|
|
"plan": [
|
|
f"分析 EP{episode_number - 1} 的记忆系统" if analyze_previous_memory and episode_number > 1 else "跳过记忆分析",
|
|
f"生成 EP{episode_number} 大纲",
|
|
f"创作 EP{episode_number} 对话内容",
|
|
f"执行质量审核",
|
|
f"更新记忆系统"
|
|
],
|
|
"status": "idle",
|
|
"current_step_index": 4
|
|
})
|
|
|
|
await safe_send("text", {"content": f"\n\n✅ EP{episode_number} 创作完成！质量分数: {episode.qualityScore or 0}\n"})
|
|
|
|
# 广播更新
|
|
await safe_send("episode_updated", {
|
|
"number": episode_number,
|
|
"title": episode.title,
|
|
"status": episode.status
|
|
})
|
|
|
|
logger.info(f"EP{episode_number} 后台创作完成，已保存到数据库")
|
|
|
|
except Exception as e:
|
|
logger.error(f"执行剧集创作失败: {str(e)}", exc_info=True)
|
|
try:
|
|
await manager.send_to_project(project_id, {
|
|
"type": "error",
|
|
"data": {
|
|
"message": f"EP{episode_number} 创作失败: {str(e)}",
|
|
"episode_number": episode_number
|
|
}
|
|
})
|
|
await manager.send_to_project(project_id, {
|
|
"type": "text",
|
|
"data": {"content": f"\n\n❌ EP{episode_number} 创作失败: {str(e)}\n"}
|
|
})
|
|
except Exception:
|
|
# WebSocket 可能已断开，忽略错误
|
|
pass
|
|
finally:
|
|
# 清理任务跟踪
|
|
if task_key in _background_tasks:
|
|
del _background_tasks[task_key]
|
|
logger.info(f"后台创作任务结束: {task_key}")
|
|
|
|
|
|
# ============================================
|
|
# 辅助函数 - 用于从其他模块发送消息
|
|
# ============================================
|
|
|
|
async def broadcast_stage_update(
|
|
project_id: str,
|
|
episode_number: int,
|
|
stage: str,
|
|
data: Dict[str, Any]
|
|
):
|
|
message = {
|
|
"type": "stage_update",
|
|
"data": {
|
|
"project_id": project_id,
|
|
"episode_number": episode_number,
|
|
"stage": stage,
|
|
**data
|
|
},
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
await manager.send_to_project(project_id, message)
|
|
|
|
|
|
async def broadcast_episode_complete(
|
|
project_id: str,
|
|
episode_number: int,
|
|
success: bool,
|
|
quality_score: float,
|
|
data: Dict[str, Any]
|
|
):
|
|
message = {
|
|
"type": "episode_complete",
|
|
"data": {
|
|
"project_id": project_id,
|
|
"episode_number": episode_number,
|
|
"success": success,
|
|
"quality_score": quality_score,
|
|
**data
|
|
},
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
await manager.send_to_project(project_id, message)
|
|
|
|
|
|
async def broadcast_batch_progress(
|
|
batch_id: str,
|
|
current_episode: int,
|
|
total_episodes: int,
|
|
completed: int,
|
|
failed: int,
|
|
data: Dict[str, Any]
|
|
):
|
|
message = {
|
|
"type": "batch_progress",
|
|
"data": {
|
|
"batch_id": batch_id,
|
|
"current_episode": current_episode,
|
|
"total_episodes": total_episodes,
|
|
"completed_episodes": completed,
|
|
"failed_episodes": failed,
|
|
"progress_percentage": (completed / total_episodes * 100) if total_episodes > 0 else 0,
|
|
**data
|
|
},
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
await manager.send_to_batch(batch_id, message)
|
|
|
|
|
|
async def broadcast_error(
|
|
project_id: str,
|
|
episode_number: Optional[int],
|
|
error: str,
|
|
error_type: str = "execution_error"
|
|
):
|
|
message = {
|
|
"type": "error",
|
|
"data": {
|
|
"project_id": project_id,
|
|
"episode_number": episode_number,
|
|
"error": error,
|
|
"error_type": error_type
|
|
},
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
await manager.send_to_project(project_id, message)
|
|
|
|
|
|
async def broadcast_batch_complete(
|
|
batch_id: str,
|
|
summary: Dict[str, Any]
|
|
):
|
|
message = {
|
|
"type": "batch_complete",
|
|
"data": {
|
|
"batch_id": batch_id,
|
|
**summary
|
|
},
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
await manager.send_to_batch(batch_id, message)
|
|
|
|
|
|
async def broadcast_to_project(
|
|
project_id: str,
|
|
message_type: str,
|
|
data: Dict[str, Any]
|
|
):
|
|
"""向项目的所有连接广播消息"""
|
|
message = {
|
|
"type": message_type,
|
|
"data": data,
|
|
"timestamp": datetime.now().isoformat()
|
|
}
|
|
await manager.send_to_project(project_id, message)
|
|
|
|
|
|
# 导出连接管理器和辅助函数
|
|
__all__ = [
|
|
"manager",
|
|
"broadcast_stage_update",
|
|
"broadcast_episode_complete",
|
|
"broadcast_batch_progress",
|
|
"broadcast_error",
|
|
"broadcast_batch_complete",
|
|
"broadcast_to_project"
|
|
]
|