1+ '''
2+ Author: ai-business-hql [email protected] 3+ Date: 2025-08-25 20:16:18
4+ LastEditors: ai-business-hql [email protected] 5+ LastEditTime: 2025-08-25 21:08:36
6+ FilePath: /ComfyUI-Copilot/backend/service/workflow_rewrite_agent_simple.py
7+ Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
8+ '''
9+
10+ import json
11+ from typing import Dict , Any , Optional
12+ import asyncio
13+ from pydantic import BaseModel
14+ from openai import OpenAI
15+
16+ from ..agent_factory import create_agent
17+ from ..utils .globals import WORKFLOW_MODEL_NAME , get_comfyui_copilot_api_key , LLM_DEFAULT_BASE_URL
18+ from ..utils .request_context import get_rewrite_context , RewriteContext
19+ from ..utils .logger import log
20+
21+
22+ class RewriteResponse (BaseModel ):
23+ """
24+ 重写后的工作流数据，必须是严格的JSON字符串
25+ """
26+ workflow_data : str
27+
28+ def rewrite_workflow_simple (rewrite_context : RewriteContext ) -> str :
29+ """
30+ 使用简化的方式重写工作流，直接调用OpenAI API
31+
32+ Args:
33+ rewrite_context: 包含所有重写所需信息的上下文
34+
35+ Returns:
36+ 改写后的API工作流(JSON字符串)
37+ """
38+ try :
39+ # 构建给LLM的完整上下文信息
40+ context_info = f"""
41+ ## 重写意图
42+ { rewrite_context .rewrite_intent }
43+
44+ ## 当前工作流 (ComfyUI API格式)
45+ { rewrite_context .current_workflow }
46+
47+ ## 节点信息
48+ { json .dumps (rewrite_context .node_infos or {}, ensure_ascii = False )}
49+
50+ ## 专家经验
51+ { rewrite_context .rewrite_expert or "无特定专家经验" }
52+ """
53+
54+ # 构建系统提示词
55+ system_prompt = """你是专业的ComfyUI工作流重写专家。你需要根据用户的重写意图，对现有的ComfyUI API格式工作流进行改写。
56+
57+ ## ComfyUI工作流基础知识
58+ - ComfyUI工作流是基于节点的图形化系统，每个节点代表一个操作
59+ - 节点通过输入输出端口进行数据流转，形成有向无环图(DAG)
60+ - 节点的输入输出类型必须严格匹配(如image, latent, model, string, int, float等)
61+ - 每个节点的必需输入必须有有效连接，否则会导致工作流运行失败
62+
63+ ## API格式结构
64+ ComfyUI API格式工作流是一个JSON对象，其中：
65+ - 键是节点ID (字符串形式的数字，如"1", "2", "3")
66+ - 值是节点对象，包含：
67+ - class_type: 节点类型
68+ - inputs: 输入参数对象，可以是：
69+ - 直接值(字符串、数字等)
70+ - 连接引用 [node_id, output_index] 形式的数组
71+
72+ ## 重写原则
73+ 1. **保持结构完整性**：确保所有节点的必需输入都有连接
74+ 2. **类型匹配**：输入输出类型必须严格匹配
75+ 3. **连接正确性**：确保节点间的连接关系正确
76+ 4. **功能实现**：确保改写后的工作流能实现用户的重写意图
77+
78+ ## 重要：输出格式要求
79+ 你必须严格按照以下JSON格式返回结果，不要添加任何其他说明文字：
80+ {
81+ "workflow_data": "这里是完整的工作流JSON字符串"
82+ }
83+ """
84+
85+ # 创建OpenAI客户端
86+ client = OpenAI (
87+ base_url = LLM_DEFAULT_BASE_URL ,
88+ api_key = get_comfyui_copilot_api_key () or ""
89+ )
90+
91+ # 调用LLM
92+ completion = client .chat .completions .parse (
93+ model = WORKFLOW_MODEL_NAME ,
94+ messages = [
95+ {"role" : "system" , "content" : system_prompt },
96+ {"role" : "user" , "content" : context_info }
97+ ],
98+ max_tokens = 8192 ,
99+ temperature = 0.1 , # 降低随机性，确保输出一致性
100+ response_format = RewriteResponse # 要求返回JSON格式
101+ )
102+
103+ result = completion .choices [0 ].message .parsed
104+ log .info (f"workflow simple rewrite LLM response: { result } )
105+
106+ # 解析返回的JSON
107+ # result = json.loads(result_text)
108+
109+ # 验证返回格式
110+ if result .workflow_data is None :
111+ return "{}"
112+
113+ # 验证workflow_data是有效的JSON字符串
114+ if isinstance (result .workflow_data , str ):
115+ # 尝试解析以验证有效性
116+ json .loads (result .workflow_data )
117+ return result .workflow_data
118+ else :
119+ # 如果不是字符串，转换为字符串
120+ return json .dumps (result .workflow_data , ensure_ascii = False )
121+
122+ except Exception as e :
123+ log .error (f"简化工作流重写失败: { str (e )} )
124+ return f"工作流重写失败:{ str (e )}
0 commit comments