433 lines
12 KiB
Python
433 lines
12 KiB
Python
"""
|
||
异步 AI 辅助生成 API
|
||
|
||
将原有的同步 AI 生成改为异步任务模式
|
||
"""
|
||
from fastapi import APIRouter, HTTPException, Depends, BackgroundTasks
|
||
from typing import Dict, Any, Optional, List
|
||
from pydantic import BaseModel
|
||
import asyncio
|
||
|
||
from app.core.llm.glm_client import get_glm_client
|
||
from app.core.skills.skill_manager import get_skill_manager
|
||
from app.core.task_manager import get_task_manager, TaskManager
|
||
from app.models.task import TaskType, TaskStatus
|
||
from app.api.v1.ai_assistant import build_enhanced_system_prompt
|
||
from app.utils.logger import get_logger
|
||
|
||
logger = get_logger(__name__)
|
||
|
||
router = APIRouter(prefix="/ai-assistant/async", tags=["AI 异步生成"])
|
||
|
||
|
||
# ============================================================================
|
||
# 请求模型
|
||
# ============================================================================
|
||
|
||
class SkillInfo(BaseModel):
|
||
"""Skill 信息（由前端传递）"""
|
||
id: str
|
||
name: str
|
||
behavior: str # behavior_guide 内容
|
||
|
||
|
||
# ============================================================================
|
||
# 任务执行器
|
||
# ============================================================================
|
||
|
||
async def execute_generate_characters(
|
||
task_manager: TaskManager,
|
||
task_id: str,
|
||
params: Dict[str, Any]
|
||
) -> Dict[str, Any]:
|
||
"""执行人物生成任务"""
|
||
try:
|
||
glm_client = get_glm_client()
|
||
skill_manager = get_skill_manager()
|
||
|
||
# 更新进度
|
||
task_manager.update_task_progress(
|
||
task_id, 10, 100,
|
||
"正在构建提示词..."
|
||
)
|
||
|
||
# 构建系统提示词
|
||
base_role = "你是专业的剧集创作专家，擅长创作丰富立体的人物角色。"
|
||
skills = params.get("skills", [])
|
||
system_prompt = await build_enhanced_system_prompt(
|
||
base_role=base_role,
|
||
skills=skills,
|
||
skill_manager=skill_manager
|
||
)
|
||
|
||
# 构建用户提示
|
||
extra_info = ""
|
||
if params.get("projectName"):
|
||
extra_info += f"\n项目名称：{params['projectName']}"
|
||
if params.get("totalEpisodes"):
|
||
extra_info += f"\n总集数：{params['totalEpisodes']}"
|
||
if params.get("genre"):
|
||
extra_info += f"\n类型：{params['genre']}"
|
||
|
||
custom_requirements = ""
|
||
if params.get("customPrompt"):
|
||
custom_requirements = f"\n【用户自定义要求】\n{params['customPrompt']}\n"
|
||
|
||
prompt = f"""请根据以下想法生成 3-5 个主要人物设定：
|
||
|
||
用户想法：{params['idea']}{extra_info}{custom_requirements}
|
||
|
||
要求：
|
||
1. 每个人物包含：姓名、身份、性格、说话风格、背景故事
|
||
2. 人物之间要有关系冲突
|
||
3. 每个人物 50-100 字
|
||
4. 格式：姓名：身份 - 性格 - 说话风格 - 背景故事
|
||
5. 严格遵守上面【应用技能指导】中的要求
|
||
|
||
请按以下格式输出：
|
||
【人物1】
|
||
姓名：xxx
|
||
身份：xxx
|
||
性格：xxx
|
||
说话风格：xxx
|
||
背景故事：xxx
|
||
【人物2】
|
||
...
|
||
"""
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 30, 100,
|
||
"正在调用 AI 生成..."
|
||
)
|
||
|
||
response = await glm_client.chat(
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
temperature=0.9
|
||
)
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 90, 100,
|
||
"正在处理结果..."
|
||
)
|
||
|
||
content = response["choices"][0]["message"]["content"]
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 100, 100,
|
||
"生成完成！"
|
||
)
|
||
|
||
return {
|
||
"success": True,
|
||
"characters": content,
|
||
"usage": response.get("usage")
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"人物生成失败: {str(e)}")
|
||
raise
|
||
|
||
|
||
async def execute_generate_outline(
|
||
task_manager: TaskManager,
|
||
task_id: str,
|
||
params: Dict[str, Any]
|
||
) -> Dict[str, Any]:
|
||
"""执行大纲生成任务"""
|
||
try:
|
||
glm_client = get_glm_client()
|
||
skill_manager = get_skill_manager()
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 10, 100,
|
||
"正在构建提示词..."
|
||
)
|
||
|
||
base_role = "你是专业的剧集创作专家，擅长构建引人入胜的剧情结构和故事节奏。"
|
||
skills = params.get("skills", [])
|
||
system_prompt = await build_enhanced_system_prompt(
|
||
base_role=base_role,
|
||
skills=skills,
|
||
skill_manager=skill_manager
|
||
)
|
||
|
||
custom_requirements = ""
|
||
if params.get("customPrompt"):
|
||
custom_requirements = f"\n【用户自定义要求】\n{params['customPrompt']}\n"
|
||
|
||
prompt = f"""请根据以下想法生成完整的剧集大纲：
|
||
|
||
用户想法：{params['idea']}
|
||
总集数：{params['totalEpisodes']}
|
||
类型：{params['genre']}
|
||
{f"项目名称：{params['projectName']}" if params.get('projectName') else ''}{custom_requirements}
|
||
|
||
要求：
|
||
1. 将故事分为 4-5 个阶段
|
||
2. 每个阶段包含具体的集数范围
|
||
3. 标注每个阶段的关键事件和转折点
|
||
4. 字数 200-400 字
|
||
5. 严格遵守上面【应用技能指导】中的要求
|
||
|
||
请按以下格式输出：
|
||
【阶段1】EPxx-EPxx：阶段名称
|
||
内容概要...
|
||
|
||
【阶段2】EPxx-EPxx：阶段名称
|
||
...
|
||
"""
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 30, 100,
|
||
"正在调用 AI 生成..."
|
||
)
|
||
|
||
response = await glm_client.chat(
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
temperature=0.9
|
||
)
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 90, 100,
|
||
"正在处理结果..."
|
||
)
|
||
|
||
content = response["choices"][0]["message"]["content"]
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 100, 100,
|
||
"生成完成！"
|
||
)
|
||
|
||
return {
|
||
"success": True,
|
||
"outline": content,
|
||
"usage": response.get("usage")
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"大纲生成失败: {str(e)}")
|
||
raise
|
||
|
||
|
||
async def execute_generate_world(
|
||
task_manager: TaskManager,
|
||
task_id: str,
|
||
params: Dict[str, Any]
|
||
) -> Dict[str, Any]:
|
||
"""执行世界观生成任务"""
|
||
try:
|
||
glm_client = get_glm_client()
|
||
skill_manager = get_skill_manager()
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 10, 100,
|
||
"正在构建提示词..."
|
||
)
|
||
|
||
base_role = "你是专业的世界观设定专家，擅长构建架空世界的背景设定。"
|
||
skills = params.get("skills", [])
|
||
system_prompt = await build_enhanced_system_prompt(
|
||
base_role=base_role,
|
||
skills=skills,
|
||
skill_manager=skill_manager
|
||
)
|
||
|
||
custom_requirements = ""
|
||
if params.get("customPrompt"):
|
||
custom_requirements = f"\n【用户自定义要求】\n{params['customPrompt']}\n"
|
||
|
||
prompt = f"""请根据以下想法生成世界观设定：
|
||
|
||
用户想法：{params['idea']}
|
||
类型：{params['genre']}
|
||
{f"项目名称：{params['projectName']}" if params.get('projectName') else ''}{custom_requirements}
|
||
|
||
要求：
|
||
1. 描述时代背景（朝代、架空世界等）
|
||
2. 描述地理环境和主要场景
|
||
3. 描述社会结构（权力体系、阶级关系）
|
||
4. 描述文化特色（习俗、服饰、语言等）
|
||
5. 字数 200-500 字
|
||
6. 严格遵守上面【应用技能指导】中的要求
|
||
|
||
请输出详细的世界观设定：
|
||
"""
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 30, 100,
|
||
"正在调用 AI 生成..."
|
||
)
|
||
|
||
response = await glm_client.chat(
|
||
messages=[
|
||
{"role": "system", "content": system_prompt},
|
||
{"role": "user", "content": prompt}
|
||
],
|
||
temperature=0.9
|
||
)
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 90, 100,
|
||
"正在处理结果..."
|
||
)
|
||
|
||
content = response["choices"][0]["message"]["content"]
|
||
|
||
task_manager.update_task_progress(
|
||
task_id, 100, 100,
|
||
"生成完成！"
|
||
)
|
||
|
||
return {
|
||
"success": True,
|
||
"worldSetting": content,
|
||
"usage": response.get("usage")
|
||
}
|
||
|
||
except Exception as e:
|
||
logger.error(f"世界观生成失败: {str(e)}")
|
||
raise
|
||
|
||
|
||
# ============================================================================
|
||
# 异步任务创建端点
|
||
# ============================================================================
|
||
|
||
class GenerateCharactersRequest(BaseModel):
|
||
"""生成人物设定请求（异步）"""
|
||
idea: str
|
||
projectName: Optional[str] = None
|
||
totalEpisodes: Optional[int] = None
|
||
genre: Optional[str] = "古风"
|
||
skills: Optional[List[SkillInfo]] = None
|
||
customPrompt: Optional[str] = None
|
||
projectId: Optional[str] = None # 关联项目ID
|
||
|
||
|
||
class GenerateOutlineRequest(BaseModel):
|
||
"""生成大纲请求（异步）"""
|
||
idea: str
|
||
totalEpisodes: int = 30
|
||
genre: str = "古风"
|
||
projectName: Optional[str] = None
|
||
skills: Optional[List[SkillInfo]] = None
|
||
customPrompt: Optional[str] = None
|
||
projectId: Optional[str] = None
|
||
|
||
|
||
class GenerateWorldRequest(BaseModel):
|
||
"""生成世界观请求（异步）"""
|
||
idea: str
|
||
projectName: Optional[str] = None
|
||
genre: Optional[str] = "古风"
|
||
skills: Optional[List[SkillInfo]] = None
|
||
customPrompt: Optional[str] = None
|
||
projectId: Optional[str] = None
|
||
|
||
|
||
@router.post("/generate/characters")
|
||
async def generate_characters_async(
|
||
request: GenerateCharactersRequest,
|
||
background_tasks: BackgroundTasks,
|
||
task_manager: TaskManager = Depends(get_task_manager)
|
||
):
|
||
"""
|
||
异步生成人物设定
|
||
|
||
返回任务ID，需要通过轮询 /tasks/{task_id} 获取结果
|
||
"""
|
||
# 创建任务
|
||
task = task_manager.create_task(
|
||
task_type=TaskType.GENERATE_CHARACTERS,
|
||
params=request.dict(exclude={"projectId"}),
|
||
project_id=request.projectId
|
||
)
|
||
|
||
# 在后台执行
|
||
async def run_task():
|
||
await task_manager.execute_task_async(
|
||
task.id,
|
||
lambda p: execute_generate_characters(task_manager, task.id, p)
|
||
)
|
||
|
||
asyncio.create_task(run_task())
|
||
|
||
return {
|
||
"success": True,
|
||
"taskId": task.id,
|
||
"message": "任务已创建，正在后台执行"
|
||
}
|
||
|
||
|
||
@router.post("/generate/outline")
|
||
async def generate_outline_async(
|
||
request: GenerateOutlineRequest,
|
||
task_manager: TaskManager = Depends(get_task_manager)
|
||
):
|
||
"""
|
||
异步生成大纲
|
||
|
||
返回任务ID，需要通过轮询 /tasks/{task_id} 获取结果
|
||
"""
|
||
# 创建任务
|
||
task = task_manager.create_task(
|
||
task_type=TaskType.GENERATE_OUTLINE,
|
||
params=request.dict(exclude={"projectId"}),
|
||
project_id=request.projectId
|
||
)
|
||
|
||
# 在后台执行
|
||
async def run_task():
|
||
await task_manager.execute_task_async(
|
||
task.id,
|
||
lambda p: execute_generate_outline(task_manager, task.id, p)
|
||
)
|
||
|
||
asyncio.create_task(run_task())
|
||
|
||
return {
|
||
"success": True,
|
||
"taskId": task.id,
|
||
"message": "任务已创建，正在后台执行"
|
||
}
|
||
|
||
|
||
@router.post("/generate/world")
|
||
async def generate_world_async(
|
||
request: GenerateWorldRequest,
|
||
task_manager: TaskManager = Depends(get_task_manager)
|
||
):
|
||
"""
|
||
异步生成世界观设定
|
||
|
||
返回任务ID，需要通过轮询 /tasks/{task_id} 获取结果
|
||
"""
|
||
# 创建任务
|
||
task = task_manager.create_task(
|
||
task_type=TaskType.GENERATE_WORLD,
|
||
params=request.dict(exclude={"projectId"}),
|
||
project_id=request.projectId
|
||
)
|
||
|
||
# 在后台执行
|
||
async def run_task():
|
||
await task_manager.execute_task_async(
|
||
task.id,
|
||
lambda p: execute_generate_world(task_manager, task.id, p)
|
||
)
|
||
|
||
asyncio.create_task(run_task())
|
||
|
||
return {
|
||
"success": True,
|
||
"taskId": task.id,
|
||
"message": "任务已创建，正在后台执行"
|
||
}
|