Spaces:
Running
Running
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)}" |