diff --git a/src/crewai/flow/config.py b/src/crewai/flow/config.py new file mode 100644 index 000000000..ddaddc7a8 --- /dev/null +++ b/src/crewai/flow/config.py @@ -0,0 +1,46 @@ +DARK_GRAY = "#333333" +CREWAI_ORANGE = "#FF5A50" +GRAY = "#666666" +WHITE = "#FFFFFF" + +COLORS = { + "bg": WHITE, + "start": CREWAI_ORANGE, + "method": DARK_GRAY, + "router": DARK_GRAY, + "router_border": CREWAI_ORANGE, + "edge": GRAY, + "router_edge": CREWAI_ORANGE, + "text": WHITE, +} + +NODE_STYLES = { + "start": { + "color": COLORS["start"], + "shape": "box", + "font": {"color": COLORS["text"]}, + "margin": {"top": 10, "bottom": 8, "left": 10, "right": 10}, + }, + "method": { + "color": COLORS["method"], + "shape": "box", + "font": {"color": COLORS["text"]}, + "margin": {"top": 10, "bottom": 8, "left": 10, "right": 10}, + }, + "router": { + "color": { + "background": COLORS["router"], + "border": COLORS["router_border"], + "highlight": { + "border": COLORS["router_border"], + "background": COLORS["router"], + }, + }, + "shape": "box", + "font": {"color": COLORS["text"]}, + "borderWidth": 3, + "borderWidthSelected": 4, + "shapeProperties": {"borderDashes": [5, 5]}, + "margin": {"top": 10, "bottom": 8, "left": 10, "right": 10}, + }, +} diff --git a/src/crewai/flow/flow_visualizer.py b/src/crewai/flow/flow_visualizer.py index 8b00d8822..468b390d1 100644 --- a/src/crewai/flow/flow_visualizer.py +++ b/src/crewai/flow/flow_visualizer.py @@ -1,62 +1,26 @@ # flow_visualizer.py -import base64 import os -import re from pyvis.network import Network -DARK_GRAY = "#333333" -CREWAI_ORANGE = "#FF5A50" -GRAY = "#666666" -WHITE = "#FFFFFF" +from crewai.flow.config import COLORS, NODE_STYLES +from crewai.flow.html_template_handler import HTMLTemplateHandler +from crewai.flow.legend_generator import generate_legend_items_html, get_legend_items +from crewai.flow.utils import calculate_node_levels +from crewai.flow.visualization_utils import ( + add_edges, + add_nodes_to_network, + compute_positions, +) class FlowVisualizer: def __init__(self, flow): self.flow = flow - self.colors = { - "bg": WHITE, - "start": CREWAI_ORANGE, - "method": DARK_GRAY, - "router": DARK_GRAY, - "router_border": CREWAI_ORANGE, - "edge": GRAY, - "router_edge": CREWAI_ORANGE, - "text": WHITE, - } - self.node_styles = { - "start": { - "color": self.colors["start"], - "shape": "box", - "font": {"color": self.colors["text"]}, - "margin": {"top": 10, "bottom": 8, "left": 10, "right": 10}, - }, - "method": { - "color": self.colors["method"], - "shape": "box", - "font": {"color": self.colors["text"]}, - "margin": {"top": 10, "bottom": 8, "left": 10, "right": 10}, - }, - "router": { - "color": { - "background": self.colors["router"], - "border": self.colors["router_border"], - "highlight": { - "border": self.colors["router_border"], - "background": self.colors["router"], - }, - }, - "shape": "box", - "font": {"color": self.colors["text"]}, - "borderWidth": 3, - "borderWidthSelected": 4, - "shapeProperties": {"borderDashes": [5, 5]}, - "margin": {"top": 10, "bottom": 8, "left": 10, "right": 10}, - }, - } + self.colors = COLORS + self.node_styles = NODE_STYLES - # TODO: DROP LIB FOLDER POST GENERATION def visualize(self, filename): net = Network( directed=True, @@ -67,172 +31,16 @@ class FlowVisualizer: ) # Calculate levels for nodes - node_levels = self._calculate_node_levels() - - # Assign positions to nodes based on levels - y_spacing = 150 - x_spacing = 150 - level_nodes = {} - - # Store node positions for edge calculations - node_positions = {} - - for method_name, level in node_levels.items(): - level_nodes.setdefault(level, []).append(method_name) + node_levels = calculate_node_levels(self.flow) # Compute positions - 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) + node_positions = compute_positions(self.flow, node_levels) - method = self.flow._methods.get(method_name) - if hasattr(method, "__is_start_method__"): - node_style = self.node_styles["start"] - elif hasattr(method, "__is_router__"): - node_style = self.node_styles["router"] - else: - node_style = self.node_styles["method"] + # Add nodes to the network + add_nodes_to_network(net, self.flow, node_positions, self.node_styles) - net.add_node( - method_name, - label=method_name, - x=x, - y=y, - fixed=True, - physics=False, - **node_style, - ) - - ancestors = self._build_ancestor_dict() - parent_children = self._build_parent_children_dict() - - # Add edges - for method_name in self.flow._listeners: - condition_type, trigger_methods = self.flow._listeners[method_name] - is_and_condition = condition_type == "AND" - - for trigger in trigger_methods: - if ( - trigger in self.flow._methods - or trigger in self.flow._routers.values() - ): - is_router_edge = any( - trigger in paths for paths in self.flow._router_paths.values() - ) - edge_color = ( - self.colors["router_edge"] - if is_router_edge - else self.colors["edge"] - ) - - # Determine if this edge forms a cycle - is_cycle_edge = self._is_ancestor(trigger, method_name, ancestors) - - # Determine if parent has multiple children - parent_has_multiple_children = ( - len(parent_children.get(trigger, [])) > 1 - ) - - # Edge curvature logic - needs_curvature = is_cycle_edge or parent_has_multiple_children - - if needs_curvature: - # Get node positions - 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] - - if dx <= 0: - # Child is left or directly below - smooth_type = "curvedCCW" # Curve left and down - else: - # Child is to the right - smooth_type = "curvedCW" # Curve right and down - - index = self._get_child_index( - trigger, method_name, parent_children - ) - edge_smooth = { - "type": smooth_type, - "roundness": 0.2 + (0.1 * index), - } - else: - # Fallback curvature - edge_smooth = {"type": "cubicBezier"} - else: - edge_smooth = False # Draw straight line - - edge_style = { - "color": edge_color, - "width": 2, - "arrows": "to", - "dashes": True if is_router_edge or is_and_condition else False, - "smooth": edge_smooth, - } - - net.add_edge(trigger, method_name, **edge_style) - - # Add edges from router methods to their possible paths - for router_method_name, paths in self.flow._router_paths.items(): - for path in paths: - for listener_name, ( - condition_type, - trigger_methods, - ) in self.flow._listeners.items(): - if path in trigger_methods: - is_cycle_edge = self._is_ancestor( - trigger, method_name, ancestors - ) - - # Determine if parent has multiple children - parent_has_multiple_children = ( - len(parent_children.get(router_method_name, [])) > 1 - ) - - # Edge curvature logic - needs_curvature = is_cycle_edge or parent_has_multiple_children - - if needs_curvature: - # Get node positions - 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] - - if dx <= 0: - # Child is left or directly below - smooth_type = "curvedCCW" # Curve left and down - else: - # Child is to the right - smooth_type = "curvedCW" # Curve right and down - - index = self._get_child_index( - router_method_name, listener_name, parent_children - ) - edge_smooth = { - "type": smooth_type, - "roundness": 0.2 + (0.1 * index), - } - else: - # Fallback curvature - edge_smooth = {"type": "cubicBezier"} - else: - edge_smooth = False # Straight line - - edge_style = { - "color": self.colors["router_edge"], - "width": 2, - "arrows": "to", - "dashes": True, - "smooth": edge_smooth, - } - net.add_edge(router_method_name, listener_name, **edge_style) + # Add edges to the network + add_edges(net, self.flow, node_positions, self.colors) # Set options to disable physics net.set_options( @@ -246,227 +54,31 @@ class FlowVisualizer: ) network_html = net.generate_html() - - # Extract just the body content from the generated HTML - match = re.search("