#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ 选题生成引擎 V2 - 不访问数据库，接收完整数据 - 使用 PromptRegistry 管理 prompt - 统一依赖注入 """ import logging from typing import Dict, Any, Optional, List from .base import BaseAIGCEngine, EngineResult logger = logging.getLogger(__name__) class TopicGenerateEngineV2(BaseAIGCEngine): """ 选题生成引擎 V2 改进: 1. 不访问数据库，所有数据由调用方传入 2. 使用 PromptRegistry 管理 prompt 3. 接收完整对象而非 ID """ engine_id = "topic_generate" engine_name = "选题生成" version = "2.0.0" description = "根据景区、产品、风格等信息生成营销选题" def __init__(self): super().__init__() self._prompt_registry = None def get_param_schema(self) -> Dict[str, Any]: """ 定义参数结构 V2.2: 合并 scenic_spot 和 product 为 subject """ return { # 基础参数 "num_topics": { "type": "int", "required": False, "default": 5, "desc": "生成选题数量", }, "month": { "type": "str", "required": True, "desc": "目标日期/月份 (如 '2024-12' 或 '12月5日')", }, # 主体信息 (景区+产品合并) "subject": { "type": "object", "required": False, "desc": "主体信息 {id, name, type, description, location, products: [...]}", }, # 兼容旧字段 "scenic_spot": { "type": "object", "required": False, "desc": "[兼容] 景区信息，建议使用 subject", }, "product": { "type": "object", "required": False, "desc": "[兼容] 产品信息，建议使用 subject.products", }, "style": { "type": "object", "required": False, "desc": "风格信息对象 {id, name}", }, "audience": { "type": "object", "required": False, "desc": "受众信息对象 {id, name}", }, # 热点信息 "hot_topics": { "type": "object", "required": False, "desc": "热点信息 {events: [], festivals: [], trending: []}", }, # 可选: 多选列表 "styles_list": { "type": "list", "required": False, "desc": "可选风格列表 [{id, name}, ...]", }, "audiences_list": { "type": "list", "required": False, "desc": "可选受众列表 [{id, name}, ...]", }, # Prompt 版本控制 "prompt_version": { "type": "str", "required": False, "default": "latest", "desc": "使用的 prompt 版本", }, } def estimate_duration(self, params: Dict[str, Any]) -> int: """预估执行时间""" num_topics = params.get('num_topics', 5) return 15 + num_topics * 2 async def execute(self, params: Dict[str, Any], task_id: str = None) -> EngineResult: """ 执行选题生成 Args: params: 包含完整对象的参数 task_id: 任务 ID """ try: self.log(f"开始生成选题 (V2)") self.set_progress(task_id, 10) # 提取参数 num_topics = params.get('num_topics', 5) month = params.get('month', '') # 主体信息 (支持新旧两种格式) subject = params.get('subject') scenic_spot = params.get('scenic_spot') product = params.get('product') # 兼容处理: 如果没有 subject，从 scenic_spot + product 构建 if not subject and scenic_spot: subject = { **scenic_spot, 'type': scenic_spot.get('type', 'scenic_spot'), 'products': [product] if product else [] } style = params.get('style') audience = params.get('audience') hot_topics = params.get('hot_topics') styles_list = params.get('styles_list', []) audiences_list = params.get('audiences_list', []) prompt_version = params.get('prompt_version', 'latest') self.set_progress(task_id, 20) # 获取 PromptRegistry prompt_registry = self._get_prompt_registry() # 🎯 方案A: 检查是否有 Java 透传的 prompt_context prompt_context = params.get('prompt_context') if prompt_context: self.log("使用 Java 透传的 prompt_context") context = prompt_context # 补充必要字段 if 'num_topics' not in context: context['num_topics'] = num_topics if 'month' not in context: context['month'] = month if 'hot_topics' not in context and hot_topics: context['hot_topics'] = hot_topics else: self.log("使用本地构建的 context (兼容模式)") # 构建 prompt 上下文 context = { 'num_topics': num_topics, 'month': month, 'subject': subject, # 兼容旧 prompt 模板 'scenic_spot': subject, 'product': subject.get('products', [{}])[0] if subject and subject.get('products') else product, 'style': style, 'audience': audience, 'hot_topics': hot_topics, 'styles_list': self._format_list(styles_list), 'audiences_list': self._format_list(audiences_list), } # 渲染 prompt system_prompt, user_prompt = prompt_registry.render( 'topic_generate', context=context, version=prompt_version ) self.set_progress(task_id, 30) # 获取模型参数 prompt_config = prompt_registry.get('topic_generate', prompt_version) model_params = prompt_config.get_model_params() # 调用 LLM self.log("调用 LLM 生成选题...") raw_result, input_tokens, output_tokens, time_cost = await self.llm.generate( system_prompt=system_prompt, user_prompt=user_prompt, **model_params ) self.set_progress(task_id, 70) # 解析结果 topics = self._parse_topics(raw_result) if not topics: return EngineResult( success=False, error="选题生成失败，无法解析结果", error_code="PARSE_ERROR" ) self.set_progress(task_id, 90) # 增强选题信息 (添加原始对象引用) enhanced_topics = self._enhance_topics( topics, subject, style, audience ) self.set_progress(task_id, 100) return EngineResult( success=True, data={ "topics": enhanced_topics, "count": len(enhanced_topics), }, metadata={ "input_tokens": input_tokens, "output_tokens": output_tokens, "time_cost": time_cost, "prompt_version": prompt_version, } ) except Exception as e: self.log(f"选题生成异常: {e}", level='error') return EngineResult( success=False, error=str(e), error_code="EXECUTION_ERROR" ) def _get_prompt_registry(self): """获取 PromptRegistry""" if self._prompt_registry: return self._prompt_registry from domain.prompt import PromptRegistry self._prompt_registry = PromptRegistry('prompts') return self._prompt_registry def _format_list(self, items: List[Dict]) -> str: """格式化列表为字符串""" if not items: return "" lines = [] for item in items: name = item.get('name', '') desc = item.get('description', '') if name: lines.append(f"- {name}: {desc}" if desc else f"- {name}") return "\n".join(lines) def _parse_topics(self, raw_result: str) -> List[Dict[str, Any]]: """解析 LLM 返回的选题""" import json import re # 尝试提取 JSON json_match = re.search(r'\[[\s\S]*\]', raw_result) if json_match: try: return json.loads(json_match.group()) except json.JSONDecodeError: pass # 尝试 json_repair try: import json_repair return json_repair.loads(raw_result) except: pass self.log("无法解析选题结果", level='error') return [] def _enhance_topics(self, topics: List[Dict], subject: Optional[Dict], style: Optional[Dict], audience: Optional[Dict]) -> List[Dict]: """增强选题信息""" enhanced = [] for topic in topics: enhanced_topic = dict(topic) # 添加原始对象 ID 引用 if subject: enhanced_topic['subject_id'] = subject.get('id') # 从 products 中提取第一个产品的 ID (如果有) products = subject.get('products', []) if products and len(products) > 0: enhanced_topic['product_id'] = products[0].get('id') if style: enhanced_topic['style_id'] = style.get('id') if audience: enhanced_topic['audience_id'] = audience.get('id') enhanced.append(enhanced_topic) return enhanced