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 
 
-def generate_wbs_diagram(json_input: str) -> str:
+def generate_wbs_diagram(json_input: str, output_format: str) -> str:
     """
     Generates a Work Breakdown Structure (WBS) Diagram from JSON input.
 
@@ -232,10 +232,9 @@ def generate_wbs_diagram(json_input: str) -> str:
             if 'tasks' in phase and isinstance(phase['tasks'], list):
                 _add_wbs_nodes_recursive(phase_id, phase['tasks'], phase_depth + 1)
 
-        # Save to temporary file
-        with NamedTemporaryFile(delete=False, suffix='.png') as tmp:
-            dot.render(tmp.name, format='png', cleanup=True)
-            return tmp.name + '.png'
+        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"