griptape-nodes 0.58.1__py3-none-any.whl → 0.59.0__py3-none-any.whl

griptape_nodes/machines/control_flow.py

@@ -8,7 +8,7 @@ from typing import TYPE_CHECKING
 from griptape_nodes.exe_types.core_types import Parameter, ParameterTypeBuiltin
 from griptape_nodes.exe_types.node_types import CONTROL_INPUT_PARAMETER, LOCAL_EXECUTION, BaseNode, NodeResolutionState
 from griptape_nodes.machines.fsm import FSM, State
-from griptape_nodes.machines.parallel_resolution import ParallelResolutionMachine
+from griptape_nodes.machines.parallel_resolution import ExecuteDagState, ParallelResolutionMachine
 from griptape_nodes.machines.sequential_resolution import SequentialResolutionMachine
 from griptape_nodes.retained_mode.events.base_events import ExecutionEvent, ExecutionGriptapeNodeEvent
 from griptape_nodes.retained_mode.events.execution_events import (
@@ -20,6 +20,11 @@ from griptape_nodes.retained_mode.events.execution_events import (
 from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
 from griptape_nodes.retained_mode.managers.settings import WorkflowExecutionMode
 
+if TYPE_CHECKING:
+    from griptape_nodes.exe_types.connections import Connections
+    from griptape_nodes.exe_types.flow import ControlFlow
+    from griptape_nodes.exe_types.node_types import NodeGroup
+
 
 @dataclass
 class NextNodeInfo:
@@ -45,6 +50,8 @@ class ControlFlowContext:
     paused: bool = False
     flow_name: str
     pickle_control_flow_result: bool
+    node_to_proxy_map: dict[BaseNode, BaseNode]
+    end_node: BaseNode | None = None
 
     def __init__(
         self,
@@ -67,6 +74,7 @@ class ControlFlowContext:
             self.resolution_machine = SequentialResolutionMachine()
         self.current_nodes = []
         self.pickle_control_flow_result = pickle_control_flow_result
+        self.node_to_proxy_map = {}
 
     def get_next_nodes(self, output_parameter: Parameter | None = None) -> list[NextNodeInfo]:
         """Get all next nodes from the current nodes.
@@ -171,7 +179,6 @@ class ResolveNodeState(State):
         # Resolve nodes - pass first node for sequential resolution
         current_node = context.current_nodes[0] if context.current_nodes else None
         await context.resolution_machine.resolve_node(current_node)
-
         if context.resolution_machine.is_complete():
             # Get the last resolved node from the DAG and set it as current
             if isinstance(context.resolution_machine, ParallelResolutionMachine):
@@ -179,6 +186,8 @@ class ResolveNodeState(State):
                 if last_resolved_node:
                     context.current_nodes = [last_resolved_node]
                 return CompleteState
+            if context.end_node == current_node:
+                return CompleteState
             return NextNodeState
         return None
 
@@ -261,6 +270,7 @@ class CompleteState(State):
                     )
                 )
             )
+        context.end_node = None
         logger.info("Flow is complete.")
         return None
 
@@ -284,14 +294,40 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
         )
         super().__init__(context)
 
-    async def start_flow(self, start_node: BaseNode, debug_mode: bool = False) -> None:  # noqa: FBT001, FBT002
+    async def start_flow(
+        self, start_node: BaseNode, end_node: BaseNode | None = None, *, debug_mode: bool = False
+    ) -> None:
+        # FIRST: Scan all nodes in the flow and create node groups BEFORE any resolution
+        flow_manager = GriptapeNodes.FlowManager()
+        flow = flow_manager.get_flow_by_name(self._context.flow_name)
+        logger.debug("Scanning flow '%s' for node groups before execution", self._context.flow_name)
+
+        try:
+            node_to_proxy_map = self._identify_and_create_node_group_proxies(flow, flow_manager.get_connections())
+            if node_to_proxy_map:
+                logger.info(
+                    "Created %d proxy nodes for %d grouped nodes in flow '%s'",
+                    len(set(node_to_proxy_map.values())),
+                    len(node_to_proxy_map),
+                    self._context.flow_name,
+                )
+            # Store the mapping in context so it can be used by resolution machines
+            self._context.node_to_proxy_map = node_to_proxy_map
+        except ValueError as e:
+            logger.error("Failed to process node groups: %s", e)
+            raise
+
+        # Determine the actual start node (use proxy if it's part of a group)
+        actual_start_node = node_to_proxy_map.get(start_node, start_node)
+
         # If using DAG resolution, process data_nodes from queue first
         if isinstance(self._context.resolution_machine, ParallelResolutionMachine):
-            current_nodes = await self._process_nodes_for_dag(start_node)
+            current_nodes = await self._process_nodes_for_dag(actual_start_node)
         else:
-            current_nodes = [start_node]
+            current_nodes = [actual_start_node]
             # For control flow/sequential: emit all nodes in flow as involved
         self._context.current_nodes = current_nodes
+        self._context.end_node = end_node
         # Set entry control parameter for initial node (None for workflow start)
         for node in current_nodes:
             node.set_entry_control_parameter(None)
@@ -299,12 +335,14 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
         self._context.paused = debug_mode
         flow_manager = GriptapeNodes.FlowManager()
         flow = flow_manager.get_flow_by_name(self._context.flow_name)
-        involved_nodes = list(flow.nodes.keys())
-        GriptapeNodes.EventManager().put_event(
-            ExecutionGriptapeNodeEvent(
-                wrapped_event=ExecutionEvent(payload=InvolvedNodesEvent(involved_nodes=involved_nodes))
+        if start_node != end_node:
+            # This blocks all nodes in the entire flow from running. If we're just resolving one node, we don't want to block that.
+            involved_nodes = list(flow.nodes.keys())
+            GriptapeNodes.EventManager().put_event(
+                ExecutionGriptapeNodeEvent(
+                    wrapped_event=ExecutionEvent(payload=InvolvedNodesEvent(involved_nodes=involved_nodes))
+                )
             )
-        )
         await self.start(ResolveNodeState)  # Begins the flow
 
     async def update(self) -> None:
@@ -347,7 +385,7 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
         ):
             await self.update()
 
-    async def _process_nodes_for_dag(self, start_node: BaseNode) -> list[BaseNode]:
+    async def _process_nodes_for_dag(self, start_node: BaseNode) -> list[BaseNode]:  # noqa: C901
         """Process data_nodes from the global queue to build unified DAG.
 
         This method identifies data_nodes in the execution queue and processes
@@ -361,7 +399,11 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
         if dag_builder is None:
             msg = "DAG builder is not initialized."
             raise ValueError(msg)
-        # Build with the first node:
+
+        # Use the node-to-proxy map that was created in start_flow
+        node_to_proxy_map = self._context.node_to_proxy_map
+
+        # Build with the first node (it should already be the proxy if it's part of a group)
         dag_builder.add_node_with_dependencies(start_node, start_node.name)
         queue_items = list(flow_manager.global_flow_queue.queue)
         start_nodes = [start_node]
@@ -372,17 +414,165 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
             if item.dag_execution_type in (DagExecutionType.CONTROL_NODE, DagExecutionType.START_NODE):
                 node = item.node
                 node.state = NodeResolutionState.UNRESOLVED
-                dag_builder.add_node_with_dependencies(node, node.name)
+                # Use proxy node if this node is part of a group, otherwise use original node
+                if node in node_to_proxy_map:
+                    node_to_add = node_to_proxy_map[node]
+                else:
+                    node_to_add = node
+                # Only add if not already added (proxy might already be in DAG)
+                if node_to_add.name not in dag_builder.node_to_reference:
+                    dag_builder.add_node_with_dependencies(node_to_add, node_to_add.name)
+                    if node_to_add not in start_nodes:
+                        start_nodes.append(node_to_add)
                 flow_manager.global_flow_queue.queue.remove(item)
-                start_nodes.append(node)
             elif item.dag_execution_type == DagExecutionType.DATA_NODE:
                 node = item.node
                 node.state = NodeResolutionState.UNRESOLVED
-                # Build here.
-                dag_builder.add_node_with_dependencies(node, node.name)
+                # Use proxy node if this node is part of a group, otherwise use original node
+                if node in node_to_proxy_map:
+                    node_to_add = node_to_proxy_map[node]
+                else:
+                    node_to_add = node
+                # Only add if not already added (proxy might already be in DAG)
+                if node_to_add.name not in dag_builder.node_to_reference:
+                    dag_builder.add_node_with_dependencies(node_to_add, node_to_add.name)
                 flow_manager.global_flow_queue.queue.remove(item)
+
         return start_nodes
 
+    def _identify_and_create_node_group_proxies(
+        self, flow: ControlFlow, connections: Connections
+    ) -> dict[BaseNode, BaseNode]:
+        """Scan all nodes in flow, identify groups, and create proxy nodes.
+
+        Returns:
+            Dictionary mapping original nodes to their proxy nodes (only for grouped nodes)
+        """
+        from griptape_nodes.exe_types.node_types import NodeGroup, NodeGroupProxyNode
+
+        # Step 1: Identify groups by scanning all nodes in the flow
+        groups: dict[str, NodeGroup] = {}
+        for node in flow.nodes.values():
+            group_id = node.get_parameter_value("job_group")
+
+            # Skip nodes without group assignment, empty group ID, or locked nodes
+            if not group_id or group_id == "" or node.lock:
+                continue
+
+            # Create group if it doesn't exist
+            if group_id not in groups:
+                groups[group_id] = NodeGroup(group_id=group_id)
+
+            # Add node to group
+            groups[group_id].add_node(node)
+
+        if not groups:
+            return {}
+
+        # Step 2: Analyze connections for each group
+        for group in groups.values():
+            self._analyze_group_connections(group, connections)
+
+        # Step 3: Validate each group
+        for group in groups.values():
+            group.validate_no_intermediate_nodes(connections.connections)
+
+        # Step 4: Create proxy nodes and build mapping
+        node_to_proxy_map: dict[BaseNode, BaseNode] = {}
+        for group_id, group in groups.items():
+            # Create proxy node
+            proxy_name = f"__group_proxy_{group_id}"
+            proxy_node = NodeGroupProxyNode(name=proxy_name, node_group=group)
+
+            # Register the proxy node with ObjectManager so it can be found during parameter updates
+            obj_manager = GriptapeNodes.ObjectManager()
+            obj_manager.add_object_by_name(proxy_name, proxy_node)
+
+            # Map all grouped nodes to this proxy
+            for node in group.nodes.values():
+                node_to_proxy_map[node] = proxy_node
+
+            # Remap connections to point to proxy
+            self._remap_connections_to_proxy_node(group, proxy_node, connections)
+
+            # Now create proxy parameters (after remapping so original references are saved)
+            proxy_node.create_proxy_parameters()
+
+        return node_to_proxy_map
+
+    def _analyze_group_connections(self, group: NodeGroup, connections: Connections) -> None:
+        """Analyze and categorize connections for a node group."""
+        node_names_in_group = group.nodes.keys()
+
+        # Analyze all connections in the flow
+        for conn in connections.connections.values():
+            source_in_group = conn.source_node.name in node_names_in_group
+            target_in_group = conn.target_node.name in node_names_in_group
+
+            if source_in_group and target_in_group:
+                # Both endpoints in group - internal connection
+                group.internal_connections.append(conn)
+            elif source_in_group and not target_in_group:
+                # From group to outside - external outgoing
+                group.external_outgoing_connections.append(conn)
+            elif not source_in_group and target_in_group:
+                # From outside to group - external incoming
+                group.external_incoming_connections.append(conn)
+
+    def _remap_connections_to_proxy_node(
+        self, group: NodeGroup, proxy_node: BaseNode, connections: Connections
+    ) -> None:
+        """Remap external connections from group nodes to the proxy node."""
+        # Remap external incoming connections (from outside -> group becomes outside -> proxy)
+        for conn in group.external_incoming_connections:
+            conn_id = id(conn)
+
+            # Save original target node before remapping (for cleanup later)
+            original_target_node = conn.target_node
+            group.original_incoming_targets[conn_id] = original_target_node
+
+            # Remove old incoming index entry
+            if (
+                conn.target_node.name in connections.incoming_index
+                and conn.target_parameter.name in connections.incoming_index[conn.target_node.name]
+            ):
+                connections.incoming_index[conn.target_node.name][conn.target_parameter.name].remove(conn_id)
+
+            # Update connection target to proxy
+            conn.target_node = proxy_node
+
+            # Create proxy parameter name using original node name
+            sanitized_node_name = original_target_node.name.replace(" ", "_")
+            proxy_param_name = f"{sanitized_node_name}__{conn.target_parameter.name}"
+
+            # Add new incoming index entry with proxy parameter name
+            connections.incoming_index.setdefault(proxy_node.name, {}).setdefault(proxy_param_name, []).append(conn_id)
+
+        # Remap external outgoing connections (group -> outside becomes proxy -> outside)
+        for conn in group.external_outgoing_connections:
+            conn_id = id(conn)
+
+            # Save original source node before remapping (for cleanup later)
+            original_source_node = conn.source_node
+            group.original_outgoing_sources[conn_id] = original_source_node
+
+            # Remove old outgoing index entry
+            if (
+                conn.source_node.name in connections.outgoing_index
+                and conn.source_parameter.name in connections.outgoing_index[conn.source_node.name]
+            ):
+                connections.outgoing_index[conn.source_node.name][conn.source_parameter.name].remove(conn_id)
+
+            # Update connection source to proxy
+            conn.source_node = proxy_node
+
+            # Create proxy parameter name using original node name
+            sanitized_node_name = original_source_node.name.replace(" ", "_")
+            proxy_param_name = f"{sanitized_node_name}__{conn.source_parameter.name}"
+
+            # Add new outgoing index entry with proxy parameter name
+            connections.outgoing_index.setdefault(proxy_node.name, {}).setdefault(proxy_param_name, []).append(conn_id)
+
     async def cancel_flow(self) -> None:
         """Cancel all nodes in the flow by delegating to the resolution machine."""
         await self.resolution_machine.cancel_all_nodes()
@@ -391,6 +581,22 @@ class ControlFlowMachine(FSM[ControlFlowContext]):
         self._context.reset(cancel=cancel)
         self._current_state = None
 
+    def cleanup_proxy_nodes(self) -> None:
+        """Cleanup all proxy nodes and restore original connections."""
+        if not self._context.node_to_proxy_map:
+            # If we're calling cleanup, but it's already been cleaned up, we just want to return.
+            return
+
+        # Get all unique proxy nodes
+        proxy_nodes = set(self._context.node_to_proxy_map.values())
+
+        # Cleanup each proxy node using the existing method
+        for proxy_node in proxy_nodes:
+            ExecuteDagState._cleanup_proxy_node(proxy_node)
+
+        # Clear the proxy mapping
+        self._context.node_to_proxy_map.clear()
+
     @property
     def resolution_machine(self) -> ParallelResolutionMachine | SequentialResolutionMachine:
         return self._context.resolution_machine

griptape_nodes/machines/dag_builder.py

@@ -13,7 +13,7 @@ if TYPE_CHECKING:
     import asyncio
 
     from griptape_nodes.exe_types.connections import Connections
-    from griptape_nodes.exe_types.node_types import BaseNode
+    from griptape_nodes.exe_types.node_types import BaseNode, NodeGroup
 
 logger = logging.getLogger("griptape_nodes")
 
@@ -224,3 +224,214 @@ class DagBuilder:
             for node_name in self.graph_to_nodes[graph_name]:
                 self.node_to_reference.pop(node_name, None)
             self.graph_to_nodes.pop(graph_name, None)
+
+    def identify_and_process_node_groups(self) -> dict[str, BaseNode]:
+        """Identify node groups, validate them, and replace with proxy nodes.
+
+        Scans all nodes in the DAG for non-empty node_group parameter values,
+        creates NodeGroup instances, validates they have no intermediate ungrouped nodes,
+        and replaces them with NodeGroupProxyNode instances in the DAG.
+
+        Returns:
+            Dictionary mapping group IDs to their proxy nodes
+
+        Raises:
+            ValueError: If validation fails (e.g., ungrouped nodes between grouped nodes)
+        """
+        from griptape_nodes.retained_mode.griptape_nodes import GriptapeNodes
+
+        connections = GriptapeNodes.FlowManager().get_connections()
+
+        # Step 1: Identify groups by scanning all nodes
+        groups = self._identify_node_groups(connections)
+
+        if not groups:
+            return {}
+
+        # Step 2: Validate each group
+        for group in groups.values():
+            group.validate_no_intermediate_nodes(connections.connections)
+
+        # Step 3: Create proxy nodes and replace groups in DAG
+        proxy_nodes = {}
+        for group_id, group in groups.items():
+            proxy_node = self._create_and_install_proxy_node(group_id, group, connections)
+            proxy_nodes[group_id] = proxy_node
+
+        return proxy_nodes
+
+    def _identify_node_groups(self, connections: Connections) -> dict[str, NodeGroup]:
+        """Identify and build NodeGroup instances from nodes in the DAG.
+
+        Args:
+            connections: Connections object for analyzing connection topology
+
+        Returns:
+            Dictionary mapping group IDs to NodeGroup instances
+        """
+        from griptape_nodes.exe_types.node_types import NodeGroup
+
+        groups: dict[str, NodeGroup] = {}
+
+        # Scan all nodes in DAG for group membership
+        for dag_node in self.node_to_reference.values():
+            node = dag_node.node_reference
+            group_id = node.get_parameter_value("job_group")
+
+            # Skip nodes without group assignment or empty group ID
+            if not group_id or group_id == "":
+                continue
+
+            # Create group if it doesn't exist
+            if group_id not in groups:
+                groups[group_id] = NodeGroup(group_id=group_id)
+
+            # Add node to group
+            groups[group_id].add_node(node)
+
+        # Analyze connections for each group
+        for group in groups.values():
+            self._analyze_group_connections(group, connections)
+
+        return groups
+
+    def _analyze_group_connections(self, group: NodeGroup, connections: Connections) -> None:
+        """Analyze and categorize connections for a node group.
+
+        Categorizes all connections involving group nodes as either:
+        - Internal: Both endpoints within the group
+        - External incoming: From outside node to group node
+        - External outgoing: From group node to outside node
+
+        Args:
+            group: NodeGroup to analyze
+            connections: Connections object containing all flow connections
+        """
+        node_names_in_group = set(group.nodes.keys())
+
+        # Analyze all connections in the flow
+        for conn in connections.connections.values():
+            source_in_group = conn.source_node.name in node_names_in_group
+            target_in_group = conn.target_node.name in node_names_in_group
+
+            if source_in_group and target_in_group:
+                # Both endpoints in group - internal connection
+                group.internal_connections.append(conn)
+            elif source_in_group and not target_in_group:
+                # From group to outside - external outgoing
+                group.external_outgoing_connections.append(conn)
+            elif not source_in_group and target_in_group:
+                # From outside to group - external incoming
+                group.external_incoming_connections.append(conn)
+
+    def _create_and_install_proxy_node(self, group_id: str, group: NodeGroup, connections: Connections) -> BaseNode:
+        """Create a proxy node for a group and install it in the DAG.
+
+        Creates a NodeGroupProxyNode, adds it to the DAG, remaps all external
+        connections to point to the proxy, and removes the original grouped
+        nodes from the DAG.
+
+        Args:
+            group_id: Unique identifier for the group
+            group: NodeGroup instance to replace
+            connections: Connections object for remapping connections
+
+        Returns:
+            The created NodeGroupProxyNode
+        """
+        from griptape_nodes.exe_types.node_types import NodeGroupProxyNode
+
+        # Create proxy node with unique name
+        proxy_name = f"__group_proxy_{group_id}"
+        proxy_node = NodeGroupProxyNode(name=proxy_name, node_group=group)
+
+        # Determine which graph to add proxy to (use first grouped node's graph)
+        target_graph_name = None
+        for graph_name, node_set in self.graph_to_nodes.items():
+            if any(node_name in node_set for node_name in group.nodes):
+                target_graph_name = graph_name
+                break
+
+        if target_graph_name is None:
+            target_graph_name = "default"
+
+        # Add proxy node to DAG
+        self.add_node(proxy_node, target_graph_name)
+
+        # Remap external connections to proxy
+        self._remap_connections_to_proxy(group, proxy_node, connections)
+
+        # Remove grouped nodes from DAG
+        self._remove_grouped_nodes_from_dag(group)
+
+        return proxy_node
+
+    def _remap_connections_to_proxy(self, group: NodeGroup, proxy_node: BaseNode, connections: Connections) -> None:
+        """Remap external connections from group nodes to the proxy node.
+
+        Updates the connection indices and Connection objects to redirect
+        external connections through the proxy node instead of the original
+        grouped nodes.
+
+        Args:
+            group: NodeGroup being replaced
+            proxy_node: Proxy node that will handle external connections
+            connections: Connections object to update
+        """
+        # Remap external incoming connections (from outside -> group becomes outside -> proxy)
+        for conn in group.external_incoming_connections:
+            conn_id = id(conn)
+
+            # Remove old incoming index entry
+            if (
+                conn.target_node.name in connections.incoming_index
+                and conn.target_parameter.name in connections.incoming_index[conn.target_node.name]
+            ):
+                connections.incoming_index[conn.target_node.name][conn.target_parameter.name].remove(conn_id)
+
+            # Update connection target to proxy
+            conn.target_node = proxy_node
+
+            # Add new incoming index entry
+            connections.incoming_index.setdefault(proxy_node.name, {}).setdefault(
+                conn.target_parameter.name, []
+            ).append(conn_id)
+
+        # Remap external outgoing connections (group -> outside becomes proxy -> outside)
+        for conn in group.external_outgoing_connections:
+            conn_id = id(conn)
+
+            # Remove old outgoing index entry
+            if (
+                conn.source_node.name in connections.outgoing_index
+                and conn.source_parameter.name in connections.outgoing_index[conn.source_node.name]
+            ):
+                connections.outgoing_index[conn.source_node.name][conn.source_parameter.name].remove(conn_id)
+
+            # Update connection source to proxy
+            conn.source_node = proxy_node
+
+            # Add new outgoing index entry
+            connections.outgoing_index.setdefault(proxy_node.name, {}).setdefault(
+                conn.source_parameter.name, []
+            ).append(conn_id)
+
+    def _remove_grouped_nodes_from_dag(self, group: NodeGroup) -> None:
+        """Remove all nodes in a group from the DAG graphs and references.
+
+        Args:
+            group: NodeGroup whose nodes should be removed from the DAG
+        """
+        for node_name in group.nodes:
+            # Remove from node_to_reference
+            if node_name in self.node_to_reference:
+                del self.node_to_reference[node_name]
+
+            # Remove from all graphs
+            for graph in self.graphs.values():
+                if node_name in graph.nodes():
+                    graph.remove_node(node_name)
+
+            # Remove from graph_to_nodes tracking
+            for node_set in self.graph_to_nodes.values():
+                node_set.discard(node_name)