Update wbs_diagram_generator.py

wbs_diagram_generator.py

@@ -1,3 +1,246 @@
+# import graphviz
+# import json
+# from tempfile import NamedTemporaryFile
+# import os
+# from graph_generator_utils import add_nodes_and_edges 
+
+# def generate_wbs_diagram(json_input: str, output_format: str) -> str:
+#     """
+#     Generates a Work Breakdown Structure (WBS) Diagram from JSON input.
+
+#     Args:
+#         json_input (str): A JSON string describing the WBS structure.
+#                           It must follow the Expected JSON Format Example below.
+
+#     Expected JSON Format Example:
+#     {
+#       "project_title": "AI Model Development Project",
+#       "phases": [
+#         {
+#           "id": "phase_prep",
+#           "label": "Preparation",
+#           "tasks": [
+#             {
+#               "id": "task_1_1_vision",
+#               "label": "Identify Vision",
+#               "subtasks": [
+#                 {
+#                   "id": "subtask_1_1_1_design_staff",
+#                   "label": "Design & Staffing",
+#                   "sub_subtasks": [
+#                     {
+#                       "id": "ss_task_1_1_1_1_env_setup",
+#                       "label": "Environment Setup",
+#                       "sub_sub_subtasks": [
+#                         {
+#                           "id": "sss_task_1_1_1_1_1_lib_install",
+#                           "label": "Install Libraries",
+#                           "final_level_tasks": [
+#                             {"id": "ft_1_1_1_1_1_1_data_access", "label": "Grant Data Access"}
+#                           ]
+#                         }
+#                       ]
+#                     }
+#                   ]
+#                 }
+#               ]
+#             }
+#           ]
+#         },
+#         {
+#           "id": "phase_plan",
+#           "label": "Planning",
+#           "tasks": [
+#             {
+#               "id": "task_2_1_cost_analysis",
+#               "label": "Cost Analysis",
+#               "subtasks": [
+#                 {
+#                   "id": "subtask_2_1_1_benefit_analysis",
+#                   "label": "Benefit Analysis",
+#                   "sub_subtasks": [
+#                     {
+#                       "id": "ss_task_2_1_1_1_risk_assess",
+#                       "label": "AI Risk Assessment",
+#                       "sub_sub_subtasks": [
+#                         {
+#                           "id": "sss_task_2_1_1_1_1_model_selection",
+#                           "label": "Model Selection",
+#                           "final_level_tasks": [
+#                             {"id": "ft_2_1_1_1_1_1_data_strategy", "label": "Data Strategy"}
+#                           ]
+#                         }
+#                       ]
+#                     }
+#                   ]
+#                 }
+#               ]
+#             }
+#           ]
+#         },
+#         {
+#           "id": "phase_dev",
+#           "label": "Development",
+#           "tasks": [
+#             {
+#               "id": "task_3_1_change_mgmt",
+#               "label": "Data Preprocessing",
+#               "subtasks": [
+#                 {
+#                   "id": "subtask_3_1_1_implementation",
+#                   "label": "Feature Engineering",
+#                   "sub_subtasks": [
+#                     {
+#                       "id": "ss_task_3_1_1_1_beta_testing",
+#                       "label": "Model Training",
+#                       "sub_sub_subtasks": [
+#                         {
+#                           "id": "sss_task_3_1_1_1_1_other_task",
+#                           "label": "Model Evaluation",
+#                           "final_level_tasks": [
+#                             {"id": "ft_3_1_1_1_1_1_hyperparam_tune", "label": "Hyperparameter Tuning"}
+#                           ]
+#                         }
+#                       ]
+#                     }
+#                   ]
+#                 }
+#               ]
+#             }
+#           ]
+#         }
+#       ]
+#     }
+
+#     Returns:
+#         str: The filepath to the generated PNG image file.
+#     """
+#     try:
+#         if not json_input.strip():
+#             return "Error: Empty input"
+            
+#         data = json.loads(json_input)
+        
+#         if 'project_title' not in data or 'phases' not in data:
+#             raise ValueError("Missing required fields: project_title or phases")
+
+#         dot = graphviz.Digraph(
+#             name='WBSDiagram',
+#             format='png',
+#             graph_attr={
+#                 'rankdir': 'TB',        # Top-to-Bottom hierarchy
+#                 'splines': 'ortho',     # Straight lines
+#                 'bgcolor': 'white',     # White background
+#                 'pad': '0.5',           # Padding
+#                 'ranksep': '0.6',       # Adjust vertical separation between ranks
+#                 'nodesep': '0.5'        # Adjust horizontal separation between nodes
+#             }
+#         )
+        
+#         base_color = '#19191a'
+
+#         # Main node
+#         dot.node(
+#             'project_root',
+#             data['project_title'],
+#             shape='box',
+#             style='filled,rounded',
+#             fillcolor=base_color,
+#             fontcolor='white',
+#             fontsize='18'
+#         )
+
+#         # Helper for color and font based on depth for WBS
+#         def get_gradient_color(depth, base_hex_color, lightening_factor=0.12):
+#             base_r = int(base_hex_color[1:3], 16)
+#             base_g = int(base_hex_color[3:5], 16)
+#             base_b = int(base_hex_color[5:7], 16)
+
+#             current_r = base_r + int((255 - base_r) * depth * lightening_factor)
+#             current_g = base_g + int((255 - base_g) * depth * lightening_factor)
+#             current_b = base_b + int((255 - base_b) * depth * lightening_factor)
+            
+#             return f'#{min(255, current_r):02x}{min(255, current_g):02x}{min(255, current_b):02x}'
+
+#         def get_font_color_for_background(depth, base_hex_color, lightening_factor=0.12):
+#             base_r = int(base_hex_color[1:3], 16)
+#             base_g = int(base_hex_color[3:5], 16)
+#             base_b = int(base_hex_color[5:7], 16)
+#             current_r = base_r + (255 - base_r) * depth * lightening_factor
+#             current_g = base_g + (255 - base_g) * depth * lightening_factor
+#             current_b = base_b + (255 - base_b) * depth * lightening_factor
+            
+#             luminance = (0.2126 * current_r + 0.7152 * current_g + 0.0722 * current_b) / 255
+#             return 'white' if luminance < 0.5 else 'black'
+
+#         def _add_wbs_nodes_recursive(parent_id, current_level_tasks, current_depth):
+#             for task_data in current_level_tasks:
+#                 task_id = task_data.get('id')
+#                 task_label = task_data.get('label')
+                
+#                 if not all([task_id, task_label]):
+#                     raise ValueError(f"Invalid task data at depth {current_depth}: {task_data}")
+
+#                 node_fill_color = get_gradient_color(current_depth, base_color)
+#                 node_font_color = get_font_color_for_background(current_depth, base_color)
+#                 font_size = max(9, 14 - (current_depth * 2))
+
+#                 dot.node(
+#                     task_id,
+#                     task_label,
+#                     shape='box',
+#                     style='filled,rounded',
+#                     fillcolor=node_fill_color,
+#                     fontcolor=node_font_color,
+#                     fontsize=str(font_size)
+#                 )
+#                 dot.edge(parent_id, task_id, color='#4a4a4a', arrowhead='none')
+
+#                 # Recursively call for next level of tasks (subtasks, sub_subtasks, etc.)
+#                 # This handles arbitrary nested keys like 'subtasks', 'sub_subtasks', 'final_level_tasks'
+#                 next_level_keys = ['tasks', 'subtasks', 'sub_subtasks', 'sub_sub_subtasks', 'final_level_tasks']
+#                 for key_idx, key in enumerate(next_level_keys):
+#                     if key in task_data and isinstance(task_data[key], list):
+#                         _add_wbs_nodes_recursive(task_id, task_data[key], current_depth + 1)
+#                         break # Only process the first found sub-level key
+
+#         # Process phases (level 1 from project_root)
+#         phase_depth = 1
+#         for phase in data['phases']:
+#             phase_id = phase.get('id')
+#             phase_label = phase.get('label')
+            
+#             if not all([phase_id, phase_label]):
+#                 raise ValueError(f"Invalid phase data: {phase}")
+
+#             phase_fill_color = get_gradient_color(phase_depth, base_color)
+#             phase_font_color = get_font_color_for_background(phase_depth, base_color)
+#             font_size_phase = max(9, 14 - (phase_depth * 2))
+
+#             dot.node(
+#                 phase_id,
+#                 phase_label,
+#                 shape='box',
+#                 style='filled,rounded',
+#                 fillcolor=phase_fill_color,
+#                 fontcolor=phase_font_color,
+#                 fontsize=str(font_size_phase)
+#             )
+#             dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none')
+
+#             # Start recursion for tasks under this phase
+#             if 'tasks' in phase and isinstance(phase['tasks'], list):
+#                 _add_wbs_nodes_recursive(phase_id, phase['tasks'], phase_depth + 1)
+
+#         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp:
+#             dot.render(tmp.name, format=output_format, cleanup=True)
+#             return f"{tmp.name}.{output_format}"
+
+#     except json.JSONDecodeError:
+#         return "Error: Invalid JSON format"
+#     except Exception as e:
+#         return f"Error: {str(e)}"
+
 import graphviz
 import json
 from tempfile import NamedTemporaryFile
