evoagentx/workflow/workflow_graph.py

import json
import inspect
import threading
from enum import Enum
import networkx as nx 
from copy import deepcopy
from networkx import MultiDiGraph
from collections import defaultdict
from pydantic import Field, field_validator, model_validator
from typing import Union, Optional, Tuple, Callable, Dict, List
from functools import wraps

from ..core.logging import logger
from ..core.module import BaseModule
from ..core.base_config import Parameter
from .action_graph import ActionGraph
from ..agents.agent import Agent 
from ..utils.utils import generate_dynamic_class_name, make_parent_folder
from ..prompts.workflow.sew_workflow import SEW_WORKFLOW
from ..prompts.workflow.structure_workflow import STRUCTURE_WORKFLOW
from ..prompts.utils import DEFAULT_SYSTEM_PROMPT
# from ..tools.tool import Toolkit, Tool


class WorkFlowNodeState(str, Enum):
    """
    Enumeration of possible states for workflow nodes.
    
    This enum defines the lifecycle states of a workflow node:
    - PENDING: The node is waiting to be executed
    - RUNNING: The node is currently being executed
    - COMPLETED: The node has been successfully executed
    - FAILED: The node execution has failed
    """
    PENDING="pending"
    RUNNING="running"
    COMPLETED = "completed"
    FAILED = "failed"


class WorkFlowNode(BaseModule):
    """
    Represents a node in a workflow graph.
    
    A workflow node represents a specific task in the workflow with its
    inputs, outputs, and execution metadata. It can have associated agents
    that execute the task and track its execution status.
    
    Attributes:
        name: A unique identifier for the task within a workflow
        description: Detailed description of what the task does
        inputs: List of input parameters required by the task
        outputs: List of output parameters produced by the task
        reason: Optional justification for this task's existence
        agents: Optional list of agents that can execute this task
        action_graph: Optional graph of actions to execute this task
        status: Current execution state of the task
    """

    name: str # A short name of the task. Should be unique in a single workflow
    description: str # A detailed description of the task
    inputs: List[Parameter] # inputs for the task
    outputs: List[Parameter] # outputs of the task
    reason: Optional[str] = None
    agents: Optional[List[Union[str, dict]]] = None
    action_graph: Optional[ActionGraph] = None
    status: Optional[WorkFlowNodeState] = WorkFlowNodeState.PENDING

    @field_validator('agents', mode="before")
    @classmethod
    def check_agent_format(cls, agents: List[Union[str, dict, Agent]]):
        if agents is None:
            return None

        validated_agents = []
        for agent in agents:
            if isinstance(agent, str):
                validated_agents.append(agent)
            elif isinstance(agent, Agent):
                validated_agents.append(agent.get_config())
            elif isinstance(agent, dict):
                assert "name" in agent and "description" in agent, \
                    "must provide the name and description of an agent when specifying an agent with a dict."
                validated_agents.append(agent)
        return validated_agents

    @model_validator(mode="after")
    @classmethod
    def check_action_graph(cls, instance: "WorkFlowNode"):
        """
        Validates that:
        1. All required parameters of execute/async_execute methods are included in inputs
        2. The execute/async_execute methods return dictionaries
        3. All output parameters are present in the returned dictionaries
        """
        if instance.action_graph is None:
            return instance
        
        # Get input parameter names from the node's input parameters
        input_param_names = {param.name for param in instance.inputs if param.required}
        output_param_names = {param.name for param in instance.outputs if param.required}
        
        def check_method_signature(method, method_name):
            """Helper function to check method signature against input parameters"""

            method_source = inspect.getsource(method)
            if "NotImplementedError" in method_source:
                return
                
            # Get method signature
            method_sig = inspect.signature(method)
            
            # Only consider parameters other than self, *args, and **kwargs as required
            required_params = []
            for name, param in method_sig.parameters.items():
                if name != 'self' and param.kind not in (param.VAR_POSITIONAL, param.VAR_KEYWORD):
                    if param.default == param.empty:  # This is a required parameter
                        required_params.append(name)
            
            # Check if all required parameters are in inputs
            missing_inputs = set(required_params) - input_param_names
            if missing_inputs:
                raise ValueError(f"`{method_name}` method requires parameters that are not in `inputs`: {missing_inputs}")
        
        # Check execute method
        check_method_signature(instance.action_graph.execute, "execute")
        # Check async_execute method if it exists
        check_method_signature(instance.action_graph.async_execute, "async_execute")

        # Monkey-patch execute and async_execute to check returns at runtime
        original_execute = instance.action_graph.execute
        original_async_execute = instance.action_graph.async_execute
        
        def check_method_return(method_name, result):
            if not isinstance(result, dict):
                raise TypeError(f"{method_name} must return a dictionary, got {type(result)}")
            
            # Check if all output keys are in the result
            missing_outputs = output_param_names - set(result.keys())
            if missing_outputs:
                raise ValueError(f"{method_name} return value is missing required outputs: {missing_outputs}")
            
            return result
        
        @wraps(original_execute)
        def patched_execute(*args, **kwargs):
            result = original_execute(*args, **kwargs)
            return check_method_return("execute", result)
        
        @wraps(original_async_execute)
        async def patched_async_execute(*args, **kwargs):
            result = await original_async_execute(*args, **kwargs)
            return check_method_return("async_execute", result)
        
        # Replace the methods with our patched versions
        instance.action_graph.execute = patched_execute
        instance.action_graph.async_execute = patched_async_execute
        
        return instance

    def to_dict(self, exclude_none: bool = True, ignore: List[str] = [], **kwargs) -> dict:

        data = super().to_dict(exclude_none=exclude_none, ignore=ignore, **kwargs)
        for agent in data.get("agents", []):
            # for CustomizeAgent 
            if isinstance(agent, dict) and "parse_func" in agent and isinstance(agent["parse_func"], Callable):
                agent["parse_func"] = agent["parse_func"].__name__
        return data

    def get_agents(self) -> List[str]:
        """
        Return the names of all agents associated with this node.
        """
        agent_names = []
        if not self.agents:
            return []
        
        for agent in self.agents:
            if isinstance(agent, str):
                agent_names.append(agent)
            elif isinstance(agent, dict):
                agent_names.append(agent["name"])
            else:
                raise TypeError(f"{type(agent)} is an unknown agent type!")
        return agent_names
    
    def set_agents(self, agents: List[Union[str, dict]]):
        self.agents = agents

    def get_status(self) -> WorkFlowNodeState:
        return self.status
    
    def set_status(self, state: WorkFlowNodeState):
        self.status = state
    
    @property
    def is_complete(self) -> bool:
        return self.status == WorkFlowNodeState.COMPLETED
    
    def get_task_info(self) -> str:

        def format_parameters(params: List[Parameter]) -> str:
            if not params:
                return "None"
            return "\n".join(f"  - {param.name} ({param.type}): {param.description}" for param in params)
        
        desc = (
            f"Name: {self.name}\n"
            f"Description: {self.description}\n"
            f"Inputs:\n{format_parameters(self.inputs)}\n"
            f"Outputs:\n{format_parameters(self.outputs)}\n"
        )
        return desc

    def get_input_names(self, required: bool = False) -> List[str]:
        if required:
            return [param.name for param in self.inputs if param.required]
        else:
            return [param.name for param in self.inputs]
    
    def get_output_names(self, required: bool = False) -> List[str]:
        if required:
            return [param.name for param in self.outputs if param.required]
        else:
            return [param.name for param in self.outputs]

