crewAI/src/crewai/flow/visualization_utils.py

import ast
import inspect
import os
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, cast

from pyvis.network import Network

from crewai.flow.flow import Flow

from .core_flow_utils import is_ancestor
from .flow_visual_utils import (
    build_ancestor_dict,
    build_parent_children_dict,
    get_child_index,
)
from .path_utils import safe_path_join, validate_file_path


def method_calls_crew(method: Optional[Callable[..., Any]]) -> bool:
    """Check if the method contains a .crew() call in its implementation.
    
    Analyzes the method's source code using AST to detect if it makes any
    calls to the .crew() method, which indicates crew involvement in the
    flow execution.
    
    Args:
        method: The method to analyze for crew calls, can be None
        
    Returns:
        bool: True if the method contains a .crew() call, False otherwise
        
    Raises:
        TypeError: If input is not None and not a callable method
        ValueError: If method source code cannot be parsed
        RuntimeError: If unexpected error occurs during parsing
    """
    if method is None:
        return False
    if not callable(method):
        raise TypeError("Input must be a callable method")
        
    try:
        source = inspect.getsource(method)
        source = inspect.cleandoc(source)
        tree = ast.parse(source)
    except (TypeError, ValueError, OSError) as e:
        raise ValueError(f"Could not parse method {getattr(method, '__name__', str(method))}: {e}")
    except Exception as e:
        raise RuntimeError(f"Unexpected error parsing method: {e}")

    class CrewCallVisitor(ast.NodeVisitor):
        """AST visitor to detect .crew() method calls in source code.
        
        A specialized AST visitor that analyzes Python source code to precisely
        identify calls to the .crew() method, enabling accurate detection of
        crew involvement in flow methods.
        
        Attributes:
            found (bool): Indicates whether a .crew() call was found
        """
        def __init__(self):
            self.found = False

        def visit_Call(self, node):
            if isinstance(node.func, ast.Attribute):
                if node.func.attr == "crew":
                    self.found = True
            self.generic_visit(node)

    visitor = CrewCallVisitor()
    visitor.visit(tree)
    return visitor.found


def add_nodes_to_network(net: Network, flow: Flow[Any], 
                     node_positions: Dict[str, Tuple[float, float]], 
                     node_styles: Dict[str, dict],
                     output_dir: Optional[str] = None) -> None:
    """Add nodes to the network visualization with precise styling and positioning.
    
    Creates and styles nodes in the visualization network based on their type
    (start, router, crew, or regular method) with fine-grained control over
    appearance and positioning.
    
    Args:
        net: The network visualization object to add nodes to
        flow: Flow object containing method definitions and relationships
        node_positions: Dictionary mapping method names to (x,y) coordinates
        node_styles: Dictionary mapping node types to their visual styles
        output_dir: Optional directory path for saving visualization assets
        
    Returns:
        None
        
    Raises:
        ValueError: If flow object is invalid or required styles are missing
        TypeError: If input arguments have incorrect types
        OSError: If output directory operations fail
        
    Note:
        Node styles are applied with precise control over shape, font, color,
        and positioning to ensure accurate visual representation of the flow.
        If output_dir is provided, it will be validated and created if needed.
    """
    if not hasattr(flow, '_methods'):
        raise ValueError("Invalid flow object: missing '_methods' attribute")
    if not isinstance(node_positions, dict):
        raise TypeError("node_positions must be a dictionary")
    if not isinstance(node_styles, dict):
        raise TypeError("node_styles must be a dictionary")
    
    required_styles = {'start', 'router', 'crew', 'method'}
    missing_styles = required_styles - set(node_styles.keys())
    if missing_styles:
        raise ValueError(f"Missing required node styles: {missing_styles}")
        
    # Validate and create output directory if specified
    if output_dir:
        try:
            output_dir = validate_file_path(output_dir, must_exist=False)
            os.makedirs(output_dir, exist_ok=True)
        except (ValueError, OSError) as e:
            raise OSError(f"Failed to create or validate output directory: {e}")
    def human_friendly_label(method_name: str) -> str: 
        """Convert method name to human-readable format.
        
        Args:
            method_name: Original method name with underscores
            
        Returns:
            str: Formatted method name with spaces and title case
        """
        return method_name.replace("_", " ").title()

    for method_name, (x, y) in node_positions.items():
        method = flow._methods.get(method_name)
        if hasattr(method, "__is_start_method__"):
            node_style = node_styles["start"]
        elif hasattr(method, "__is_router__"):
            node_style = node_styles["router"]
        elif method_calls_crew(method):
            node_style = node_styles["crew"]
        else:
            node_style = node_styles["method"]

        node_style = node_style.copy()
        label = human_friendly_label(method_name)

        # Handle file-based assets if output directory is provided
        if output_dir and node_style.get("image"):
            try:
                image_path = node_style["image"]
                safe_image_path = safe_path_join(output_dir, Path(image_path).name)
                node_style["image"] = str(safe_image_path)
            except (ValueError, OSError) as e:
                raise OSError(f"Failed to process node image path: {e}")

        node_style.update(
            {
                "label": label,
                "shape": "box",
                "font": {
                    "multi": "html",
                    "color": node_style.get("font", {}).get("color", "#FFFFFF"),
                },
            }
        )

        net.add_node(
            method_name,
            x=x,
            y=y,
            fixed=True,
            physics=False,
            **node_style,
        )


