100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
|
#!/usr/bin/env python3
|
|||
|
|
# -*- coding: utf-8 -*-
|
|||
|
|
|
|||
|
|
"""
|
|||
|
|
海报文案内容生成器
|
|||
|
|
"""
|
|||
|
|
import logging
|
|||
|
|
import json
|
|||
|
|
from typing import Dict, Any, Optional, List
|
|||
|
|
|
|||
|
|
from core.ai import AIAgent
|
|||
|
|
|
|||
|
|
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 self._parse_json_response(raw_response)
|
|||
|
|
|
|||
|
|
except Exception as e:
|
|||
|
|
self.logger.error(f"调用AI生成文案时发生严重错误: {e}", exc_info=True)
|
|||
|
|
return None
|
|||
|
|
|
|||
|
|
def _parse_json_response(self, text: str) -> Optional[Dict[str, Any]]:
|
|||
|
|
"""
|
|||
|
|
从AI返回的文本中提取并解析JSON。
|
|||
|
|
"""
|
|||
|
|
try:
|
|||
|
|
# 找到第一个 '{' 和最后一个 '}' 来提取JSON字符串
|
|||
|
|
start_index = text.find('{')
|
|||
|
|
end_index = text.rfind('}')
|
|||
|
|
if start_index != -1 and end_index != -1:
|
|||
|
|
json_str = text[start_index : end_index + 1]
|
|||
|
|
return json.loads(json_str)
|
|||
|
|
else:
|
|||
|
|
self.logger.error("在AI响应中未找到有效的JSON对象。")
|
|||
|
|
return None
|
|||
|
|
except json.JSONDecodeError as e:
|
|||
|
|
self.logger.error(f"解析AI返回的JSON失败: {e}")
|
|||
|
|
self.logger.debug(f"原始响应文本: {text}")
|
|||
|
|
return None
|