class WorkFlowEdge(BaseModule):
    """
    Represents a directed edge in a workflow graph.
    
    Workflow edges connect tasks (nodes) in the workflow graph, establishing
    execution dependencies and data flow relationships. Each edge has a source
    node, target node, and optional priority to influence execution order.
    
    Attributes:
        source: Name of the source node (where the edge starts)
        target: Name of the target node (where the edge ends)
        priority: Numeric priority value for this edge (higher means higher priority)
    """

    source: str 
    target: str 
    priority: int = 0 

    def __init__(self, edge_tuple: Optional[tuple]=(), **kwargs):
        """
        Initialize a WorkFlowEdge instance with either a tuple or keyword arguments.

        Parameters:
        ----------
            edge_tuple (tuple): a tuple containing the edge attributes in the format: (source, target, priority[optional]). 
                - source (str): the source of the edge. 
                - target (str): the target of the edge. 
                - priority (int, optional): The priority of the edge. Defaults to 0 if not provided.
            
            kwargs (dict): Key-value pairs specifying the edge attributes. These values will override those provided in `args` if both are supplied.

        Notes:
        ----------
            - Attributes provided via `kwargs` take precedence over those from the `args` tuple.
            - If `args` is empty or not provided, only `kwargs` will be used to initialize the instance.
        """
        data = self.init_from_tuple(edge_tuple)
        data.update(kwargs)
        super().__init__(**data)
    
    def init_from_tuple(self, edge_tuple: tuple) -> dict:
        if not edge_tuple:
            return {}
        keys = ["source", "target", "priority"]
        data = {k: v for k, v in zip(keys, edge_tuple)}
        return data
    
    def compare_attrs(self):
        return (self.source, self.target, self.priority)
    
    def __eq__(self, other: "WorkFlowEdge"):
        if not isinstance(other, WorkFlowEdge):
            return NotImplemented
        self_compare_attrs = self.compare_attrs()
        other_compare_attrs = other.compare_attrs()
        return all(self_attr==other_attr for self_attr, other_attr in zip(self_compare_attrs, other_compare_attrs))

    def __hash__(self):
        return hash(self.compare_attrs())