def compute_positions(flow: Flow[Any], node_levels: Dict[str, int], 
                   y_spacing: float = 150, x_spacing: float = 150) -> Dict[str, Tuple[float, float]]:
    """Calculate precise x,y coordinates for each node in the flow diagram.
    
    Computes optimal node positions with fine-grained control over spacing
    and alignment, ensuring clear visualization of flow hierarchy and
    relationships.
    
    Args:
        flow: Flow object containing method definitions
        node_levels: Dictionary mapping method names to their hierarchy levels
        y_spacing: Vertical spacing between hierarchy levels (default: 150)
        x_spacing: Horizontal spacing between nodes at same level (default: 150)
        
    Returns:
        dict[str, tuple[float, float]]: Dictionary mapping method names to
            their calculated (x,y) coordinates in the visualization
            
    Note:
        Positions are calculated to maintain clear hierarchical structure while
        ensuring optimal spacing and readability of the flow diagram.
    """
    if not hasattr(flow, '_methods'):
        raise ValueError("Invalid flow object: missing '_methods' attribute")
    if not isinstance(node_levels, dict):
        raise TypeError("node_levels must be a dictionary")
    if not isinstance(y_spacing, (int, float)) or y_spacing <= 0:
        raise ValueError("y_spacing must be a positive number")
    if not isinstance(x_spacing, (int, float)) or x_spacing <= 0:
        raise ValueError("x_spacing must be a positive number")
        
    if not node_levels:
        raise ValueError("node_levels dictionary cannot be empty")
    level_nodes: Dict[int, List[str]] = {}
    node_positions: Dict[str, Tuple[float, float]] = {}

    for method_name, level in node_levels.items():
        level_nodes.setdefault(level, []).append(method_name)

    for level, nodes in level_nodes.items():
        x_offset = -(len(nodes) - 1) * x_spacing / 2  # Center nodes horizontally
        for i, method_name in enumerate(nodes):
            x = x_offset + i * x_spacing
            y = level * y_spacing
            node_positions[method_name] = (x, y)

    return node_positions


