137 lines
4.5 KiB
Python
137 lines
4.5 KiB
Python
import re
|
|
|
|
|
|
class WorkflowWrapper(dict):
|
|
def __init__(self, workflow: dict):
|
|
super().__init__(workflow)
|
|
|
|
def get_tagged_nodes(self, tag: str = None):
|
|
tagged_nodes = []
|
|
|
|
for node_id, node_data in self.items():
|
|
title = node_data.get("_meta", {}).get("title", "")
|
|
if title:
|
|
found_tags = re.findall(r"\[(.*?)\]", title)
|
|
if found_tags:
|
|
tagged_nodes.append(
|
|
{"id": node_id, "node": node_data, "tags": found_tags}
|
|
)
|
|
|
|
if tag:
|
|
tagged_nodes = [n for n in tagged_nodes if tag in n["tags"]]
|
|
|
|
return tagged_nodes
|
|
|
|
def get_node_tags(self, node_id: str):
|
|
tagged_nodes = self.get_tagged_nodes()
|
|
for tagged_node in tagged_nodes:
|
|
if tagged_node["id"] == node_id:
|
|
return tagged_node["tags"]
|
|
|
|
return []
|
|
|
|
def get_tagged_inputs(self):
|
|
inputs_by_tag = {}
|
|
tagged_nodes = self.get_tagged_nodes()
|
|
|
|
for tagged_node in tagged_nodes:
|
|
node = tagged_node["node"]
|
|
tags = tagged_node["tags"]
|
|
|
|
for raw_tag in tags:
|
|
if raw_tag.startswith("!"):
|
|
continue
|
|
|
|
tag_name, filters = self._parse_tag(raw_tag)
|
|
|
|
if tag_name not in inputs_by_tag:
|
|
inputs_by_tag[tag_name] = {}
|
|
|
|
node_inputs = node.get("inputs", {})
|
|
for input_key, input_value in node_inputs.items():
|
|
if not filters or input_key in filters:
|
|
inputs_by_tag[tag_name][input_key] = type(input_value).__name__
|
|
|
|
return inputs_by_tag
|
|
|
|
def update_tagged_nodes_input(
|
|
self, tag: str, input_key: str, new_value: any
|
|
) -> None:
|
|
tagged_inputs = self.get_tagged_inputs()
|
|
|
|
if tag not in tagged_inputs:
|
|
raise ValueError(
|
|
f"No inputs available for tag [{tag}]. Tag does not exist, or is excluded with '!'."
|
|
)
|
|
|
|
if input_key not in tagged_inputs[tag]:
|
|
raise ValueError(
|
|
f"The input '{input_key}' does not exist for tag [{tag}] "
|
|
f"(perhaps it's not in the filtered list?)."
|
|
)
|
|
|
|
tagged_nodes = self.get_tagged_nodes(tag)
|
|
if not tagged_nodes:
|
|
raise ValueError(f"Could not update node, no node with tag [{tag}] found.")
|
|
|
|
for tagged_node in tagged_nodes:
|
|
node_id = tagged_node["id"]
|
|
node = self[node_id]
|
|
node_inputs = node.get("inputs", {})
|
|
|
|
if input_key not in node_inputs:
|
|
raise ValueError(
|
|
f"Could not update node, no input '{input_key}' in node with tag "
|
|
f"[{tag}] (node found at id {node_id})."
|
|
)
|
|
|
|
node_inputs[input_key] = new_value
|
|
|
|
def bypass_nodes(self, tag: str) -> None:
|
|
tagged_nodes = self.get_tagged_nodes(tag)
|
|
|
|
for skip_node_data in tagged_nodes:
|
|
skip_node_id = skip_node_data["id"]
|
|
if not skip_node_id:
|
|
continue
|
|
|
|
skip_node = self[skip_node_id]
|
|
|
|
# Supprime le noeud du workflow
|
|
del self[skip_node_id]
|
|
|
|
target_node_id = None
|
|
for node_id, node_data in self.items():
|
|
for input_value in node_data.get("inputs", {}).values():
|
|
if isinstance(input_value, list) and len(input_value) > 0:
|
|
if input_value[0] == skip_node_id:
|
|
target_node_id = node_id
|
|
break
|
|
if target_node_id:
|
|
break
|
|
|
|
if not target_node_id:
|
|
continue
|
|
|
|
target_node = self[target_node_id]
|
|
for target_input_key in target_node.get("inputs", {}):
|
|
singular_input_key = target_input_key[:-1]
|
|
|
|
if singular_input_key in skip_node.get("inputs", {}):
|
|
target_node["inputs"][target_input_key] = skip_node["inputs"][
|
|
singular_input_key
|
|
]
|
|
break
|
|
|
|
@staticmethod
|
|
def _parse_tag(tag_str: str):
|
|
if ":" in tag_str:
|
|
parts = tag_str.split(":", 1)
|
|
tag_name = parts[0].strip()
|
|
filters_str = parts[1].strip()
|
|
filters = [f.strip() for f in filters_str.split(",") if f.strip()]
|
|
else:
|
|
tag_name = tag_str.strip()
|
|
filters = []
|
|
return tag_name, filters
|