Graphify / process_flow_generator.py
ZahirJS's picture
Update process_flow_generator.py
6fc15fa verified
raw
history blame
5.31 kB
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
"""
Generates a Process Flow Diagram (Flowchart) from JSON input.
Args:
json_input (str): A JSON string describing the process flow 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:
{
"start_node": "Start 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"
}
],
"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"}
]
}
"""
try:
if not json_input.strip():
return "Error: Empty input"
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")
dot = graphviz.Digraph(
name='ProcessFlowDiagram',
format='png',
graph_attr={
'rankdir': 'TB', # Top-to-Bottom flow is common for flowcharts
'splines': 'ortho', # Straight lines with 90-degree bends
'bgcolor': 'white', # White background
'pad': '0.5', # Padding around the graph
'nodesep': '0.6', # Spacing between nodes on same rank
'ranksep': '0.8' # Spacing between ranks
}
)
# 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
}
# 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"}
for node_data in data.get('nodes', []):
all_nodes_data[node_data['id']] = node_data
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
# 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
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'
)
# Add connections (edges)
for connection in data.get('connections', []):
dot.edge(
connection['from'],
connection['to'],
label=connection.get('label', ''),
color='#4a4a4a', # Dark gray for lines
fontcolor='#4a4a4a',
fontsize='10'
)
# 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)}"