Update wbs_diagram_generator.py

wbs_diagram_generator.py

@@ -2,7 +2,7 @@ import graphviz
 import json
 from tempfile import NamedTemporaryFile
 import os
-from graph_generator_utils import add_nodes_and_edges # Keeping this import for consistency, though WBS will use its own recursive logic.
+from graph_generator_utils import add_nodes_and_edges 
 
 def generate_wbs_diagram(json_input: str) -> str:
     """
@@ -12,26 +12,93 @@ def generate_wbs_diagram(json_input: str) -> str:
         json_input (str): A JSON string describing the WBS structure.
                           It must follow the Expected JSON Format Example below.
 
-    Returns:
-        str: The filepath to the generated PNG image file.
-
     Expected JSON Format Example:
     {
-      "project_title": "Software Development Project",
+      "project_title": "AI Model Development Project",
       "phases": [
         {
           "id": "phase_prep",
-          "label": "1. Preparation",
+          "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_vision", "label": "1.1. Identify Vision",
+            {
+              "id": "task_2_1_cost_analysis",
+              "label": "Cost Analysis",
               "subtasks": [
-                {"id": "subtask_1_1_1", "label": "1.1.1. Problem Definition",
+                {
+                  "id": "subtask_2_1_1_benefit_analysis",
+                  "label": "Benefit Analysis",
                   "sub_subtasks": [
-                    {"id": "ss_task_1_1_1_1", "label": "1.1.1.1. Req. Analysis",
+                    {
+                      "id": "ss_task_2_1_1_1_risk_assess",
+                      "label": "AI Risk Assessment",
                       "sub_sub_subtasks": [
-                        {"id": "sss_task_1_1_1_1_1", "label": "1.1.1.1.1. User Stories",
+                        {
+                          "id": "sss_task_2_1_1_1_1_model_selection",
+                          "label": "Model Selection",
                           "final_level_tasks": [
-                            {"id": "ft_1_1_1_1_1_1", "label": "1.1.1.1.1.1. Interview Users"}
+                            {"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"}
                           ]
                         }
                       ]
@@ -44,6 +111,9 @@ def generate_wbs_diagram(json_input: str) -> str:
         }
       ]
     }
+
+    Returns:
+        str: The filepath to the generated PNG image file.
     """
     try:
         if not json_input.strip():