class WorkFlowGraph(BaseModule):
    """
    Represents a complete workflow as a directed graph.
    
    WorkFlowGraph models a workflow as a directed graph where nodes represent tasks
    and edges represent dependencies and data flow between tasks. It provides
    methods for constructing, validating, traversing, and executing workflows.
    
    The graph structure supports advanced features like detecting and handling loops,
    determining execution order, and tracking execution state.
    
    Attributes:
        goal: The high-level objective of this workflow
        nodes: List of WorkFlowNode instances representing tasks
        edges: List of WorkFlowEdge instances representing dependencies
        graph: Internal NetworkX MultiDiGraph or another WorkFlowGraph
    """

    goal: str
    nodes: Optional[List[WorkFlowNode]] = []
    edges: Optional[List[WorkFlowEdge]] = []
    graph: Optional[Union[MultiDiGraph, "WorkFlowGraph"]] = Field(default=None, exclude=True)

    def init_module(self):
        self._lock = threading.Lock()
        if not self.graph:
            self._init_from_nodes_and_edges(self.nodes, self.edges)
        elif isinstance(self.graph, MultiDiGraph):
            self._init_from_multidigraph(self.graph, self.nodes, self.edges)
        elif isinstance(self.graph, WorkFlowGraph):
            self._init_from_workflowgraph(self.graph, self.nodes, self.edges)
        else:
            raise TypeError(f"{type(self.graph)} is an unknown type for graph. Supported types: [MultiDiGraph, WorkFlowGraph]")
        self._validate_workflow_structure()
        self.update_graph()
    
    def update_graph(self):
        # call this function when modifying nodes or edges!
        self._loops = self._find_all_loops()

    def _init_from_nodes_and_edges(self, nodes: List[WorkFlowNode] = [], edges: List[WorkFlowEdge] = []):

        """
        Initialize the WorkFlowGraph from a set of nodes and edges. 
        """
        
        if edges and not nodes:
            raise ValueError("edges cannot be passed without nodes or a graph")
        
        self.nodes = []
        self.edges = []
        self.graph = MultiDiGraph()
        self.add_nodes(*nodes, update_graph=False)
        self.add_edges(*edges, update_graph=False)

    def _init_from_multidigraph(self, graph: MultiDiGraph, nodes: List[WorkFlowNode] = [], edges: List[WorkFlowEdge] = []):

        graph_nodes = [deepcopy(node_attrs["ref"]) for _, node_attrs in graph.nodes(data=True)]
        graph_edges = [deepcopy(edge_attrs["ref"]) for *_, edge_attrs in graph.edges(data=True)]
        graph_nodes = self.merge_nodes(graph_nodes, nodes)
        graph_edges = self.merge_edges(graph_edges, edges)
        self._init_from_nodes_and_edges(nodes=graph_nodes, edges=graph_edges)

    def _init_from_workflowgraph(self, graph: "WorkFlowGraph", nodes: List[WorkFlowNode] = [], edges: List[WorkFlowEdge] = []):

        graph_nodes = deepcopy(graph.nodes)
        graph_edges = deepcopy(graph.edges)
        graph_nodes = self.merge_nodes(graph_nodes, nodes)
        graph_edges = self.merge_edges(graph_edges, edges)
        self._init_from_nodes_and_edges(nodes=graph_nodes, edges=graph_edges)
    
    def _validate_workflow_structure(self):

        isolated_nodes = list(nx.isolates(self.graph))
        if len(self.graph.nodes) > 1 and isolated_nodes:
            logger.warning(f"The workflow contains isolated nodes: {isolated_nodes}")
        
        initial_nodes = self.find_initial_nodes()
        if len(self.graph.nodes) > 1 and not initial_nodes:
            error_message = "There are no initial nodes in the workflow!"
            logger.error(error_message)
            raise ValueError(error_message)

        end_nodes = self.find_end_nodes()
        if len(self.graph.nodes) > 1 and not end_nodes:
            logger.warning("There are no end nodes in the workflow")
    
    def find_initial_nodes(self) -> List[str]:
        initial_nodes = [node for node, in_degree in self.graph.in_degree() if in_degree==0]
        return initial_nodes
    
    def find_end_nodes(self) -> List[str]:
        end_nodes = [node for node, out_degree in self.graph.out_degree() if out_degree==0]
        return end_nodes
    
    def _find_loops(self, start_node: Union[str, WorkFlowNode]) -> Dict[str, list]:

        if isinstance(start_node, str):
            start_node = self.get_node(node_name=start_node)
        start_node_name = start_node.name

        loops = defaultdict(list)
        def dfs(current_node_name: str, path: List[str]):
            if current_node_name in path:
                # a loop exists
                loops[current_node_name].append(path[path.index(current_node_name):])
                return
            path.append(current_node_name)
            children = self.get_node_children(current_node_name)
            if children:
                for child in children:
                    dfs(child, path)
            path.pop()
        
        dfs(start_node_name, [])
        return loops

    def _find_all_loops(self) -> Dict[str, list]:

        initial_nodes = self.find_initial_nodes()
        if not initial_nodes:
            return {} 
        
        def contain_loop(loops: List[List[str]], new_loop: List[str]):
            if not loops:
                return False
            return frozenset(new_loop) in [frozenset(loop) for loop in loops]
        
        # merge loops from different nodes 
        all_loops = defaultdict(list)
        for initial_node in initial_nodes:
            loops_from_init_node = self._find_loops(initial_node)
            for start_node, loops in loops_from_init_node.items():
                for loop in loops:
                    if not contain_loop(all_loops[start_node], loop):
                        # 合并从相同的start_node开始的环
                        all_loops[start_node].append(loop)
        
        if len(all_loops) <= 1:
            return all_loops
        
        # merge same loops with different starts (因为同一个环可能在之前的遍历中有不同的start_node)
        loop_to_start_nodes = defaultdict(dict)
        for start_node, loops in all_loops.items():
            for loop in loops:
                normalized_loop = frozenset(loop)
                loop_to_start_nodes[normalized_loop][start_node] = loop
        
        all_paths: List[List[str]] = [] 
        # 用深度遍历来判断一个环中的开始节点
        for initial_node in initial_nodes:
            all_paths.extend(self.get_all_paths_from_node(initial_node))
        
        def rank_nodes(nodes: List[str]):
            if len(nodes) == 1:
                return nodes[0]
            path_contain_nodes = None
            for path in all_paths:
                if all(node in path for node in nodes):
                    path_contain_nodes = path
                    break
            if path_contain_nodes is None:
                raise ValueError(f"Couldn't find a path that contain nodes: {nodes}")
            node_indices = [path.index(node) for node in nodes]
            return nodes[node_indices.index(min(node_indices))]
        
        all_loops = defaultdict(list)
        for start_node_loop in loop_to_start_nodes.values():
            first_node = rank_nodes(list(start_node_loop.keys()))
            all_loops[first_node].append(start_node_loop[first_node])
        
        return all_loops

    def add_node(self, node: WorkFlowNode, update_graph: bool = True, **kwargs):

        if not isinstance(node, WorkFlowNode):
            raise ValueError(f"{node} is not a valid WorkFlowNode instance!")
        if self.node_exists(node.name):
            raise ValueError(f"Duplicate node names are not allowed! Found duplicate node name: {node.name}")

        self.nodes.append(node)
        self.graph.add_node(node.name, ref=node)
        if update_graph:
            self.update_graph()

    def add_edge(self, edge: WorkFlowEdge, update_graph: bool = True, **kwargs):

        if not isinstance(edge, WorkFlowEdge):
            raise ValueError(f"{edge} is not a valid WorkFlowEdge instance!")
        for attr, node_name in zip(["source", "target"], [edge.source, edge.target]):
            if not self.node_exists(node_name):
                raise ValueError(f"{attr} node {node_name} does not exists!")
        if self.edge_exists(edge):
            raise ValueError(f"Duplicate edges are not allowed! Found duplicate edges: {edge}")
        
        # check the inputs and outputs of the edge
        source_node = self.get_node(edge.source)
        target_node = self.get_node(edge.target)
        source_output_names = set(param.name for param in source_node.outputs)
        target_input_names = set(param.name for param in target_node.inputs)
        if len(source_output_names & target_input_names) == 0:
            logger.warning(f"The edge ({edge.source}, {edge.target}) has no matching inputs and outputs! You may need to check the inputs and outputs of the nodes to ensure that at least one input of the target node is the output of the source node.")
        
        self.edges.append(edge)
        self.graph.add_edge(edge.source, edge.target, ref=edge)
        if update_graph:
            self.update_graph()

    def add_nodes(self, *nodes: WorkFlowNode, update_graph: bool = True, **kwargs):

        nodes: list = list(nodes)
        nodes.extend([kwargs.pop(var) for var in ["node", "nodes"] if var in kwargs])

        for node in nodes:
            if isinstance(node, (tuple, list)):
                for n in node:
                    self.add_node(n, update_graph=update_graph, **kwargs)
            else:
                self.add_node(node, update_graph=update_graph, **kwargs)

    def add_edges(self, *edges: WorkFlowEdge, update_graph: bool = True, **kwargs):

        edges: list = list(edges)
        edges.extend([kwargs.pop(var) for var in ["edge", "edges"] if var in kwargs])

        for edge in edges:
            if isinstance(edge, (tuple, list)):
                for e in edge:
                    self.add_edge(e, update_graph=update_graph, **kwargs)
            else:
                self.add_edge(edge, update_graph=update_graph, **kwargs)

    def node_exists(self, node: Union[str, WorkFlowNode]) -> bool:
        if isinstance(node, str):
            return node in self.graph.nodes
        elif isinstance(node, WorkFlowNode):
            return node.name in self.graph.nodes
        else:
            raise TypeError("node must be a str or WorkFlowNode instance")
    
    def _edge_exists(self, source: str, target: str, **attr_filters) -> bool:

        if not self.graph.has_edge(source, target):
            return False
        if attr_filters:
            for key, value in attr_filters.items():
                if key not in self.graph[source][target] or self.graph[source][target][key] != value:
                    return False
        return True
    
    def edge_exists(self, edge: Union[Tuple[str, str], WorkFlowEdge], **attr_filters) -> bool:

        """
        Check whether an edge exists in the workflow graph. The input `edge` can either be a tuple or a WorkFlowEdge instance.

        1. If a tuple is passed, it should be (source, target). The function will only determin whether there is an edge between the source node and the target node. 
        If attr_filters is passed, they will also be used to match the edge attributes. 
        2. If a WorkFlowEdge is passed, it will use the __eq__ method in WorkFlowEdge to determine 

        Parameters:
        ----------
            edge (Union[Tuple[str, str], WorkFlowEdge]):
                - If a tuple is provided, it should be in the format `(source, target)`. 
                The method will check whether there is an edge between the source and target nodes.
                If `attr_filters` are provided, they will be used to match edge attributes.
                - If a WorkFlowEdge instance is provided, the method will use the `__eq__` method in WorkFlowEdge 
                to determine whether the edge exists.

            attr_filters (dict, optional):
                Additional attributes to filter edges when `edge` is a tuple.

        Returns:
        -------
            bool: True if the edge exists and matches the filters (if provided); False otherwise.
        """
        if isinstance(edge, tuple):
            assert len(edge) == 2, "edge must be a tuple (source, target) or WorkFlowEdge instance"
            source, target = edge 
            return self._edge_exists(source, target, **attr_filters)
        elif isinstance(edge, WorkFlowEdge):
            return edge in self.edges 
        else:
            raise TypeError("edge must be a tuple (source, target) or WorkFlowEdge instance")
    
    def is_loop_start(self, node: Union[str, WorkFlowNode]) -> bool:
        if len(self._loops) == 0:
            return False
        node_name = node if isinstance(node, str) else node.name
        return node_name in self._loops
    
    def is_loop_end(self, node: Union[str, WorkFlowNode]) -> bool:
        if len(self._loops) == 0:
            return False
        loop_end_nodes = set()
        node_name = node if isinstance(node, str) else node.name
        for loops in self._loops.values():
            loop_end_nodes.update([loop[-1] for loop in loops])
        return node_name in loop_end_nodes
    
    def find_loops_with_start_and_end(self, start_node: Union[str, WorkFlowNode], end_node: Union[str, WorkFlowNode]) -> List[List[str]]:
        if len(self._loops) == 0:
            return []
        start_node_name = start_node if isinstance(start_node, str) else start_node.name
        end_node_name = end_node if isinstance(end_node, str) else end_node.name
        if start_node_name not in self._loops:
            return [] 
        target = []
        for loop in self._loops[start_node_name]:
            if loop[-1] == end_node_name:
                target.append(loop)
        return target

    def merge_nodes(self, nodes: List[WorkFlowNode], new_nodes: List[WorkFlowNode]):

        node_names = {node.name for node in nodes}
        for node in new_nodes:
            if node.name in node_names:
                continue
            nodes.append(node)
        return nodes
    
    def merge_edges(self, edges: List[WorkFlowEdge], new_edges: List[WorkFlowEdge]):

        for edge in new_edges:
            if edge in edges:
                continue
            edges.append(edge)
        return edges 

    def list_nodes(self) -> List[str]:
        """
        return the names of all nodes 
        """
        return [node.name for node in self.nodes]

    def get_node(self, node_name: str) -> WorkFlowNode:
        """
        return a WorkFlowNode instance based on its name.
        """
        if not self.node_exists(node=node_name):
            raise KeyError(f"{node_name} is an invalid node name. Currently available node names: {self.list_nodes()}")
        return self.graph.nodes[node_name]["ref"]
    
    def get_node_status(self, node: Union[str, WorkFlowNode]) -> WorkFlowNodeState:
        if isinstance(node, str):
            node = self.get_node(node_name=node)
        return node.get_status()
    
    @property
    def is_complete(self):
        # node_complete_list = [node.is_complete for node in self.nodes]
        leaf_nodes = [self.get_node(name) for name in self.find_end_nodes()]
        node_complete_list = [node.is_complete for node in leaf_nodes]
        if len(node_complete_list) == 0:
            return True
        if all(node_complete_list):
            return True
        return False
    
    def reset_graph(self):
        """
        set the status of all nodes to pending
        """
        for node in self.nodes:
            node.set_status(WorkFlowNodeState.PENDING)

    def set_node_status(self, node: Union[str, WorkFlowNode], new_state: WorkFlowNodeState) -> bool:
        """
        Update the state of a specific node. 

        Args:
            node (Union[str, WorkFlowNode]): The name of a node or the node instance.
            new_state (WorkFlowNodeState): The new state to set.
        
        Returns:
            bool: True if the state was updated successfully, False otherwise.
        """
        flag = False
        try:
            if isinstance(node, str):
                node = self.get_node(node_name=node)
            node.set_status(new_state)
            flag = True
        except Exception as e:
            raise ValueError(f"An error occurs when setting node status: {e}")
        return flag
    
    def pending(self, node: Union[str, WorkFlowNode]) -> bool:
        return self.set_node_status(node=node, new_state=WorkFlowNodeState.PENDING)
    
    def running(self, node: Union[str, WorkFlowNode]) -> bool:
        return self.set_node_status(node=node, new_state=WorkFlowNodeState.RUNNING)
    
    def completed(self, node: Union[str, WorkFlowNode]) -> bool:
        return self.set_node_status(node=node, new_state=WorkFlowNodeState.COMPLETED)
    
    def failed(self, node: Union[str, WorkFlowNode]) -> bool:
        return self.set_node_status(node=node, new_state=WorkFlowNodeState.FAILED)
    
    def get_node_children(self, node: Union[str, WorkFlowNode]) -> List[str]:
        node_name = node if isinstance(node, str) else node.name
        if not self.node_exists(node=node):
            raise ValueError(f"Node `{node_name}` does not exists!")
        children = list(self.graph.successors(node_name))
        return children
    
    def get_node_predecessors(self, node: Union[str, WorkFlowNode]) -> List[str]:
        node_name = node if isinstance(node, str) else node.name
        if not self.node_exists(node=node):
            raise ValueError(f"Node `{node_name}` does not exists!")
        predecessors = list(self.graph.predecessors(node_name))
        return predecessors

    def get_uncomplete_initial_nodes(self) -> List[str]:
        initial_nodes = self.find_initial_nodes()
        are_initial_nodes_complete = [self.get_node(node_name).is_complete for node_name in initial_nodes]
        uncomplete_initial_nodes = []
        for node_name, is_complete in zip(initial_nodes, are_initial_nodes_complete):
            if not is_complete:
                uncomplete_initial_nodes.append(node_name)
        return uncomplete_initial_nodes 
    
    def get_all_paths_from_node(self, start_node: Union[str, WorkFlowNode]) -> List[List[str]]:

        if isinstance(start_node, str):
            start_node = self.get_node(node_name=start_node)
        start_node_name = start_node.name

        all_paths = []
        visited = set() # handle loop in the graph

        def dfs(current_node_name: str, path: List[str]):
            if current_node_name in visited:
                # 如果一个loop的end node只有指向loop的start node的边，那么添加这条路径
                if path and len(self.get_node_children(path[-1])) == 1:
                    all_paths.append(path.copy())
                return
            path.append(current_node_name)
            visited.add(current_node_name)
            children = self.get_node_children(current_node_name)
            if not children:
                all_paths.append(path.copy())
            else:
                for child in children:
                    dfs(child, path)
            
            path.pop()
            visited.remove(current_node_name)
        
        dfs(start_node_name, [])
        return all_paths

    def find_completed_leaf_nodes(self, start_node: Union[str, WorkFlowNode]) -> List[str]:

        if isinstance(start_node, str):
            start_node = self.get_node(node_name=start_node)
        start_node_name = start_node.name

        paths_starting_from_node = self.get_all_paths_from_node(start_node=start_node_name)
        last_completed_nodes = [] 
        for path in paths_starting_from_node:
            if not path:
                continue
            completed_node = None
            for path_node in path:
                if self.get_node(path_node).is_complete:
                    completed_node = path_node
                else:
                    break
            if completed_node and completed_node not in last_completed_nodes:
                last_completed_nodes.append(completed_node)
        last_completed_nodes = last_completed_nodes[::-1]
        return last_completed_nodes
    
    def find_completed_leaf_nodes_start_from_initial_nodes(self) -> List[str]:

        initial_nodes = self.find_initial_nodes()
        completed_leaf_nodes = []
        for initial_node in initial_nodes:
            for complete_node in self.find_completed_leaf_nodes(start_node=initial_node):
                if complete_node not in completed_leaf_nodes:
                    completed_leaf_nodes.append(complete_node)
        return completed_leaf_nodes
    
    def get_all_children_nodes(self, nodes: List[Union[str, WorkFlowNode]]) -> List[str]:

        node_names = [node if isinstance(node, str) else node.name for node in nodes]
        children_nodes = [] 
        for node_name in node_names:
            for child in self.get_node_children(node_name):
                if child not in children_nodes:
                    children_nodes.append(child)
        return children_nodes
    
    def filter_completed_nodes(self, nodes: List[Union[str, WorkFlowNode]]) -> List[str]:
        """
        remove completed nodes from `nodes`
        """
        node_names = [node if isinstance(node, str) else node.name for node in nodes]
        uncompleted_nodes = [] 
        for node_name in node_names:
            if self.get_node(node_name).is_complete:
                continue
            uncompleted_nodes.append(node_name)
        return uncompleted_nodes
    
    def get_candidate_children_nodes(self, completed_nodes: List[Union[str, WorkFlowNode]]) -> List[str]:
        """
        Return the next set of possible tasks to execute. If there are no loops in the graph, consider only the uncompleted children. 
        If there exists loops, also consider the previous completed tasks.

        Args:
            completed_nodes (List[Union[str, WorkFlowNode]]): A list of completed nodes.
            
        Returns:
            List[str]: List of node names that are candidates for execution.
        """
        node_names = [node if isinstance(node, str) else node.name for node in completed_nodes]
        has_loop = (len(self._loops) > 0)
        if has_loop:
            # if there exists loops, we need to check the completed nodes and their children nodes
            uncompleted_children_nodes = []
            for node_name in node_names:
                children_nodes = self.get_all_children_nodes(nodes=[node_name])
                if self.is_loop_end(node=node_name):
                    current_uncompleted_children_nodes = [] 
                    for child in children_nodes:
                        if self.is_loop_start(node=child):
                            # node_name是一个环的结束的时候，如果它的子节点是环的开始，那么无论它是否completed，都添加到下一步可执行的操作
                            current_uncompleted_children_nodes.append(child)
                        else:
                            # node_name是环的结束，但是子节点不是环的开始时，需要检查child是否已经completed，只添加未完成的任务
                            current_uncompleted_children_nodes.extend(self.filter_completed_nodes(nodes=[child]))
                else:
                    current_uncompleted_children_nodes = self.filter_completed_nodes(nodes=children_nodes)
                for child in current_uncompleted_children_nodes:
                    if child not in uncompleted_children_nodes:
                        uncompleted_children_nodes.append(child)
        else:
            # 不存在环的时候直接得到所有的子节点，并去掉其中已完成的部分
            children_nodes = self.get_all_children_nodes(nodes=node_names)
            uncompleted_children_nodes = self.filter_completed_nodes(nodes=children_nodes)

        return uncompleted_children_nodes
    
    def are_dependencies_complete(self, node_name: str) -> bool:
        """
        Check if all predecessors for a node are complete.

        Args:
            node_name (str): The name of the task/node to check.
        
        Returns:
            bool: True if all predecessors are complete, False otherwise.
        """
        has_loop = (len(self._loops) > 0)
        predecessors = self.get_node_predecessors(node=node_name)
        if has_loop and self.is_loop_start(node=node_name):
            flag = True 
            for pre in predecessors:
                if self.is_loop_end(pre):
                    pass 
                else:
                    flag &= self.get_node(pre).is_complete
        else:
            flag = all(self.get_node(pre).is_complete for pre in predecessors)
        return flag

    def filter_nodes_with_uncompleted_predecessors(self, nodes: List[Union[str, WorkFlowNode]]) -> List[str]:
        node_names = [node if isinstance(node, str) else node.name for node in nodes]
        nodes_with_completed_predecessors = [] 
        for node_name in node_names:
            if self.are_dependencies_complete(node_name=node_name):
                nodes_with_completed_predecessors.append(node_name)
        return nodes_with_completed_predecessors

    def get_next_candidate_nodes(self) -> List[str]:

        uncomplete_initial_nodes = self.get_uncomplete_initial_nodes()
        if len(uncomplete_initial_nodes) > 0:
            return uncomplete_initial_nodes
        
        # find the last completed nodes in all paths starting from initial nodes. 
        completed_leaf_nodes = self.find_completed_leaf_nodes_start_from_initial_nodes()

        # obtain children nodes of last completed nodes which are uncompleted (consider previous completed tasks if there exists loops)
        # children_nodes = self.get_all_children_nodes(nodes=completed_leaf_nodes)
        # uncompleted_children_nodes = self.filter_completed_nodes(nodes=children_nodes)
        candidate_children_nodes = self.get_candidate_children_nodes(completed_nodes=completed_leaf_nodes)

        # check whether all the predecessors are completed
        # children_nodes_with_complete_predecessors = self.filter_nodes_with_uncompleted_predecessors(uncompleted_children_nodes)
        children_nodes_with_complete_predecessors = self.filter_nodes_with_uncompleted_predecessors(candidate_children_nodes)

        return children_nodes_with_complete_predecessors

    def next(self) -> List[WorkFlowNode]:
        if self.is_complete:
            return [] 
        candidate_node_names = self.get_next_candidate_nodes()
        candidate_tasks = [self.get_node(node_name=node_name) for node_name in candidate_node_names]
        return candidate_tasks

    def step(self, source_node: Union[str, WorkFlowNode], target_node: Union[str, WorkFlowNode]):

        if source_node is None:
            self.running(target_node)
            return
        
        source_node_name = source_node if isinstance(source_node, str) else source_node.name
        target_node_name = target_node if isinstance(target_node, str) else target_node.name
        source_node_status = self.get_node_status(source_node_name)
        if source_node_status != WorkFlowNodeState.COMPLETED:
            raise ValueError(f"The state of `source_node` should be WorkFlowNodeState.COMPLETED, but found {source_node_status}")
        # set the state of `target_node` to WorkFlowNodeState.RUNNING
        if self.is_loop_end(source_node_name) and self.is_loop_start(target_node_name):
            loops = self.find_loops_with_start_and_end(
                start_node=target_node_name, end_node=source_node_name
            )
            loop_nodes = set(sum(loops, []))
            for loop_node in loop_nodes:
                self.pending(node=loop_node)
        if not self.edge_exists(edge=(source_node_name, target_node_name)):
            # the execution doesn't follow an edge means re-executing a previous subtask due to errors or incomplete output
            # find a path that contains both source_node and target node and set them as "pending"
            all_paths = self.get_all_paths_from_node(start_node=target_node_name)
            for path in all_paths:
                if source_node_name in path:
                    for node_name in path:
                        self.pending(node=node_name)
        self.running(node=target_node_name)
    
    def get_node_description(self, node: Union[str, WorkFlowNode]) -> str:

        if isinstance(node, str):
            node = self.get_node(node_name=node)
        
        def format_parameters(params: List[Parameter]) -> str:
            if not params:
                return "  - None"
            # return "\n".join(f"  - {param.name} ({param.type}): {param.description}" for param in params)
            return "\n".join(f"  - {param.name} ({param.type})" for param in params)
        
        def format_agents(agent_names: List[str]) -> str:
            if not agent_names:
                return "None"
            return "\n".join(f"  - {name}" for name in agent_names)
        
        def format_action_graph(action_graph: ActionGraph) -> str:
            if action_graph is None:
                return "  - None"
            return type(action_graph).__name__
        
        desc = (
            f"Name: {node.name}\n"
            # f"Description: {node.description}\n"
            f"Inputs:\n{format_parameters(node.inputs)}\n"
            f"Outputs:\n{format_parameters(node.outputs)}\n"
            f"Agents:\n{format_agents(node.get_agents())}\n"
            f"Action Graph:\n{format_action_graph(node.action_graph)}"
        )
        return desc

    def display(self):
        """
        Display the workflow graph with node and edge attributes.
        Nodes are colored based on their status.
        """
        import matplotlib.pyplot as plt

        # Define colors for node statuses
        status_colors = {
            WorkFlowNodeState.PENDING: 'lightgray',
            WorkFlowNodeState.RUNNING: 'orange',
            WorkFlowNodeState.COMPLETED: 'green',
            WorkFlowNodeState.FAILED: 'red'
        }

        if not self.graph.nodes:
            print("Graph is empty. No nodes to display.")
            return

        # Get node colors based on their statuses
        node_colors = [status_colors.get(self.get_node_status(node), 'lightgray') for node in self.graph.nodes]

        # Prepare node labels with additional information
        node_labels = {node: self.get_node_description(data["ref"]) for node, data in self.graph.nodes(data=True)}
        
        # Draw the graph
        # pos = nx.shell_layout(self.graph)
        if len(self.graph.nodes) == 1:
            single_node = list(self.graph.nodes)[0]
            pos = {single_node: (0, 0)}  # Place the single node at the center
        else:
            pos = nx.shell_layout(self.graph)
        
        plt.figure(figsize=(12, 8))
        nx.draw(
            self.graph, pos, with_labels=False, node_color=node_colors, edge_color='black',
            node_size=1500, font_size=8, font_color='black', font_weight='bold'
        )

        if len(self.graph.nodes) == 1:
            for node, (x, y) in pos.items():
                plt.text(x+0.005, y, node_labels[node], ha='left', va='center', fontsize=9, bbox=dict(facecolor='white', alpha=0.7))
        else:
            # Draw node labels next to the nodes (left-aligned)
            # text_offsets = {node: (pos[node][0]-0.2, pos[node][1]-0.22) for node in self.graph.nodes}
            y_positions = [y for _, y in pos.values()]
            y_min, y_max = min(y_positions), max(y_positions)
            lower_third_boundary = y_min + (y_max - y_min) / 3

            # Adjust text offsets based on node position in the graph
            text_offsets = {}
            for node, (x, y) in pos.items():
                if y < lower_third_boundary:  # If in the lower third, display label above the node
                    text_offsets[node] = (x-0.2, y + 0.23)
                else:  # Otherwise, display label below the node
                    text_offsets[node] = (x-0.2, y - 0.23)
            
            for node, (x, y) in text_offsets.items():
                plt.text(x, y, node_labels[node], ha='left', va='center', fontsize=9, bbox=dict(facecolor='white', alpha=0.7))

        # Draw edge labels for priorities
        edge_labels = nx.get_edge_attributes(self.graph, 'priority')
        nx.draw_networkx_edge_labels(self.graph, pos, edge_labels=edge_labels)

        # Add a legend to show node status colors
        legend_elements = [
            plt.Line2D([0], [0], marker='o', color='w', label=status.name, markersize=10, markerfacecolor=color)
            for status, color in status_colors.items()
        ]
        plt.legend(handles=legend_elements, title="Workflow Node Status", loc='upper left', fontsize='medium')

        plt.title("Workflow Graph")
        plt.show()

    def get_workflow_description(self) -> str:

        def format_param_requirement(required: bool):
            return "required" if required else "optional"
        
        def format_parameters(params: List[Parameter]) -> str:
            if not params:
                return "None"
            return "\n".join(
                f"  - {param.name} ({param.type}, {format_param_requirement(param.required)}): "
                f"{param.description}" for param in params
            )
        
        subtask_texts = [] 
        for node in self.nodes:
            text = (
                f"Task Name: {node.name}\n"
                f"Description: {node.description}\n"
                f"Inputs:\n{format_parameters(node.inputs)}\n"
                f"Outputs:\n{format_parameters(node.outputs)}"
            )
            subtask_texts.append(text)
        workflow_desc = "\n\n".join(subtask_texts)
        return workflow_desc
    
    def _infer_edges_from_nodes(self, nodes: List[WorkFlowNode]) -> List[WorkFlowEdge]:

        if not nodes:
            return []
        edges: List[WorkFlowEdge] = []
        for node in nodes:
            for another_node in nodes:
                if node.name == another_node.name:
                    continue
                node_output_params = [param.name for param in node.outputs]
                another_node_input_params = [param.name for param in another_node.inputs]
                if any([param in another_node_input_params for param in node_output_params]):
                    edges.append(WorkFlowEdge(edge_tuple=(node.name, another_node.name)))
        return edges
    
    def get_config(self) -> dict:
        """
        Get a dictionary containing all necessary configuration to recreate this workflow graph.
        
        Returns:
            dict: A configuration dictionary that can be used to initialize a new WorkFlowGraph instance
            with the same properties as this one.
        """
        config = self.to_dict() 
        config.pop("graph", None)
        return config