def add_edges(net: Network, flow: Flow[Any], 
            node_positions: Dict[str, Tuple[float, float]], 
            colors: Dict[str, str],
            asset_dir: Optional[str] = None) -> None:
    if not hasattr(flow, '_methods'):
        raise ValueError("Invalid flow object: missing '_methods' attribute")
    if not hasattr(flow, '_listeners'):
        raise ValueError("Invalid flow object: missing '_listeners' attribute")
    if not hasattr(flow, '_router_paths'):
        raise ValueError("Invalid flow object: missing '_router_paths' attribute")
        
    if not isinstance(node_positions, dict):
        raise TypeError("node_positions must be a dictionary")
    if not isinstance(colors, dict):
        raise TypeError("colors must be a dictionary")
        
    required_colors = {'edge', 'router_edge'}
    missing_colors = required_colors - set(colors.keys())
    if missing_colors:
        raise ValueError(f"Missing required edge colors: {missing_colors}")
        
    # Validate asset directory if provided
    if asset_dir:
        try:
            asset_dir = validate_file_path(asset_dir, must_exist=False)
            os.makedirs(asset_dir, exist_ok=True)
        except (ValueError, OSError) as e:
            raise OSError(f"Failed to create or validate asset directory: {e}")
    ancestors = build_ancestor_dict(flow)
    parent_children = build_parent_children_dict(flow)

    # Edges for normal listeners
    for method_name in flow._listeners:
        condition_type, trigger_methods = flow._listeners[method_name]
        is_and_condition = condition_type == "AND"

        for trigger in trigger_methods:
            # Check if nodes exist before adding edges
            if trigger in node_positions and method_name in node_positions:
                is_router_edge = any(
                    trigger in paths for paths in flow._router_paths.values()
                )
                edge_color = colors["router_edge"] if is_router_edge else colors["edge"]

                is_cycle_edge = is_ancestor(trigger, method_name, ancestors)
                parent_has_multiple_children = len(parent_children.get(trigger, [])) > 1
                needs_curvature = is_cycle_edge or parent_has_multiple_children

                if needs_curvature:
                    source_pos = node_positions.get(trigger)
                    target_pos = node_positions.get(method_name)

                    if source_pos and target_pos:
                        dx = target_pos[0] - source_pos[0]
                        smooth_type = "curvedCCW" if dx <= 0 else "curvedCW"
                        index = get_child_index(trigger, method_name, parent_children)
                        edge_config = {
                            "type": smooth_type,
                            "roundness": 0.2 + (0.1 * index),
                        }
                    else:
                        edge_config = {"type": "cubicBezier"}
                else:
                    edge_config = {"type": "straight"}

                edge_props: Dict[str, Any] = {
                    "color": edge_color,
                    "width": 2,
                    "arrows": "to",
                    "dashes": True if is_router_edge or is_and_condition else False,
                    "smooth": edge_config,
                }

                net.add_edge(trigger, method_name, **edge_props)
            else:
                # Nodes not found in node_positions. Check if it's a known router outcome and a known method.
                is_router_edge = any(
                    trigger in paths for paths in flow._router_paths.values()
                )
                # Check if method_name is a known method
                method_known = method_name in flow._methods

                # If it's a known router edge and the method is known, don't warn.
                # This means the path is legitimate, just not reflected as nodes here.
                if not (is_router_edge and method_known):
                    print(
                        f"Warning: No node found for '{trigger}' or '{method_name}'. Skipping edge."
                    )

    # Edges for router return paths
    for router_method_name, paths in flow._router_paths.items():
        for path in paths:
            for listener_name, (
                condition_type,
                trigger_methods,
            ) in flow._listeners.items():
                if path in trigger_methods:
                    if (
                        router_method_name in node_positions
                        and listener_name in node_positions
                    ):
                        is_cycle_edge = is_ancestor(
                            router_method_name, listener_name, ancestors
                        )
                        parent_has_multiple_children = (
                            len(parent_children.get(router_method_name, [])) > 1
                        )
                        needs_curvature = is_cycle_edge or parent_has_multiple_children

                        if needs_curvature:
                            source_pos = node_positions.get(router_method_name)
                            target_pos = node_positions.get(listener_name)

                            if source_pos and target_pos:
                                dx = target_pos[0] - source_pos[0]
                                smooth_type = "curvedCCW" if dx <= 0 else "curvedCW"
                                index = get_child_index(
                                    router_method_name, listener_name, parent_children
                                )
                                edge_config = {
                                    "type": smooth_type,
                                    "roundness": 0.2 + (0.1 * index),
                                }
                            else:
                                edge_config = {"type": "cubicBezier"}
                        else:
                            edge_config = {"type": "straight"}

                        router_edge_props: Dict[str, Any] = {
                            "color": colors["router_edge"],
                            "width": 2,
                            "arrows": "to",
                            "dashes": True,
                            "smooth": edge_config,
                        }
                        net.add_edge(router_method_name, listener_name, **router_edge_props)
                    else:
                        # Same check here: known router edge and known method?
                        method_known = listener_name in flow._methods
                        if not method_known:
                            print(
                                f"Warning: No node found for '{router_method_name}' or '{listener_name}'. Skipping edge."
                            )