Graphify / network_graph_generator.py
ZahirJS's picture
Update network_graph_generator.py
bcd6fa4 verified
raw
history blame
10 kB
# 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.
# Expected JSON Format Example:
# {
# "nodes": [
# {
# "id": "server1",
# "label": "Web Server",
# "type": "server"
# },
# {
# "id": "db1",
# "label": "Database",
# "type": "database"
# },
# {
# "id": "user1",
# "label": "User",
# "type": "user"
# },
# {
# "id": "service1",
# "label": "API Service",
# "type": "service"
# }
# ],
# "connections": [
# {
# "from": "user1",
# "to": "server1",
# "label": "HTTP Request",
# "weight": 2
# },
# {
# "from": "server1",
# "to": "service1",
# "label": "API Call",
# "weight": 3
# },
# {
# "from": "service1",
# "to": "db1",
# "label": "Query",
# "weight": 1
# }
# ]
# }
# Returns:
# str: The filepath to the generated PNG 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': '0.5',
# 'layout': 'neato'
# },
# node_attr={
# 'fixedsize': 'false'
# }
# )
# 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='10',
# penwidth=penwidth
# )
# 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)}"
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.
Expected JSON Format Example:
{
"nodes": [
{
"id": "server1",
"label": "Web Server",
"type": "server"
},
{
"id": "db1",
"label": "Database",
"type": "database"
},
{
"id": "user1",
"label": "User",
"type": "user"
},
{
"id": "service1",
"label": "API Service",
"type": "service"
}
],
"connections": [
{
"from": "user1",
"to": "server1",
"label": "HTTP Request",
"weight": 2
},
{
"from": "server1",
"to": "service1",
"label": "API Call",
"weight": 3
},
{
"from": "service1",
"to": "db1",
"label": "Query",
"weight": 1
}
]
}
Returns:
str: The filepath to the generated PNG 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': 'curved',
'bgcolor': 'white',
'pad': '1.0',
'layout': 'neato',
'sep': '+20',
'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)}"