@@ -128,30 +371,31 @@ def generate_wbs_diagram(json_input: str, output_format: str) -> str:
             name='WBSDiagram',
             format='png',
             graph_attr={
-                'rankdir': 'TB',        # Top-to-Bottom hierarchy
-                'splines': 'ortho',     # Straight lines
-                'bgcolor': 'white',     # White background
-                'pad': '0.5',           # Padding
-                'ranksep': '0.6',       # Adjust vertical separation between ranks
-                'nodesep': '0.5'        # Adjust horizontal separation between nodes
+                'rankdir': 'TB',
+                'splines': 'ortho',
+                'bgcolor': 'white',
+                'pad': '0.5',
+                'ranksep': '0.6',
+                'nodesep': '0.5'
             }
         )
         
-        base_color = '#19191a'
+        base_color = '#BEBEBE'
 
-        # Main node
         dot.node(
             'project_root',
             data['project_title'],
             shape='box',
             style='filled,rounded',
             fillcolor=base_color,
-            fontcolor='white',
+            fontcolor='black',
             fontsize='18'
         )
 
-        # Helper for color and font based on depth for WBS
-        def get_gradient_color(depth, base_hex_color, lightening_factor=0.12):
+        def get_gradient_color(depth, base_hex_color, lightening_factor=0.06):
+            if not isinstance(base_hex_color, str) or not base_hex_color.startswith('#') or len(base_hex_color) != 7:
+                base_hex_color = '#BEBEBE'
+                
             base_r = int(base_hex_color[1:3], 16)
             base_g = int(base_hex_color[3:5], 16)
             base_b = int(base_hex_color[5:7], 16)
