import graphviz import json 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: # base_color is now correctly used """ 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. base_color (str): The hexadecimal color string (e.g., '#19191a') for the base color of the nodes, from which a gradient will be generated. Returns: str: The filepath to the generated PNG image file. Expected JSON Format Example: { "project_title": "Software Development Project", "phases": [ { "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"} ] } ] } """ 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.8', # Adjust vertical separation between ranks 'nodesep': '0.5' # Adjust horizontal separation between nodes } ) # This line was REMOVED to ensure the passed base_color is used: base_color = '#19191a' # Project Title node (main node) dot.node( 'project_root', data['project_title'], shape='box', style='filled,rounded', fillcolor=base_color, # Use the selected 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 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') tasks = phase.get('tasks', []) if not all([phase_id, phase_label]): raise ValueError(f"Invalid phase: {phase}") 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, shape='box', style='filled,rounded', fillcolor=phase_fill_color, fontcolor=phase_font_color, fontsize='14' ) dot.edge('project_root', phase_id, color='#4a4a4a', arrowhead='none') # Connect to root 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, 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}') # 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: dot.render(tmp.name, format='png', cleanup=True) return tmp.name + '.png' except json.JSONDecodeError: return "Error: Invalid JSON format" except Exception as e: return f"Error: {str(e)}"