Update wbs_diagram_generator.py

wbs_diagram_generator.py

@@ -2,9 +2,9 @@ import graphviz
 import json
 from tempfile import NamedTemporaryFile
 import os
-from graph_generator_utils import add_nodes_and_edges
+from graph_generator_utils import add_nodes_and_edges # Keeping this import for consistency, though WBS will use its own recursive logic.
 
-def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
+def generate_wbs_diagram(json_input: str) -> str:
     """
     Generates a Work Breakdown Structure (WBS) Diagram from JSON input.
 
@@ -23,17 +23,23 @@ def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
           "id": "phase_prep",
           "label": "1. Preparation",
           "tasks": [
-            {"id": "task_vision", "label": "1.1. Identify Vision"},
-            {"id": "task_design", "label": "1.2. Design & Staffing"}
-          ]
-        },
-        {
-          "id": "phase_plan",
-          "label": "2. Planning",
-          "tasks": [
-            {"id": "task_cost", "label": "2.1. Cost Analysis"},
-            {"id": "task_benefit", "label": "2.2. Benefit Analysis"},
-            {"id": "task_risk", "label": "2.3. Risk Assessment"}
+            {"id": "task_vision", "label": "1.1. Identify Vision",
+              "subtasks": [
+                {"id": "subtask_1_1_1", "label": "1.1.1. Problem Definition",
+                  "sub_subtasks": [
+                    {"id": "ss_task_1_1_1_1", "label": "1.1.1.1. Req. Analysis",
+                      "sub_sub_subtasks": [
+                        {"id": "sss_task_1_1_1_1_1", "label": "1.1.1.1.1. User Stories",
+                          "final_level_tasks": [
+                            {"id": "ft_1_1_1_1_1_1", "label": "1.1.1.1.1.1. Interview Users"}
+                          ]
+                        }
+                      ]
+                    }
+                  ]
+                }
+              ]
+            }
           ]
         }
       ]
@@ -69,13 +75,12 @@ def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
             data['project_title'],
             shape='box',
             style='filled,rounded',
-            fillcolor=base_color, # Use the selected base color
+            fillcolor=base_color,
             fontcolor='white',
             fontsize='18'
         )
 
-        # Logic for color gradient within WBS specific nodes (Phases and Tasks)
-        # This ensures the gradient works correctly with the hardcoded base_color
+        # 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)
@@ -88,8 +93,6 @@ def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
             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):
-            # Calculate brightness/lightness of the node color at this depth
-            # and return black/white for text accordingly
             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)
@@ -97,23 +100,53 @@ def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
             current_g = base_g + (255 - base_g) * depth * lightening_factor
             current_b = base_b + (255 - base_b) * depth * lightening_factor
             
-            # Simple luminance check (ITU-R BT.709 coefficients)
             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))
 
-        current_depth = 1 # Depth for phases
+                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')
-            tasks = phase.get('tasks', [])
-
+            
             if not all([phase_id, phase_label]):
-                raise ValueError(f"Invalid phase: {phase}")
+                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))
 
-            phase_fill_color = get_gradient_color(current_depth, base_color)
-            phase_font_color = get_font_color_for_background(current_depth, base_color)
-            
             dot.node(
                 phase_id,
                 phase_label,
@@ -121,38 +154,14 @@ def generate_wbs_diagram(json_input: str) -> str: # Removed base_color parameter
                 style='filled,rounded',
                 fillcolor=phase_fill_color,
                 fontcolor=phase_font_color,
-                fontsize='14'
+                fontsize=str(font_size_phase)
             )
-            dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none') # Connect to root
+            dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none')
 
-            task_depth = current_depth + 1 # Depth for tasks
-            task_font_size = max(9, 14 - (task_depth * 2))
+            # 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)
 
-            task_nodes_in_phase = []
-            for task in tasks:
-                task_id = task.get('id')
-                task_label = task.get('label')
-                if not all([task_id, task_label]):
-                    raise ValueError(f"Invalid task: {task}")
-                
-                task_fill_color = get_gradient_color(task_depth, base_color)
-                task_font_color = get_font_color_for_background(task_depth, base_color)
-
-                dot.node(
-                    task_id,
-                    task_label,
-                    shape='box',
-                    style='filled,rounded',
-                    fillcolor=task_fill_color,
-                    fontcolor=task_font_color,
-                    fontsize=str(task_font_size)
-                )
-                dot.edge(phase_id, task_id, color='#4a4a4a', arrowhead='none') # Connect task to phase
-                task_nodes_in_phase.append(task_id)
-            
-            # Use subgraph to enforce vertical alignment for tasks within a phase
-            if task_nodes_in_phase: # Only create subgraph if there are tasks
-                dot.subgraph(name=f'cluster_{phase_id}')
         # Save to temporary file
         with NamedTemporaryFile(delete=False, suffix='.png') as tmp:
             dot.render(tmp.name, format='png', cleanup=True)