""" 提示词增强节点 - V3 版本 V3 迁移说明： - 继承 LLMNodeBase（工具基类 Mixin）+ io.ComfyNode（V3 节点基类） - INPUT_TYPES → define_schema()，返回 io.Schema - IS_CHANGED → fingerprint_inputs() - def enhance(self, ...) → @classmethod execute(cls, ...) - 返回 io.NodeOutput(val) 代替 (val,) - hidden unique_id 通过 cls.hidden.unique_id 访问 - 不再导出 NODE_CLASS_MAPPINGS，由顶层 __init__.py 的 ComfyExtension 统一注册 """ import hashlib from comfy.model_management import InterruptProcessingException from comfy_api.latest import io from ..services.llm import LLMService from ..utils.common import ( format_api_error, format_model_with_thinking, generate_request_id, log_prepare, log_error, TASK_EXPAND, SOURCE_NODE ) from ..services.thinking_control import build_thinking_suppression from .base import LLMNodeBase class PromptExpand(LLMNodeBase, io.ComfyNode): """ 提示词增强节点（V3） - 输入 "source_text"，根据所选规则模板或自定义规则进行增强/扩写 - 仅包含一个字符串输入和一个字符串输出 """ @classmethod def define_schema(cls): """定义节点 Schema（V3 替代 INPUT_TYPES + 类属性）""" # 从 config_manager 获取系统提示词配置 from ..config_manager import config_manager system_prompts = config_manager.get_system_prompts() # 获取所有 expand_prompts 作为下拉选项 expand_prompts = {} active_expand_id = None if system_prompts: expand_prompts = system_prompts.get('expand_prompts', {}) or {} active_expand_id = system_prompts.get('active_prompts', {}).get('expand') # 构建提示词模板选项（支持分类格式：类别/规则名称） prompt_template_options = [] id_to_display_name = {} for key, value in expand_prompts.items(): # 过滤掉不在后端显示的规则 show_in = value.get('showIn', ["frontend", "node"]) if 'node' not in show_in: continue name = value.get('name', key) category = value.get('category', '') display_name = f"{category}/{name}" if category else name id_to_display_name[key] = display_name prompt_template_options.append(display_name) # 默认选项回退 default_template_name = prompt_template_options[0] if prompt_template_options else "扩写-自然语言" if active_expand_id and active_expand_id in id_to_display_name: default_template_name = id_to_display_name[active_expand_id] if not prompt_template_options: prompt_template_options = ["扩写-自然语言"] # ---动态获取 LLM 服务/模型列表--- service_options = cls.get_llm_service_options() default_service = service_options[0] if service_options else "智谱" return io.Schema( node_id="PromptExpand", display_name="✨Prompt Enhance", category="✨Prompt Assistant", description="Enhance and expand prompts using LLM services", inputs=[ # 规则模板：来自系统配置的所有扩写规则 io.Combo.Input( "rule", options=prompt_template_options, default=default_template_name, tooltip="Choose a preset rule for prompt enhancement", ), # 临时规则开关 io.Boolean.Input( "custom_rule", default=False, label_on="Enable", label_off="Disable", tooltip="Enable to use custom rule content below instead of preset", ), # 临时规则内容输入框 io.String.Input( "custom_rule_content", multiline=True, default="", placeholder="Enter custom rule here, only effective when 'Custom Rule' is enabled", tooltip="Enter your custom rule content here", ), # 用户提示词 io.String.Input( "user_prompt", multiline=True, default="", placeholder="Enter the prompt to enhance; if source_text is also connected, both will be merged", tooltip="The original prompt to enhance", ), # 扩写服务 io.Combo.Input( "llm_service", options=service_options, default=default_service, tooltip="Select LLM service and model", ), # Ollama 自动释放显存 io.Boolean.Input( "ollama_auto_unload", default=True, label_on="Enable", label_off="Disable", tooltip="Auto unload Ollama model after generation", ), io.Int.Input( "seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True, ), # 原文输入端口（可选），默认为连接端口 io.String.Input( "source_text", optional=True, multiline=True, default="", force_input=True, placeholder="Input text to enhance...", tooltip="Optional input text", ), ], outputs=[ io.String.Output("enhanced_text"), ], hidden=[io.Hidden.unique_id], ) @classmethod def fingerprint_inputs( cls, rule=None, custom_rule=None, custom_rule_content=None, user_prompt=None, llm_service=None, ollama_auto_unload=None, seed=None, source_text=None ): """ 替代 V1 IS_CHANGED，只在输入内容真正变化时才触发重新执行 使用输入参数的哈希值作为判断依据 """ text_hash = hashlib.md5(((source_text or "")).encode('utf-8')).hexdigest() temp_rule_hash = hashlib.md5((custom_rule_content or "").encode('utf-8')).hexdigest() user_hint_hash = hashlib.md5((user_prompt or "").encode('utf-8')).hexdigest() input_hash = hash(( rule, bool(custom_rule), temp_rule_hash, user_hint_hash, llm_service, bool(ollama_auto_unload), seed, text_hash, )) return input_hash @classmethod def execute( cls, rule, custom_rule, custom_rule_content, user_prompt, llm_service, ollama_auto_unload, seed=None, source_text=None ): """ 增强/扩写文本函数（V3 classmethod 版本） 通过 cls.hidden.unique_id 访问节点唯一 ID """ # 从 cls.hidden 获取节点唯一 ID unique_id = cls.hidden.unique_id request_id = None try: # 允许原文为空，但原文与用户提示词至少有一项非空 source_text = (source_text or "").strip() user_prompt = (user_prompt or "").strip() if not source_text and not user_prompt: return io.NodeOutput("") # ---准备系统提示词（规则）--- system_message = None rule_name = "Custom Rule" if (custom_rule and custom_rule_content) else rule if custom_rule and custom_rule_content: # 使用临时规则 system_message = {"role": "system", "content": custom_rule_content} else: # 使用模板：从 config_manager 获取系统提示词配置 from ..config_manager import config_manager system_prompts = config_manager.get_system_prompts() expand_prompts = system_prompts.get('expand_prompts', {}) if system_prompts else {} # 查找选定的提示词模板（按显示名称匹配） template_found = False for key, value in expand_prompts.items(): name = value.get('name', key) category = value.get('category', '') display_name = f"{category}/{name}" if category else name if display_name == rule: system_message = {"role": value.get('role', 'system'), "content": value.get('content', '')} template_found = True break if not template_found: # 允许用规则名称或键名直接匹配（兼容旧格式） for key, value in expand_prompts.items(): if value.get('name') == rule or key == rule: system_message = {"role": value.get('role', 'system'), "content": value.get('content', '')} template_found = True break if not template_found or not system_message or not system_message.get('content'): # 回退到默认 system_message = {"role": "system", "content": "你是一名提示词扩写专家，请将用户给定文本扩写为更完整、更具可读性和可执行性的提示词。"} rule_name = "Default Rule" # ---解析服务/模型字符串--- service_id, model_name = cls.parse_service_model(llm_service) if not service_id: raise ValueError(f"Invalid service selection: {llm_service}") # ---获取服务配置--- from ..config_manager import config_manager service = config_manager.get_service(service_id) if not service: raise ValueError(f"Service config not found: {llm_service}") # ---构建 provider_config--- llm_models = service.get('llm_models', []) target_model = None if model_name: target_model = next((m for m in llm_models if m.get('name') == model_name), None) if not target_model: target_model = next( (m for m in llm_models if m.get('is_default')), llm_models[0] if llm_models else None ) if not target_model: raise ValueError(f"Service {llm_service} has no available models") provider_config = { 'provider': service_id, 'model': target_model.get('name', ''), 'base_url': service.get('base_url', ''), 'api_key': service.get('api_key', ''), 'temperature': target_model.get('temperature', 0.7), 'max_tokens': target_model.get('max_tokens', 1000), 'top_p': target_model.get('top_p', 0.9), } # Ollama 特殊处理：添加 auto_unload 配置 if service.get('type') == 'ollama': provider_config['auto_unload'] = ollama_auto_unload # 生成请求 ID request_id = generate_request_id("exp", None, unique_id) # 合并原文与用户提示词（输入端口在前，节点内文本框在后） combined_text = ( user_prompt if not source_text else (f"{source_text}\n\n{user_prompt}" if user_prompt else source_text) ) # 检查是否关闭思维链 model_full_name = provider_config.get('model') disable_thinking_enabled = service.get('disable_thinking', True) thinking_extra = build_thinking_suppression(service_id, model_full_name) if disable_thinking_enabled else None model_display = format_model_with_thinking(model_full_name, bool(thinking_extra)) # 获取服务显示名称 service_display_name = service.get('name', service_id) # 准备阶段日志 log_prepare(TASK_EXPAND, request_id, SOURCE_NODE, service_display_name, model_display, rule_name, {"长度": len(combined_text)}) # 检查 API 密钥和模型 if not provider_config.get('model', ''): raise ValueError(f"Please configure model for {llm_service}") if cls._service_requires_api_key(service) and not provider_config.get('api_key', ''): raise ValueError(f"Please configure API key and model for {llm_service}") # 执行扩写（异步线程 + 可中断） result = cls._run_llm_task( LLMService.expand_prompt, service_id, prompt=combined_text, request_id=request_id, stream_callback=None, custom_provider=service_id, custom_provider_config=provider_config, system_message_override=system_message, task_type=TASK_EXPAND, source=SOURCE_NODE ) if result and result.get('success'): expanded_text = result.get('data', {}).get('expanded', '').strip() if not expanded_text: error_msg = 'API returned empty result' log_error(TASK_EXPAND, request_id, error_msg, source=SOURCE_NODE) raise RuntimeError(f"Enhancement failed: {error_msg}") # 结果阶段日志由服务层统一输出，节点层不再重复打印 return io.NodeOutput(expanded_text) else: error_msg = result.get('error', 'Unknown error') if result else 'No result returned' if error_msg == "任务被中断": raise InterruptProcessingException() log_error(TASK_EXPAND, request_id, error_msg, source=SOURCE_NODE) raise RuntimeError(f"Enhancement failed: {error_msg}") except InterruptProcessingException: # 不打印日志，由基类统一打印 raise except Exception as e: error_msg = format_api_error(e, llm_service) log_error(TASK_EXPAND, request_id, error_msg, source=SOURCE_NODE) raise RuntimeError(f"Enhancement error: {error_msg}")