@@ -162,16 +406,8 @@ def generate_wbs_diagram(json_input: str, output_format: str) -> str:
             
             return f'#{min(255, current_r):02x}{min(255, current_g):02x}{min(255, current_b):02x}'
 
-        def get_font_color_for_background(depth, base_hex_color, lightening_factor=0.12):
-            base_r = int(base_hex_color[1:3], 16)
-            base_g = int(base_hex_color[3:5], 16)
-            base_b = int(base_hex_color[5:7], 16)
-            current_r = base_r + (255 - base_r) * depth * lightening_factor
-            current_g = base_g + (255 - base_g) * depth * lightening_factor
-            current_b = base_b + (255 - base_b) * depth * lightening_factor
-            
-            luminance = (0.2126 * current_r + 0.7152 * current_g + 0.0722 * current_b) / 255
-            return 'white' if luminance < 0.5 else 'black'
+        def get_font_color_for_background(depth, base_hex_color, lightening_factor=0.06):
+            return 'black'
 
         def _add_wbs_nodes_recursive(parent_id, current_level_tasks, current_depth):
             for task_data in current_level_tasks:
@@ -196,15 +432,12 @@ def generate_wbs_diagram(json_input: str, output_format: str) -> str:
                 )
                 dot.edge(parent_id, task_id, color='#4a4a4a', arrowhead='none')
 
-                # Recursively call for next level of tasks (subtasks, sub_subtasks, etc.)
-                # This handles arbitrary nested keys like 'subtasks', 'sub_subtasks', 'final_level_tasks'
                 next_level_keys = ['tasks', 'subtasks', 'sub_subtasks', 'sub_sub_subtasks', 'final_level_tasks']
                 for key_idx, key in enumerate(next_level_keys):
                     if key in task_data and isinstance(task_data[key], list):
                         _add_wbs_nodes_recursive(task_id, task_data[key], current_depth + 1)
-                        break # Only process the first found sub-level key
+                        break
 
-        # Process phases (level 1 from project_root)
         phase_depth = 1
         for phase in data['phases']:
             phase_id = phase.get('id')
@@ -228,7 +461,6 @@ def generate_wbs_diagram(json_input: str, output_format: str) -> str:
             )
             dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none')
 
-            # Start recursion for tasks under this phase
             if 'tasks' in phase and isinstance(phase['tasks'], list):
                 _add_wbs_nodes_recursive(phase_id, phase['tasks'], phase_depth + 1)
 
@@ -239,5 +471,4 @@ def generate_wbs_diagram(json_input: str, output_format: str) -> str:
     except json.JSONDecodeError:
         return "Error: Invalid JSON format"
     except Exception as e:
-        return f"Error: {str(e)}"
-
+        return f"Error: {str(e)}"