Update network_graph_generator.py

network_graph_generator.py

@@ -4,6 +4,62 @@
 # 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 +68,7 @@
         
 #         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 +85,12 @@
 #             }
 #         )
         
-#         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 +105,23 @@
 #                 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 +153,11 @@
 #                 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:
@@ -194,11 +242,19 @@ def generate_network_graph(json_input: str, output_format: str) -> str:
                 'overlap': 'false',
                 'splines': 'true',
                 'bgcolor': 'white',
-                'pad': '0.5',
-                'layout': 'neato'
+                'pad': '0.8',
+                'layout': 'neato',
+                'sep': '+10',
+                'esep': '+5',
+                'nodesep': '1.0',
+                'concentrate': 'true'
             },
             node_attr={
                 'fixedsize': 'false'
+            },
+            edge_attr={
+                'labeldistance': '2.0',
+                'labelangle': '0'
             }
         )
         
@@ -267,8 +323,10 @@ def generate_network_graph(json_input: str, output_format: str) -> str:
                 label=label,
                 color='#4a4a4a',
                 fontcolor='#4a4a4a',
-                fontsize='10',
-                penwidth=penwidth
+                fontsize='9',
+                penwidth=penwidth,
+                labeldistance='1.5',
+                labelangle='0'
             )
             
         with NamedTemporaryFile(delete=False, suffix=f'.{output_format}') as tmp: