Update process_flow_generator.py

process_flow_generator.py

@@ -2,9 +2,12 @@ import graphviz
 import json
 from tempfile import NamedTemporaryFile
 import os
-from graph_generator_utils import add_nodes_and_edges # Reusing common utility
 
-def generate_process_flow_diagram(json_input: str, base_color: str) -> str: # base_color is now correctly used
+# No longer importing add_nodes_and_edges from graph_generator_utils
+# as process flow diagrams have a different structure and node addition logic.
+# However, we'll re-implement the color gradient logic directly for consistency.
+
+def generate_process_flow_diagram(json_input: str, base_color: str) -> str:
     """
     Generates a Process Flow Diagram (Flowchart) from JSON input.
 
@@ -19,45 +22,25 @@ def generate_process_flow_diagram(json_input: str, base_color: str) -> str: # ba
 
     Expected JSON Format Example:
     {
-      "start_node": "Start Process",
+      "start_node": "Start Order Process",
       "nodes": [
-        {
-          "id": "step1",
-          "label": "Gather Requirements",
-          "type": "process",
-          "relationship": "Next"
-        },
-        {
-          "id": "decision1",
-          "label": "Is Data Available?",
-          "type": "decision",
-          "relationship": "Check"
-        },
-        {
-          "id": "path_yes",
-          "label": "Process Data",
-          "type": "process",
-          "relationship": "Yes"
-        },
-        {
-          "id": "path_no",
-          "label": "Collect More Data",
-          "type": "process",
-          "relationship": "No"
-        },
-        {
-          "id": "end_node",
-          "label": "End Process",
-          "type": "end"
-        }
+        {"id": "customer_order", "label": "Customer Places Order", "type": "process"},
+        {"id": "check_stock", "label": "Check Stock Availability", "type": "decision"},
+        {"id": "stock_available", "label": "Stock Available?", "type": "decision"},
+        {"id": "process_payment", "label": "Process Payment", "type": "process"},
+        {"id": "order_confirmed", "label": "Order Confirmed", "type": "process"},
+        {"id": "notify_customer_oos", "label": "Notify Customer (Out of Stock)", "type": "process"},
+        {"id": "end_order_process", "label": "End Order Process", "type": "end"}
       ],
       "connections": [
-        {"from": "start_node", "to": "step1", "label": "Initiate"},
-        {"from": "step1", "to": "decision1", "label": "Continue"},
-        {"from": "decision1", "to": "path_yes", "label": "Yes"},
-        {"from": "decision1", "to": "path_no", "label": "No"},
-        {"from": "path_yes", "to": "end_node", "label": "Done"},
-        {"from": "path_no", "to": "end_node", "label": "Done"}
+        {"from": "start_node", "to": "customer_order", "label": "Initiate"},
+        {"from": "customer_order", "to": "check_stock", "label": "Proceed"},
+        {"from": "check_stock", "to": "stock_available", "label": "Result"},
+        {"from": "stock_available", "to": "process_payment", "label": "Yes"},
+        {"from": "stock_available", "to": "notify_customer_oos", "label": "No"},
+        {"from": "process_payment", "to": "order_confirmed", "label": "Success"},
+        {"from": "order_confirmed", "to": "end_order_process", "label": "Complete"},
+        {"from": "notify_customer_oos", "to": "end_order_process", "label": "Finish"}
       ]
     }
     """
@@ -67,8 +50,20 @@ def generate_process_flow_diagram(json_input: str, base_color: str) -> str: # ba
             
         data = json.loads(json_input)
         
-        if 'central_node' not in data or 'nodes' not in data:
-            raise ValueError("Missing required fields: central_node or nodes")
+        # Validate required top-level keys for a flowchart
+        if 'start_node' not in data or 'nodes' not in data or 'connections' not in data:
+            raise ValueError("Missing required fields: 'start_node', 'nodes', or 'connections'")
+
+        # Define specific node shapes for flowchart types
+        node_shapes = {
+            "process": "box",          # Rectangle for processes
+            "decision": "diamond",     # Diamond for decisions
+            "start": "oval",           # Oval for start
+            "end": "oval",             # Oval for end
+            "io": "parallelogram",     # Input/Output
+            "document": "note",        # Document symbol
+            "default": "box"           # Fallback
+        }
 
         dot = graphviz.Digraph(
             name='ProcessFlowDiagram',
@@ -83,48 +78,66 @@ def generate_process_flow_diagram(json_input: str, base_color: str) -> str: # ba
             }
         )
         
