From 454ab55a261987f139c071a13f7fcaa6dbcfda06 Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Mon, 30 Dec 2024 20:59:06 +0000
Subject: [PATCH] fix: type annotations and import sorting in flow modules

Co-Authored-By: Joe Moura <joao@crewai.com>
---
 src/crewai/flow/flow_visual_utils.py   |  4 ++--
 src/crewai/flow/visualization_utils.py | 14 +++++++-------
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/crewai/flow/flow_visual_utils.py b/src/crewai/flow/flow_visual_utils.py
index 75e360510..604d226b9 100644
--- a/src/crewai/flow/flow_visual_utils.py
+++ b/src/crewai/flow/flow_visual_utils.py
@@ -136,8 +136,8 @@ def build_ancestor_dict(flow: Flow[Any]) -> Dict[str, Set[str]]:
         dict[str, set[str]]: Dictionary mapping each method name to a set
             of its ancestor method names
     """
-    ancestors = {node: set() for node in flow._methods}
-    visited = set()
+    ancestors: Dict[str, Set[str]] = {node: set() for node in flow._methods}
+    visited: Set[str] = set()
     for node in flow._methods:
         if node not in visited:
             dfs_ancestors(node, ancestors, visited, flow)
diff --git a/src/crewai/flow/visualization_utils.py b/src/crewai/flow/visualization_utils.py
index e30cc9c5d..e84a04f90 100644
--- a/src/crewai/flow/visualization_utils.py
+++ b/src/crewai/flow/visualization_utils.py
@@ -2,7 +2,7 @@ import ast
 import inspect
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, Optional, Tuple, cast
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, cast
 
 from pyvis.network import Network
 
@@ -205,8 +205,8 @@ def compute_positions(flow: Flow[Any], node_levels: Dict[str, int],
         Positions are calculated to maintain clear hierarchical structure while
         ensuring optimal spacing and readability of the flow diagram.
     """
-    level_nodes = {}
-    node_positions = {}
+    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)
@@ -306,9 +306,9 @@ def add_edges(net: Network, flow: Flow[Any],
                     else:
                         edge_smooth = {"type": "cubicBezier"}
                 else:
-                    edge_smooth = False
+                    edge_smooth: Dict[str, Any] = {"type": "straight"}
 
-                edge_style = {
+                edge_style: Dict[str, Any] = {
                     "color": edge_color,
                     "width": 2,
                     "arrows": "to",
@@ -369,9 +369,9 @@ def add_edges(net: Network, flow: Flow[Any],
                             else:
                                 edge_smooth = {"type": "cubicBezier"}
                         else:
-                            edge_smooth = False
+                            edge_smooth: Dict[str, Any] = {"type": "straight"}
 
-                        edge_style = {
+                        edge_style: Dict[str, Any] = {
                             "color": colors["router_edge"],
                             "width": 2,
                             "arrows": "to",