class SequentialWorkFlowGraph(WorkFlowGraph):

    """
    A linear workflow graph with a single path from start to end.

    Args:
        goal (str): The goal of the workflow.
        tasks (List[dict]): A list of tasks with their descriptions and inputs. Each task should have the following format:
            {
                "name": str,
                "description": str,
                "inputs": [{"name": str, "type": str, "required": bool, "description": str}, ...],
                "outputs": [{"name": str, "type": str, "required": bool, "description": str}, ...],
                "prompt": str, 
                "prompt_template": PromptTemplate, 
                "system_prompt" (optional): str, default is DEFAULT_SYSTEM_PROMPT,
                "output_parser" (optional): Type[ActionOutput],
                "parse_mode" (optional): str, default is "str" 
                "parse_func" (optional): Callable,
                "parse_title" (optional): str ,
                "tool_names" (optional): List[str] 
            }
    """

    def __init__(self, goal: str, tasks: List[dict], **kwargs):
        nodes = self._infer_nodes_from_tasks(tasks=tasks)
        edges = self._infer_edges_from_nodes(nodes=nodes)
        super().__init__(goal=goal, nodes=nodes, edges=edges, **kwargs)
    
    def _infer_nodes_from_tasks(self, tasks: List[dict]) -> List[WorkFlowNode]:
        nodes = [self._infer_node_from_task(task=task) for task in tasks]
        return nodes
    
    def _infer_node_from_task(self, task: dict) -> WorkFlowNode:

        node_name = task.get("name", None)
        if not node_name:
            raise ValueError("The `name` for the following task is required: {}".format(task))
        node_description = task.get("description", None)
        if not node_description:
            raise ValueError("The `description` for the following task is required: {}".format(task))
        agent_prompt = task.get("prompt", None)
        agent_prompt_template = task.get("prompt_template", None)
        if not agent_prompt and not agent_prompt_template:
            raise ValueError("The `prompt` or `prompt_template` for the following task is required: {}".format(task))
        
        inputs = task.get("inputs", [])
        outputs = task.get("outputs", [])
        agent_name = generate_dynamic_class_name(node_name+" Agent")
        agent_description = node_description # .replace("task", "agent")
        agent_system_prompt = task.get("system_prompt", DEFAULT_SYSTEM_PROMPT)
        agent_output_parser = task.get("output_parser", None)
        agent_parse_mode = task.get("parse_mode", "str")
        agent_parse_func = task.get("parse_func", None)
        agent_parse_title = task.get("parse_title", None)
        tool_names = task.get("tool_names", None)
        tools = task.get("tools",None)
        # tools = task.get("tools", [])
        # tool_names = []
        # if tools:
        #     for tool in tools:
        #         if isinstance(tool,Toolkit):
        #             tool_names.append(tool.name)
        #         elif isinstance(tool, Tool):
        #             tool_names.append(tool.name)
        #         else:
        #             tool_names.append(tool)

        node = WorkFlowNode.from_dict(
            {
                "name": node_name,
                "description": node_description,
                "inputs": inputs,
                "outputs": outputs,
                "agents": [
                    {
                        "name": agent_name,
                        "description": agent_description,
                        "prompt": agent_prompt,
                        "prompt_template": agent_prompt_template, 
                        "system_prompt": agent_system_prompt,
                        "inputs": inputs,
                        "outputs": outputs,
                        "output_parser": agent_output_parser,
                        "parse_mode": agent_parse_mode,
                        "parse_func": agent_parse_func,
                        "parse_title": agent_parse_title,
                        "tool_names": tool_names,
                        "tools":tools
                    }
                ],
            }
        )
        return node
    
    def get_graph_info(self, **kwargs) -> dict:
        """
        Get the information of the workflow graph.
        """
        config = {
            "class_name": self.__class__.__name__,
            "goal": self.goal, 
            "tasks": [
                {
                    "name": node.name,
                    "description": node.description,
                    "inputs": [param.to_dict(ignore=["class_name"]) for param in node.inputs],
                    "outputs": [param.to_dict(ignore=["class_name"]) for param in node.outputs],
                    "prompt": node.agents[0].get("prompt", None),
                    "prompt_template": node.agents[0].get("prompt_template", None).to_dict() if node.agents[0].get("prompt_template", None) else None,
                    "system_prompt": node.agents[0].get("system_prompt", None),
                    "parse_mode": node.agents[0].get("parse_mode", "str"), 
                    "parse_func": node.agents[0].get("parse_func", None).__name__ if node.agents[0].get("parse_func", None) else None,
                    "parse_title": node.agents[0].get("parse_title", None),
                    "tool_names": node.agents[0].get("tool_names", None),
                    "tools": node.agents[0].get("tools", None)
                }
                for node in self.nodes
            ]
        }
        return config
    
    def save_module(self, path: str, ignore: List[str] = [], **kwargs):
        """
        Save the workflow graph to a module file.
        """
        logger.info("Saving {} to {}", self.__class__.__name__, path)
        config = self.get_graph_info()
        for ignore_key in ignore:
            config.pop(ignore_key, None)
        make_parent_folder(path)
        with open(path, "w", encoding="utf-8") as f:
            json.dump(config, f, indent=4)
        return path
    
    def get_config(self) -> Dict:
        """
        Get a dictionary containing all necessary configuration to recreate this workflow graph.
        
        Returns:
            dict: A configuration dictionary that can be used to initialize a new SequentialWorkFlowGraph instance
            with the same properties as this one.
        """
        return self.get_graph_info()
    

class SEWWorkFlowGraph(SequentialWorkFlowGraph):

    def __init__(self, **kwargs):
        goal = kwargs.pop("goal", SEW_WORKFLOW["goal"])
        tasks = kwargs.pop("tasks", SEW_WORKFLOW["tasks"])
        super().__init__(goal=goal, tasks=tasks, **kwargs)

class STRUCTUREWorkFlowGraph(SequentialWorkFlowGraph):

    def __init__(self, **kwargs):
        goal = kwargs.pop("goal", STRUCTURE_WORKFLOW["goal"])
        tasks = kwargs.pop("tasks", STRUCTURE_WORKFLOW["tasks"])
        super().__init__(goal=goal, tasks=tasks, **kwargs)
        
class QASTRUCTUREWorkFlowGraph(SequentialWorkFlowGraph):

    def __init__(self, **kwargs):
        goal = kwargs.pop("goal", STRUCTURE_WORKFLOW["goal"])
        tasks = kwargs.pop("tasks", STRUCTURE_WORKFLOW["tasks"])
        super().__init__(goal=goal, tasks=tasks, **kwargs)