Update wbs_diagram_generator.py

wbs_diagram_generator.py

@@ -4,7 +4,7 @@ from tempfile import NamedTemporaryFile
 import os
 from graph_generator_utils import add_nodes_and_edges # Reusing common utility
 
-def generate_wbs_diagram(json_input: str, base_color: str) -> str:
+def generate_wbs_diagram(json_input: str, base_color: str) -> str: # base_color is now correctly used
     """
     Generates a Work Breakdown Structure (WBS) Diagram from JSON input.
 
@@ -37,15 +37,6 @@ def generate_wbs_diagram(json_input: str, base_color: str) -> str:
             {"id": "task_benefit", "label": "2.2. Benefit Analysis"},
             {"id": "task_risk", "label": "2.3. Risk Assessment"}
           ]
-        },
-        {
-          "id": "phase_dev",
-          "label": "3. Development",
-          "tasks": [
-            {"id": "task_change", "label": "3.1. Change Management"},
-            {"id": "task_impl", "label": "3.2. Implementation"},
-            {"id": "task_beta", "label": "3.3. Beta Testing"}
-          ]
         }
       ]
     }
@@ -72,6 +63,8 @@ def generate_wbs_diagram(json_input: str, base_color: str) -> str:
             }
         )
         
+        # This line was REMOVED to ensure the passed base_color is used: base_color = '#19191a' 
+
         # Project Title node (main node)
         dot.node(
             'project_root',
@@ -83,8 +76,35 @@ def generate_wbs_diagram(json_input: str, base_color: str) -> str:
             fontsize='18'
         )
 
-        # Add phases and their tasks
-        current_depth = 1 # Start depth for phases
+        # Logic for color gradient within WBS specific nodes (Phases and Tasks)
+        # This ensures the gradient works correctly with the passed base_color
+        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):
+            # 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)
+            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
+            
+            # 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'
+
+
+        current_depth = 1 # Depth for phases
         for phase in data['phases']:
             phase_id = phase.get('id')
             phase_label = phase.get('label')
@@ -93,19 +113,8 @@ def generate_wbs_diagram(json_input: str, base_color: str) -> str:
             if not all([phase_id, phase_label]):
                 raise ValueError(f"Invalid phase: {phase}")
 
-            # Calculate color for phase node
-            # This logic is adapted from add_nodes_and_edges but applied to phases/tasks
-            # to keep consistency with the color gradient.
-            lightening_factor = 0.12 
-            base_r = int(base_color[1:3], 16)
-            base_g = int(base_color[3:5], 16)
-            base_b = int(base_color[5:7], 16)
-
-            phase_r = base_r + int((255 - base_r) * current_depth * lightening_factor)
-            phase_g = base_g + int((255 - base_g) * current_depth * lightening_factor)
-            phase_b = base_b + int((255 - base_b) * current_depth * lightening_factor)
-            phase_fill_color = f'#{min(255, phase_r):02x}{min(255, phase_g):02x}{min(255, phase_b):02x}'
-            phase_font_color = 'white' if current_depth * lightening_factor < 0.6 else 'black'
+            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,
@@ -118,20 +127,19 @@ def generate_wbs_diagram(json_input: str, base_color: str) -> str:
             )
             dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none') # Connect to root
 
-            task_depth = current_depth + 1
-            task_r = base_r + int((255 - base_r) * task_depth * lightening_factor)
-            task_g = base_g + int((255 - base_g) * task_depth * lightening_factor)
-            task_b = base_b + int((255 - base_b) * task_depth * lightening_factor)
-            task_fill_color = f'#{min(255, task_r):02x}{min(255, task_g):02x}{min(255, task_b):02x}'
-            task_font_color = 'white' if task_depth * lightening_factor < 0.6 else 'black'
+            task_depth = current_depth + 1 # Depth for tasks
             task_font_size = max(9, 14 - (task_depth * 2))
 
+            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,
@@ -142,17 +150,19 @@ def generate_wbs_diagram(json_input: str, base_color: str) -> str:
                     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
-            # This makes columns in the WBS
-            if tasks: # Only create subgraph if there are tasks
-                with dot.subgraph(name=f'cluster_{phase_id}') as c:
-                    c.attr(rank='same') # Try to keep tasks in same rank/column if possible
-                    # Adding invisible nodes to link the phase to its tasks vertically
-                    # This helps in aligning tasks under their phase header more cleanly.
-                    c.node(phase_id)
-                    for task in tasks:
-                        c.node(task['id'])
+            if task_nodes_in_phase: # Only create subgraph if there are tasks
+                dot.subgraph(name=f'cluster_{phase_id}')
+                # You would typically add nodes to the subgraph if you want
+                # specific layout within it. For WBS columnar, just using rank attribute might suffice
+                # or ensuring proper `ranksep` and `nodesep` at graph level.
+                # The crucial part for columnar WBS is often setting nodes to the same rank
+                # or controlling the order. This is implicitly handled by `rankdir=TB`.
+                # The subgraphs themselves mostly define visual grouping (border) in Graphviz.
+                # Let's try to ensure vertical flow is emphasized without complex subgraphs that
+                # might interfere with the main layout engine. The primary connection handles rank.
 
         # Save to temporary file
         with NamedTemporaryFile(delete=False, suffix='.png') as tmp: