|
|
|
@ -1,6 +1,9 @@
|
|
|
|
|
from swarms.structs.base import BaseStructure
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import logging
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class GraphWorkflow(BaseStructure):
|
|
|
|
|
"""
|
|
|
|
|
Represents a graph-based workflow structure.
|
|
|
|
@ -31,27 +34,96 @@ class GraphWorkflow(BaseStructure):
|
|
|
|
|
self.entry_point = None
|
|
|
|
|
|
|
|
|
|
def add(self, node, node_value):
|
|
|
|
|
"""
|
|
|
|
|
Adds a node to the graph with the specified value.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
node (str): The name of the node.
|
|
|
|
|
node_value (str): The value of the node.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
"""
|
|
|
|
|
self.graph[node] = {"value": node_value, "edges": {}}
|
|
|
|
|
logging.info(f"Added node: {node}")
|
|
|
|
|
|
|
|
|
|
def start(self, node_name):
|
|
|
|
|
"""
|
|
|
|
|
Sets the starting node for the workflow.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
node_name (str): The name of the starting node.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
"""
|
|
|
|
|
self._check_node_exists(node_name)
|
|
|
|
|
|
|
|
|
|
def connect(self, from_node, to_node):
|
|
|
|
|
"""
|
|
|
|
|
Connects two nodes in the graph.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
from_node (str): The name of the source node.
|
|
|
|
|
to_node (str): The name of the target node.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
"""
|
|
|
|
|
self._check_node_exists(from_node, to_node)
|
|
|
|
|
|
|
|
|
|
def set_entry_point(self, node_name):
|
|
|
|
|
"""
|
|
|
|
|
Sets the entry point node for the workflow.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
node_name (str): The name of the entry point node.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
ValueError: If the specified node does not exist in the graph.
|
|
|
|
|
"""
|
|
|
|
|
if node_name is self.graph:
|
|
|
|
|
self.entry_point = node_name
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("Node does not exist in graph")
|
|
|
|
|
|
|
|
|
|
def add_edge(self, from_node, to_node):
|
|
|
|
|
"""
|
|
|
|
|
Adds an edge between two nodes in the graph.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
from_node (str): The name of the source node.
|
|
|
|
|
to_node (str): The name of the target node.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
ValueError: If either the source or target node does not exist in the graph.
|
|
|
|
|
"""
|
|
|
|
|
if from_node in self.graph and to_node in self.graph:
|
|
|
|
|
self.graph[from_node]["edges"][to_node] = "edge"
|
|
|
|
|
else:
|
|
|
|
|
raise ValueError("Node does not exist in graph")
|
|
|
|
|
|
|
|
|
|
def add_conditional_edges(self, from_node, condition, edge_dict):
|
|
|
|
|
"""
|
|
|
|
|
Adds conditional edges from a node to multiple nodes based on a condition.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
from_node (str): The name of the source node.
|
|
|
|
|
condition: The condition for the conditional edges.
|
|
|
|
|
edge_dict (dict): A dictionary mapping condition values to target nodes.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
None
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
ValueError: If the source node or any of the target nodes do not exist in the graph.
|
|
|
|
|
"""
|
|
|
|
|
if from_node in self.graph:
|
|
|
|
|
for condition_value, to_node in edge_dict.items():
|
|
|
|
|
if to_node in self.graph:
|
|
|
|
@ -66,16 +138,43 @@ class GraphWorkflow(BaseStructure):
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
def run(self):
|
|
|
|
|
"""
|
|
|
|
|
Runs the workflow and returns the graph.
|
|
|
|
|
|
|
|
|
|
Returns:
|
|
|
|
|
dict: The graph representing the nodes and edges.
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
ValueError: If the entry point is not set.
|
|
|
|
|
"""
|
|
|
|
|
if self.entry_point is None:
|
|
|
|
|
raise ValueError("Entry point not set")
|
|
|
|
|
return self.graph
|
|
|
|
|
|
|
|
|
|
def _check_node_exists(self, node_name):
|
|
|
|
|
"""Checks if a node exists in the graph.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
node_name (_type_): _description_
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
ValueError: _description_
|
|
|
|
|
"""
|
|
|
|
|
if node_name not in self.graph:
|
|
|
|
|
raise ValueError(
|
|
|
|
|
f"Node {node_name} does not exist in graph"
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
def _check_nodes_exist(self, from_node, to_node):
|
|
|
|
|
"""
|
|
|
|
|
Checks if the given from_node and to_node exist in the graph.
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
|
from_node: The starting node of the edge.
|
|
|
|
|
to_node: The ending node of the edge.
|
|
|
|
|
|
|
|
|
|
Raises:
|
|
|
|
|
NodeNotFoundError: If either from_node or to_node does not exist in the graph.
|
|
|
|
|
"""
|
|
|
|
|
self._check_node_exists(from_node)
|
|
|
|
|
self._check_node_exists(to_node)
|
|
|
|
|