-        # This line was REMOVED to ensure the passed base_color is used: base_color = '#19191a' 
-
-        # Determine specific node shapes for flowchart types
-        node_shapes = {
-            "process": "box",          # Rectangle for processes
-            "decision": "diamond",     # Diamond for decisions
-            "start": "oval",           # Oval for start
-            "end": "oval",             # Oval for end
-            "io": "parallelogram",     # Input/Output
-            "document": "note",        # Document symbol
-            "default": "box"           # Fallback
-        }
+        # Ensure base_color is valid, fallback if not
+        if not isinstance(base_color, str) or not base_color.startswith('#') or len(base_color) != 7:
+            base_color = '#19191a' # Fallback to default dark if invalid
 
-        # Add all nodes based on JSON structure
-        all_nodes_data = {}
-        # Special handling for start_node if defined separately
-        if 'start_node' in data:
-            all_nodes_data[data['start_node']] = {"label": data['start_node'], "type": "start"}
+        # Helper for color and font based on depth (simulated for flowcharts if needed, or consistent color)
+        # For flowcharts, often all process/decision nodes are the same color.
+        # If a gradient is desired, a 'depth' could be manually assigned or derived.
+        # For simplicity and typical flowchart appearance, we'll use a consistent base_color for all nodes.
+        fill_color_for_nodes = base_color
+        font_color_for_nodes = 'white' if base_color == '#19191a' or base_color.lower() in ['#000000', '#19191a'] else 'black'
+        
+        # Store all nodes by ID for easy lookup
+        all_defined_nodes = {node['id']: node for node in data['nodes']}
         
-        for node_data in data.get('nodes', []):
-            all_nodes_data[node_data['id']] = node_data
+        # Add start node explicitly
+        start_node_id = data['start_node']
+        dot.node(
+            start_node_id,
+            start_node_id, # Label is typically the ID itself for start/end
+            shape=node_shapes['start'],
+            style='filled,rounded',
+            fillcolor='#2196F3', # A distinct blue for Start
+            fontcolor='white',
+            fontsize='14'
+        )
+
+        # Add all other nodes (process, decision, etc.)
+        for node_id, node_info in all_defined_nodes.items():
+            if node_id == start_node_id: # Skip if it's the start node, already added
+                continue
 
-        for node_id, node_info in all_nodes_data.items():
             node_type = node_info.get("type", "default")
-            shape = node_shapes.get(node_type, "box") # Default to box if type is unknown
+            shape = node_shapes.get(node_type, "box")
 
-            # Use base_color for all nodes for consistent flowchart look
-            fill_color_for_node = base_color
-            font_color_for_node = 'white' if base_color == '#19191a' else 'black' # Keep text readable
+            node_label = node_info['label']
 
-            dot.node(
-                node_id,
-                node_info['label'],
-                shape=shape,
-                style='filled,rounded', # Keep rounded corners
-                fillcolor=fill_color_for_node, 
-                fontcolor=font_color_for_node,
-                fontsize='14'
-            )
+            # Use distinct color for end node if it exists
+            if node_type == 'end':
+                 dot.node(
+                    node_id,
+                    node_label,
+                    shape=shape,
+                    style='filled,rounded',
+                    fillcolor='#F44336', # A distinct red for End
+                    fontcolor='white',
+                    fontsize='14'
+                )
+            else: # Regular process, decision, etc. nodes use the selected base color
+                dot.node(
+                    node_id,
+                    node_label,
+                    shape=shape,
+                    style='filled,rounded',
+                    fillcolor=fill_color_for_nodes, 
+                    fontcolor=font_color_for_nodes,
+                    fontsize='14'
+                )
 
         # Add connections (edges)
-        for connection in data.get('connections', []):
+        for connection in data['connections']:
             dot.edge(
                 connection['from'],
                 connection['to'],