Update network_graph_generator.py

network_graph_generator.py

@@ -1,9 +1,182 @@
+# import graphviz
+# import json
+# from tempfile import NamedTemporaryFile
+# import os
+
+# def generate_network_graph(json_input: str, output_format: str) -> str:
+#     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'
+#             }
+#         )
+        
+#         base_color = '#19191a'
+        
+#         type_colors = {
+#             'server': base_color,
+#             'service': '#4a90e2',
+#             'database': '#a0a0a0',
+#             'user': '#f39c12',
+#             'default': base_color
+#         }
+        
+#         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'])
+            
+#             if node_type == 'database' or node_color == '#a0a0a0':
+#                 font_color = 'black'
+#             elif node_color == '#f39c12':
+#                 font_color = 'black'
+#             else:
+#                 font_color = 'white'
+            
+#             if node_type == 'server':
+#                 shape = 'box'
+#                 style = 'filled,rounded'
+#             elif node_type == 'database':
+#                 shape = 'cylinder'
+#                 style = 'filled'
+#             elif node_type == 'user':
+#                 shape = 'ellipse'
+#                 style = 'filled'
+#             elif node_type == 'service':
+#                 shape = 'hexagon'
+#                 style = 'filled'
+#             else:
+#                 shape = 'circle'
+#                 style = 'filled'
+            
+#             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"
@@ -12,7 +185,7 @@ def generate_network_graph(json_input: str, output_format: str) -> str:
         
         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',
@@ -29,14 +202,12 @@ def generate_network_graph(json_input: str, output_format: str) -> str:
             }
         )
         
-        base_color = '#19191a'
-        
         type_colors = {
-            'server': base_color,
-            'service': '#4a90e2',
-            'database': '#a0a0a0',
-            'user': '#f39c12',
-            'default': base_color
+            'server': '#BEBEBE',
+            'service': '#B8D4F1',
+            'database': '#A8E6CF',
+            'user': '#FFF9C4',
+            'default': '#BEBEBE'
         }
         
         nodes = data.get('nodes', [])
@@ -51,29 +222,23 @@ def generate_network_graph(json_input: str, output_format: str) -> str:
                 raise ValueError(f"Invalid node: {node}")
             
             node_color = type_colors.get(node_type, type_colors['default'])
-            
-            if node_type == 'database' or node_color == '#a0a0a0':
-                font_color = 'black'
-            elif node_color == '#f39c12':
-                font_color = 'black'
-            else:
-                font_color = 'white'
+            font_color = 'black'
             
             if node_type == 'server':
                 shape = 'box'
                 style = 'filled,rounded'
             elif node_type == 'database':
                 shape = 'cylinder'
-                style = 'filled'
+                style = 'filled,rounded'
             elif node_type == 'user':
                 shape = 'ellipse'
-                style = 'filled'
+                style = 'filled,rounded'
             elif node_type == 'service':
                 shape = 'hexagon'
-                style = 'filled'
+                style = 'filled,rounded'
             else:
                 shape = 'circle'
-                style = 'filled'
+                style = 'filled,rounded'
             
             dot.node(
                 node_id,
@@ -105,11 +270,11 @@ def generate_network_graph(json_input: str, output_format: str) -> str:
                 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: