Graphify / network_graph_generator.py
ZahirJS's picture
Update network_graph_generator.py
e57a611 verified
import graphviz
import json
from tempfile import NamedTemporaryFile
import os
def generate_network_graph(json_input: str, output_format: str) -> str:
"""
Generates a network graph from JSON input.
Args:
json_input (str): A JSON string describing the network graph structure.
It must follow the Expected JSON Format Example below.
output_format (str): The output format for the generated diagram.
Supported formats: "png" or "svg"
Expected JSON Format Example:
{
"nodes": [
{"id": "customers", "label": "Customers", "type": "user"},
{"id": "sellers", "label": "Sellers", "type": "user"},
{"id": "admin", "label": "Admin", "type": "user"},
{"id": "web_frontend", "label": "Web Frontend", "type": "server"},
{"id": "product_service", "label": "Product Service", "type": "service"},
{"id": "cart_service", "label": "Cart Service", "type": "service"},
{"id": "order_service", "label": "Order Service", "type": "service"},
{"id": "payment_service", "label": "Payment Service", "type": "service"},
{"id": "inventory_tracker", "label": "Inventory Tracker", "type": "service"},
{"id": "review_system", "label": "Review System", "type": "service"},
{"id": "email_sender", "label": "Email Sender", "type": "service"},
{"id": "image_uploader", "label": "Image Uploader", "type": "service"},
{"id": "product_db", "label": "Product Database", "type": "database"},
{"id": "user_db", "label": "User Database", "type": "database"},
{"id": "order_db", "label": "Order Database", "type": "database"},
{"id": "image_storage", "label": "Image Storage", "type": "database"},
{"id": "stripe_api", "label": "Stripe API", "type": "service"}
],
"connections": [
{"from": "customers", "to": "web_frontend", "label": "Browse", "weight": 5},
{"from": "sellers", "to": "web_frontend", "label": "Manage Store", "weight": 3},
{"from": "admin", "to": "web_frontend", "label": "Admin Panel", "weight": 2},
{"from": "web_frontend", "to": "product_service", "label": "Get Products", "weight": 4},
{"from": "web_frontend", "to": "cart_service", "label": "Cart Actions", "weight": 4},
{"from": "web_frontend", "to": "order_service", "label": "Place Order", "weight": 3},
{"from": "product_service", "to": "product_db", "label": "Product Data", "weight": 4},
{"from": "product_service", "to": "review_system", "label": "Product Reviews", "weight": 2},
{"from": "cart_service", "to": "inventory_tracker", "label": "Check Stock", "weight": 3},
{"from": "order_service", "to": "order_db", "label": "Store Order", "weight": 4},
{"from": "order_service", "to": "payment_service", "label": "Process Payment", "weight": 4},
{"from": "order_service", "to": "email_sender", "label": "Order Confirmation", "weight": 2},
{"from": "payment_service", "to": "stripe_api", "label": "Charge Card", "weight": 4},
{"from": "inventory_tracker", "to": "product_db", "label": "Update Stock", "weight": 3},
{"from": "review_system", "to": "user_db", "label": "User Reviews", "weight": 2},
{"from": "sellers", "to": "image_uploader", "label": "Upload Images", "weight": 2},
{"from": "image_uploader", "to": "image_storage", "label": "Store Images", "weight": 3}
]
}
Returns:
str: The filepath to the generated image file.
"""
try:
if not json_input.strip():
return "Error: Empty input"
data = json.loads(json_input)
if 'nodes' not in data or 'connections' not in data:
raise ValueError("Missing required fields: nodes or connections")
dot = graphviz.Graph(
name='NetworkGraph',
format='png',
engine='neato',
graph_attr={
'overlap': 'false',
'splines': 'true',
'bgcolor': 'white',
'pad': '1.0',
'layout': 'neato',
'sep': '+15',
'esep': '+10',
'nodesep': '1.5',
'concentrate': 'false',
'maxiter': '1000'
},
node_attr={
'fixedsize': 'false'
},
edge_attr={
'labeldistance': '3.0',
'labelangle': '15',
'labelfloat': 'true'
}
)
type_colors = {
'server': '#BEBEBE',
'service': '#B8D4F1',
'database': '#A8E6CF',
'user': '#FFF9C4',
'default': '#BEBEBE'
}
nodes = data.get('nodes', [])
connections = data.get('connections', [])
for node in nodes:
node_id = node.get('id')
label = node.get('label')
node_type = node.get('type', 'default')
if not all([node_id, label]):
raise ValueError(f"Invalid node: {node}")
node_color = type_colors.get(node_type, type_colors['default'])
font_color = 'black'
if node_type == 'server':
shape = 'box'
style = 'filled,rounded'
elif node_type == 'database':
shape = 'cylinder'
style = 'filled,rounded'
elif node_type == 'user':
shape = 'ellipse'
style = 'filled,rounded'
elif node_type == 'service':
shape = 'hexagon'
style = 'filled,rounded'
else:
shape = 'circle'
style = 'filled,rounded'
dot.node(
node_id,
label,
shape=shape,
style=style,
fillcolor=node_color,
fontcolor=font_color,
fontsize='12'
)
for connection in connections:
from_node = connection.get('from')
to_node = connection.get('to')
label = connection.get('label', '')
weight = connection.get('weight', 1)
if not all([from_node, to_node]):
raise ValueError(f"Invalid connection: {connection}")
penwidth = str(max(1, min(5, weight)))
dot.edge(
from_node,
to_node,
label=label,
color='#4a4a4a',
fontcolor='#4a4a4a',
fontsize='8',
penwidth=penwidth,
labeldistance='2.5',
labelangle='10',
labelfloat='true'
)
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"
except Exception as e:
return f"Error: {str(e)}"