128 lines
4.5 KiB
Python
128 lines
4.5 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
"""Compute pickable upstream paths for DataPicker / AI workflow tools."""
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Dict, List, Set
|
|
|
|
from modules.features.graphicalEditor.nodeDefinitions import STATIC_NODE_TYPES
|
|
from modules.features.graphicalEditor.portTypes import PORT_TYPE_CATALOG, PortSchema, parse_graph_defined_output_schema
|
|
from modules.workflows.automation2.graphUtils import buildConnectionMap
|
|
|
|
_NODE_BY_TYPE = {n["id"]: n for n in STATIC_NODE_TYPES}
|
|
|
|
|
|
def _paths_for_port_schema(schema: PortSchema, producer_node_id: str) -> List[Dict[str, Any]]:
|
|
out: List[Dict[str, Any]] = []
|
|
for field in schema.fields:
|
|
path = [field.name]
|
|
out.append(
|
|
{
|
|
"producerNodeId": producer_node_id,
|
|
"path": path,
|
|
"type": field.type,
|
|
"label": ".".join(str(p) for p in path),
|
|
"scopeOrigin": "data",
|
|
}
|
|
)
|
|
out.append(
|
|
{
|
|
"producerNodeId": producer_node_id,
|
|
"path": [],
|
|
"type": schema.name,
|
|
"label": "(whole output)",
|
|
"scopeOrigin": "data",
|
|
}
|
|
)
|
|
return out
|
|
|
|
|
|
def _paths_for_schema(schema_name: str, producer_node_id: str) -> List[Dict[str, Any]]:
|
|
if not schema_name or schema_name == "Transit":
|
|
return []
|
|
schema = PORT_TYPE_CATALOG.get(schema_name)
|
|
if not schema:
|
|
return []
|
|
return _paths_for_port_schema(schema, producer_node_id)
|
|
|
|
|
|
def compute_upstream_paths(graph: Dict[str, Any], target_node_id: str) -> List[Dict[str, Any]]:
|
|
"""
|
|
Return flattened first-level paths for every ancestor node's primary output schema.
|
|
"""
|
|
nodes = graph.get("nodes") or []
|
|
connections = graph.get("connections") or []
|
|
node_by_id = {n["id"]: n for n in nodes if n.get("id")}
|
|
if target_node_id not in node_by_id:
|
|
return []
|
|
|
|
conn_map = buildConnectionMap(connections)
|
|
# predecessors: walk backwards along edges (target -> source)
|
|
preds: Dict[str, Set[str]] = {}
|
|
for tgt, pairs in conn_map.items():
|
|
for src, _, _ in pairs:
|
|
preds.setdefault(tgt, set()).add(src)
|
|
|
|
seen: Set[str] = set()
|
|
stack = [target_node_id]
|
|
ancestors: Set[str] = set()
|
|
while stack:
|
|
cur = stack.pop()
|
|
for p in preds.get(cur, ()):
|
|
if p not in seen:
|
|
seen.add(p)
|
|
ancestors.add(p)
|
|
stack.append(p)
|
|
|
|
paths: List[Dict[str, Any]] = []
|
|
for aid in sorted(ancestors):
|
|
anode = node_by_id.get(aid)
|
|
if not anode:
|
|
continue
|
|
nt = anode.get("type", "")
|
|
ndef = _NODE_BY_TYPE.get(nt)
|
|
if not ndef:
|
|
continue
|
|
out0 = (ndef.get("outputPorts") or {}).get(0, {})
|
|
derived = parse_graph_defined_output_schema(anode, out0 if isinstance(out0, dict) else {})
|
|
if derived:
|
|
for entry in _paths_for_port_schema(derived, aid):
|
|
entry["producerLabel"] = (anode.get("title") or "").strip() or aid
|
|
paths.append(entry)
|
|
else:
|
|
raw_schema = out0.get("schema") if isinstance(out0, dict) else None
|
|
schema_name = raw_schema if isinstance(raw_schema, str) and raw_schema else "ActionResult"
|
|
for entry in _paths_for_schema(schema_name, aid):
|
|
entry["producerLabel"] = (anode.get("title") or "").strip() or aid
|
|
paths.append(entry)
|
|
|
|
# Lexical loop hints (flow.loop): any loop node in ancestors adds synthetic paths
|
|
for aid in ancestors:
|
|
anode = node_by_id.get(aid) or {}
|
|
if anode.get("type") == "flow.loop":
|
|
paths.extend(
|
|
[
|
|
{
|
|
"producerNodeId": aid,
|
|
"path": ["currentItem"],
|
|
"type": "Any",
|
|
"label": "loop.currentItem",
|
|
"scopeOrigin": "loop",
|
|
},
|
|
{
|
|
"producerNodeId": aid,
|
|
"path": ["currentIndex"],
|
|
"type": "int",
|
|
"label": "loop.currentIndex",
|
|
"scopeOrigin": "loop",
|
|
},
|
|
{
|
|
"producerNodeId": aid,
|
|
"path": ["count"],
|
|
"type": "int",
|
|
"label": "loop.count",
|
|
"scopeOrigin": "loop",
|
|
},
|
|
]
|
|
)
|
|
|
|
return paths
|