488 lines
22 KiB
Python
488 lines
22 KiB
Python
# Link Agent Tools for ComfyUI Workflow Connection Analysis
|
|
|
|
import json
|
|
import time
|
|
from typing import Dict, Optional, List
|
|
|
|
from agents.tool import function_tool
|
|
from ..utils.request_context import get_session_id
|
|
from ..dao.workflow_table import get_workflow_data, save_workflow_data
|
|
from ..utils.comfy_gateway import get_object_info, get_object_info_by_missed_class_list
|
|
from ..utils.logger import log
|
|
|
|
@function_tool
|
|
async def analyze_missing_connections() -> str:
|
|
"""
|
|
分析工作流中缺失的连接，枚举所有可能的连接选项和所需的新节点。
|
|
|
|
返回格式说明：
|
|
- missing_connections: 缺失连接的详细列表，包含节点ID、输入名称、需要的数据类型等（仅包含required输入）
|
|
- possible_connections: 现有节点可以提供的连接选项
|
|
- universal_inputs: 可以接受任意输出类型的通用输入端口
|
|
- optional_unconnected_inputs: 未连接的可选输入列表，包含节点ID、输入名称、配置信息等
|
|
- required_new_nodes: 需要创建的新节点类型列表
|
|
- connection_summary: 连接分析的统计摘要（包含optional输入的统计）
|
|
"""
|
|
try:
|
|
session_id = get_session_id()
|
|
if not session_id:
|
|
log.error("analyze_missing_connections: No session_id found in context")
|
|
return json.dumps({"error": "No session_id found in context"})
|
|
|
|
workflow_data = get_workflow_data(session_id)
|
|
if not workflow_data:
|
|
return json.dumps({"error": "No workflow data found for this session"})
|
|
|
|
log.info(f"analyze_missing_connections: workflow_data: {workflow_data}")
|
|
|
|
object_info = await get_object_info()
|
|
|
|
log.info(f"analyze_missing_connections: object_info: {object_info}")
|
|
|
|
# 缺失的node class
|
|
missed_node_class = []
|
|
# 缺失的object info
|
|
missed_object_info = {}
|
|
for node_id, node_data in workflow_data.items():
|
|
node_class = node_data.get("class_type")
|
|
log.info(f"node: node_id:{node_id}, node_class: {node_class}")
|
|
if node_class not in object_info:
|
|
log.info(f"missed node: node_id:{node_id}, node_class: {node_class}")
|
|
missed_node_class.append(node_class)
|
|
|
|
log.info(f"analyze_missing_connections: missed_node_class: {missed_node_class}")
|
|
|
|
if missed_node_class:
|
|
missed_object_info = await get_object_info_by_missed_class_list(missed_node_class)
|
|
object_info.update(missed_object_info)
|
|
|
|
analysis_result = {
|
|
"missing_connections": [],
|
|
"possible_connections": [],
|
|
"universal_inputs": [], # 可以连接任意输出的输入端口
|
|
"optional_unconnected_inputs": [], # 未连接的optional输入
|
|
"required_new_nodes": [],
|
|
"connection_summary": {
|
|
"total_missing": 0,
|
|
"auto_fixable": 0,
|
|
"requires_new_nodes": 0,
|
|
"universal_inputs_count": 0,
|
|
"optional_unconnected_count": 0
|
|
}
|
|
}
|
|
|
|
# 构建现有节点的输出映射
|
|
available_outputs = {}
|
|
for node_id, node_data in workflow_data.items():
|
|
node_class = node_data.get("class_type")
|
|
if node_class in object_info and "output" in object_info[node_class]:
|
|
outputs = object_info[node_class]["output"]
|
|
available_outputs[node_id] = {
|
|
"class_type": node_class,
|
|
"outputs": [(i, output_type) for i, output_type in enumerate(outputs)]
|
|
}
|
|
|
|
# 分析每个节点的缺失连接
|
|
for node_id, node_data in workflow_data.items():
|
|
node_class = node_data.get("class_type")
|
|
if node_class not in object_info:
|
|
continue
|
|
|
|
node_info = object_info[node_class]
|
|
if "input" not in node_info:
|
|
continue
|
|
|
|
required_inputs = node_info["input"].get("required", {})
|
|
current_inputs = node_data.get("inputs", {})
|
|
|
|
# 检查每个required input
|
|
for input_name, input_config in required_inputs.items():
|
|
if input_name not in current_inputs:
|
|
# 发现缺失的连接
|
|
missing_connection = {
|
|
"node_id": node_id,
|
|
"node_class": node_class,
|
|
"input_name": input_name,
|
|
"input_config": input_config,
|
|
"required": True
|
|
}
|
|
|
|
# 分析输入类型
|
|
if isinstance(input_config, (list, tuple)) and len(input_config) > 0:
|
|
expected_types = input_config[0] if isinstance(input_config[0], list) else [input_config[0]]
|
|
|
|
missing_connection["expected_types"] = expected_types
|
|
|
|
# 检查是否是通用输入（可以接受任意类型）
|
|
is_universal_input = "*" in expected_types
|
|
|
|
if is_universal_input:
|
|
# 这是一个通用输入端口，可以连接任意输出
|
|
universal_input = {
|
|
"node_id": node_id,
|
|
"node_class": node_class,
|
|
"input_name": input_name,
|
|
"input_config": input_config,
|
|
"can_connect_any_output": True
|
|
}
|
|
analysis_result["universal_inputs"].append(universal_input)
|
|
analysis_result["connection_summary"]["universal_inputs_count"] += 1
|
|
analysis_result["connection_summary"]["auto_fixable"] += 1
|
|
|
|
# 对于通用输入，我们不列出所有可能的连接，而是标记为通用
|
|
missing_connection["possible_matches"] = "universal"
|
|
missing_connection["is_universal"] = True
|
|
else:
|
|
# 查找具体类型匹配的连接
|
|
possible_matches = []
|
|
for source_node_id, source_info in available_outputs.items():
|
|
if source_node_id == node_id: # 不能连接自己
|
|
continue
|
|
|
|
for output_index, output_type in source_info["outputs"]:
|
|
# 检查类型匹配（排除通用类型的情况）
|
|
type_match = (output_type in expected_types or output_type == "*")
|
|
|
|
if type_match:
|
|
possible_matches.append({
|
|
"source_node_id": source_node_id,
|
|
"source_class": source_info["class_type"],
|
|
"output_index": output_index,
|
|
"output_type": output_type,
|
|
"match_confidence": "high" if output_type in expected_types else "medium"
|
|
})
|
|
|
|
missing_connection["possible_matches"] = possible_matches
|
|
missing_connection["is_universal"] = False
|
|
|
|
# 如果有可能的匹配，添加到possible_connections
|
|
if possible_matches:
|
|
analysis_result["possible_connections"].extend([
|
|
{
|
|
"target_node_id": node_id,
|
|
"target_input": input_name,
|
|
"source_node_id": match["source_node_id"],
|
|
"source_output_index": match["output_index"],
|
|
"connection": [match["source_node_id"], match["output_index"]],
|
|
"confidence": match["match_confidence"],
|
|
"types": {
|
|
"expected": expected_types,
|
|
"provided": match["output_type"]
|
|
}
|
|
} for match in possible_matches
|
|
])
|
|
analysis_result["connection_summary"]["auto_fixable"] += 1
|
|
else:
|
|
# 没有匹配的输出，需要新节点
|
|
analysis_result["connection_summary"]["requires_new_nodes"] += 1
|
|
|
|
# 分析需要什么类型的节点
|
|
required_node_types = analyze_required_node_types(expected_types, object_info)
|
|
analysis_result["required_new_nodes"].extend([{
|
|
"for_node": node_id,
|
|
"for_input": input_name,
|
|
"expected_types": expected_types,
|
|
"suggested_node_types": required_node_types
|
|
}])
|
|
|
|
analysis_result["missing_connections"].append(missing_connection)
|
|
|
|
# 检查optional inputs (未连接的可选输入)
|
|
optional_inputs = node_info["input"].get("optional", {})
|
|
for input_name, input_config in optional_inputs.items():
|
|
if input_name not in current_inputs:
|
|
# 发现未连接的optional输入
|
|
optional_unconnected = {
|
|
"node_id": node_id,
|
|
"node_class": node_class,
|
|
"input_name": input_name,
|
|
"input_config": input_config,
|
|
"required": False
|
|
}
|
|
|
|
# 分析输入类型
|
|
if isinstance(input_config, (list, tuple)) and len(input_config) > 0:
|
|
expected_types = input_config[0] if isinstance(input_config[0], list) else [input_config[0]]
|
|
else:
|
|
expected_types = ["*"] # 默认为通用类型
|
|
|
|
optional_unconnected["expected_types"] = expected_types
|
|
|
|
# 检查是否是通用输入（可以接受任意类型）
|
|
is_universal_input = "*" in expected_types
|
|
optional_unconnected["is_universal"] = is_universal_input
|
|
|
|
analysis_result["optional_unconnected_inputs"].append(optional_unconnected)
|
|
|
|
# 更新统计信息
|
|
analysis_result["connection_summary"]["total_missing"] = len(analysis_result["missing_connections"])
|
|
analysis_result["connection_summary"]["optional_unconnected_count"] = len(analysis_result["optional_unconnected_inputs"])
|
|
|
|
log.info(f"analysis_result: {json.dumps(analysis_result, ensure_ascii=False)}")
|
|
return json.dumps(analysis_result)
|
|
|
|
except Exception as e:
|
|
return json.dumps({"error": f"Failed to analyze missing connections: {str(e)}"})
|
|
|
|
def analyze_required_node_types(expected_types: List[str], object_info: Dict) -> List[Dict]:
|
|
"""分析需要什么类型的节点来提供指定的输出类型"""
|
|
suggested_nodes = []
|
|
|
|
# 常见类型到节点的映射
|
|
type_to_nodes = {
|
|
"MODEL": ["CheckpointLoaderSimple", "CheckpointLoader", "UNETLoader"],
|
|
"CLIP": ["CheckpointLoaderSimple", "CheckpointLoader", "CLIPLoader"],
|
|
"VAE": ["CheckpointLoaderSimple", "CheckpointLoader", "VAELoader"],
|
|
"CONDITIONING": ["CLIPTextEncode"],
|
|
"LATENT": ["EmptyLatentImage", "VAEEncode"],
|
|
"IMAGE": ["LoadImage", "VAEDecode"],
|
|
"MASK": ["LoadImageMask"],
|
|
"CONTROL_NET": ["ControlNetLoader"],
|
|
"LORA": ["LoraLoader"],
|
|
"IPADAPTER": ["IPAdapterModelLoader"]
|
|
}
|
|
|
|
for expected_type in expected_types:
|
|
if expected_type in type_to_nodes:
|
|
for node_class in type_to_nodes[expected_type]:
|
|
if node_class in object_info:
|
|
suggested_nodes.append({
|
|
"node_class": node_class,
|
|
"output_type": expected_type,
|
|
"confidence": "high",
|
|
"description": f"{node_class} can provide {expected_type}"
|
|
})
|
|
else:
|
|
# 搜索所有能提供该类型输出的节点
|
|
for node_class, node_info in object_info.items():
|
|
if "output" in node_info:
|
|
outputs = node_info["output"]
|
|
if expected_type in outputs:
|
|
suggested_nodes.append({
|
|
"node_class": node_class,
|
|
"output_type": expected_type,
|
|
"confidence": "medium",
|
|
"description": f"{node_class} can provide {expected_type}"
|
|
})
|
|
|
|
# 去重并排序
|
|
unique_nodes = {}
|
|
for node in suggested_nodes:
|
|
key = node["node_class"]
|
|
if key not in unique_nodes or unique_nodes[key]["confidence"] == "medium":
|
|
unique_nodes[key] = node
|
|
|
|
return list(unique_nodes.values())
|
|
|
|
def save_checkpoint_before_link_modification(session_id: str, action_description: str) -> Optional[int]:
|
|
"""在连接修改前保存checkpoint"""
|
|
try:
|
|
current_workflow = get_workflow_data(session_id)
|
|
if not current_workflow:
|
|
return None
|
|
|
|
checkpoint_id = save_workflow_data(
|
|
session_id,
|
|
current_workflow,
|
|
workflow_data_ui=None,
|
|
attributes={
|
|
"checkpoint_type": "link_agent_start",
|
|
"description": f"Checkpoint before {action_description}",
|
|
"action": "link_agent_checkpoint",
|
|
"timestamp": time.time()
|
|
}
|
|
)
|
|
log.info(f"Saved link agent checkpoint with ID: {checkpoint_id}")
|
|
return checkpoint_id
|
|
except Exception as e:
|
|
log.error(f"Failed to save checkpoint before link modification: {str(e)}")
|
|
return None
|
|
|
|
@function_tool
|
|
def apply_connection_fixes(fixes_json: str) -> str:
|
|
"""批量应用连接修复，fixes_json应为包含修复指令的JSON字符串"""
|
|
try:
|
|
session_id = get_session_id()
|
|
if not session_id:
|
|
log.error("apply_connection_fixes: No session_id found in context")
|
|
return json.dumps({"error": "No session_id found in context"})
|
|
|
|
# 在修改前保存checkpoint
|
|
checkpoint_id = save_checkpoint_before_link_modification(session_id, "batch connection fixes")
|
|
|
|
# 解析修复指令
|
|
fixes = json.loads(fixes_json) if isinstance(fixes_json, str) else fixes_json
|
|
|
|
workflow_data = get_workflow_data(session_id)
|
|
if not workflow_data:
|
|
return json.dumps({"error": "No workflow data found for this session"})
|
|
|
|
applied_fixes = []
|
|
failed_fixes = []
|
|
|
|
# 处理连接修复
|
|
connections_to_add = fixes.get("connections", [])
|
|
for conn_fix in connections_to_add:
|
|
try:
|
|
target_node_id = conn_fix["target_node_id"]
|
|
target_input = conn_fix["target_input"]
|
|
source_node_id = conn_fix["source_node_id"]
|
|
source_output_index = conn_fix["source_output_index"]
|
|
|
|
if target_node_id not in workflow_data:
|
|
failed_fixes.append({
|
|
"type": "connection",
|
|
"target": f"{target_node_id}.{target_input}",
|
|
"error": f"Target node {target_node_id} not found"
|
|
})
|
|
continue
|
|
|
|
if source_node_id not in workflow_data:
|
|
failed_fixes.append({
|
|
"type": "connection",
|
|
"target": f"{target_node_id}.{target_input}",
|
|
"error": f"Source node {source_node_id} not found"
|
|
})
|
|
continue
|
|
|
|
# 应用连接
|
|
if "inputs" not in workflow_data[target_node_id]:
|
|
workflow_data[target_node_id]["inputs"] = {}
|
|
|
|
old_value = workflow_data[target_node_id]["inputs"].get(target_input, "not connected")
|
|
new_connection = [source_node_id, source_output_index]
|
|
workflow_data[target_node_id]["inputs"][target_input] = new_connection
|
|
|
|
applied_fixes.append({
|
|
"type": "connection",
|
|
"target": f"{target_node_id}.{target_input}",
|
|
"source": f"{source_node_id}[{source_output_index}]",
|
|
"old_value": old_value,
|
|
"new_value": new_connection
|
|
})
|
|
|
|
except Exception as e:
|
|
failed_fixes.append({
|
|
"type": "connection",
|
|
"target": f"{conn_fix.get('target_node_id', 'unknown')}.{conn_fix.get('target_input', 'unknown')}",
|
|
"error": str(e)
|
|
})
|
|
|
|
# 处理新节点添加
|
|
nodes_to_add = fixes.get("new_nodes", [])
|
|
for node_spec in nodes_to_add:
|
|
try:
|
|
node_class = node_spec["node_class"]
|
|
node_id = node_spec.get("node_id", "")
|
|
inputs = node_spec.get("inputs", {})
|
|
|
|
# 生成节点ID
|
|
if not node_id:
|
|
existing_ids = set(workflow_data.keys())
|
|
node_id = "1"
|
|
while node_id in existing_ids:
|
|
node_id = str(int(node_id) + 1)
|
|
|
|
# 创建新节点
|
|
new_node = {
|
|
"class_type": node_class,
|
|
"inputs": inputs,
|
|
"_meta": {"title": node_class}
|
|
}
|
|
|
|
workflow_data[node_id] = new_node
|
|
|
|
applied_fixes.append({
|
|
"type": "add_node",
|
|
"node_id": node_id,
|
|
"node_class": node_class,
|
|
"inputs": inputs
|
|
})
|
|
|
|
# 如果指定了要自动连接的目标，进行连接
|
|
auto_connect = node_spec.get("auto_connect", [])
|
|
for auto_conn in auto_connect:
|
|
target_node_id = auto_conn["target_node_id"]
|
|
target_input = auto_conn["target_input"]
|
|
output_index = auto_conn.get("output_index", 0)
|
|
|
|
if target_node_id in workflow_data:
|
|
if "inputs" not in workflow_data[target_node_id]:
|
|
workflow_data[target_node_id]["inputs"] = {}
|
|
|
|
workflow_data[target_node_id]["inputs"][target_input] = [node_id, output_index]
|
|
|
|
applied_fixes.append({
|
|
"type": "auto_connection",
|
|
"from_new_node": node_id,
|
|
"to": f"{target_node_id}.{target_input}",
|
|
"output_index": output_index
|
|
})
|
|
|
|
except Exception as e:
|
|
failed_fixes.append({
|
|
"type": "add_node",
|
|
"node_class": node_spec.get("node_class", "unknown"),
|
|
"error": str(e)
|
|
})
|
|
|
|
# 保存更新的工作流
|
|
version_id = save_workflow_data(
|
|
session_id,
|
|
workflow_data,
|
|
attributes={
|
|
"action": "link_agent_batch_fix",
|
|
"description": f"Applied {len(applied_fixes)} connection fixes",
|
|
"fixes_applied": applied_fixes,
|
|
"fixes_failed": failed_fixes
|
|
}
|
|
)
|
|
|
|
# 构建返回数据，包含checkpoint信息
|
|
ext_data = [{
|
|
"type": "workflow_update",
|
|
"data": {
|
|
"workflow_data": workflow_data,
|
|
"changes": {
|
|
"applied_fixes": applied_fixes,
|
|
"failed_fixes": failed_fixes
|
|
}
|
|
}
|
|
}]
|
|
|
|
# 如果成功保存了checkpoint，添加修改前的checkpoint信息
|
|
if checkpoint_id:
|
|
ext_data.append({
|
|
"type": "link_agent_checkpoint",
|
|
"data": {
|
|
"checkpoint_id": checkpoint_id,
|
|
"checkpoint_type": "link_agent_start"
|
|
}
|
|
})
|
|
|
|
# 添加修改后的版本信息
|
|
ext_data.append({
|
|
"type": "link_agent_complete",
|
|
"data": {
|
|
"version_id": version_id,
|
|
"checkpoint_type": "link_agent_complete"
|
|
}
|
|
})
|
|
|
|
return json.dumps({
|
|
"success": True,
|
|
"version_id": version_id,
|
|
"applied_fixes": applied_fixes,
|
|
"failed_fixes": failed_fixes,
|
|
"summary": {
|
|
"total_fixes": len(applied_fixes),
|
|
"failed_fixes": len(failed_fixes),
|
|
"connections_added": len([f for f in applied_fixes if f["type"] == "connection"]),
|
|
"nodes_added": len([f for f in applied_fixes if f["type"] == "add_node"])
|
|
},
|
|
"message": f"Applied {len(applied_fixes)} fixes successfully, {len(failed_fixes)} failed",
|
|
"ext": ext_data
|
|
})
|
|
|
|
except Exception as e:
|
|
return json.dumps({"error": f"Failed to apply connection fixes: {str(e)}"}) |