bangbang-aigc-server/poster/text_generator.py

81 lines
2.7 KiB
Python
Raw Permalink Normal View History

2025-07-31 15:35:23 +08:00
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
海报文案内容生成器
"""
import logging
from typing import Dict, Any, Optional, List
from core.ai import AIAgent
from utils.file_io import ResourceLoader, process_llm_json_text
logger = logging.getLogger(__name__)
class PosterContentGenerator:
"""
使用AI模型为海报生成文本内容
"""
def __init__(self, ai_agent: AIAgent):
"""
初始化内容生成器
Args:
ai_agent (AIAgent): 用于与AI模型交互的代理
"""
self.ai_agent = ai_agent
self.logger = logging.getLogger(__name__)
async def generate_text_for_poster(
self,
system_prompt: str,
user_prompt: str,
context_data: Optional[Dict[str, Any]] = None,
temperature: Optional[float] = None,
top_p: Optional[float] = None,
) -> Optional[Dict[str, Any]]:
"""
为单个海报任务生成文本内容
Args:
system_prompt (str): 提供给AI的系统级指令
user_prompt (str): 提供给AI的用户级指令可包含占位符
context_data (Optional[Dict[str, Any]]): 用于填充用户指令占位符的数据
temperature (Optional[float]): AI模型温度参数
top_p (Optional[float]): AI模型top_p参数
Returns:
Optional[Dict[str, Any]]: 解析后的JSON对象包含生成的文本如果失败则返回None
"""
if context_data:
try:
# 使用上下文数据格式化用户提示
final_user_prompt = user_prompt.format(**context_data)
except KeyError as e:
self.logger.error(f"格式化用户提示失败,缺少键: {e}")
return None
else:
final_user_prompt = user_prompt
self.logger.info("正在调用AI生成海报文案...")
self.logger.debug(f"System Prompt: {system_prompt[:200]}...")
self.logger.debug(f"User Prompt: {final_user_prompt[:200]}...")
try:
raw_response, _, _, _ = await self.ai_agent.generate_text(
system_prompt=system_prompt,
user_prompt=final_user_prompt,
temperature=temperature,
top_p=top_p
)
if not raw_response:
self.logger.error("AI未能返回任何内容。")
return None
# 使用通用JSON解析函数处理响应
return process_llm_json_text(raw_response)
except Exception as e:
self.logger.error(f"调用AI生成文案时发生严重错误: {e}", exc_info=True)
return None