Update app.py

app.py

@@ -25,21 +25,22 @@ class Neo4jGraphVisualizer:
             dict: A dictionary containing nodes and relationships
         """
         with self.driver.session() as session:
-            # Fetch all nodes
+            # Fetch all nodes with elementId instead of deprecated ID()
             nodes_result = session.run("""
                 MATCH (n) 
-                RETURN ID(n) as id, 
+                RETURN elementId(n) as id, 
                        labels(n) as labels, 
                        properties(n) as properties
             """)
             
-            # Fetch all relationships
+            # Fetch all relationships using elementId
             relationships_result = session.run("""
                 MATCH (a)-[r]->(b) 
                 RETURN 
-                    ID(a) as source_id, 
-                    ID(b) as target_id, 
-                    type(r) as relationship_type
+                    elementId(a) as source_id, 
+                    elementId(b) as target_id, 
+                    type(r) as relationship_type,
+                    properties(r) as relationship_properties
             """)
             
             # Process nodes
@@ -57,7 +58,8 @@ class Neo4jGraphVisualizer:
                 {
                     'source': record['source_id'], 
                     'target': record['target_id'], 
-                    'type': record['relationship_type']
+                    'type': record['relationship_type'],
+                    'properties': dict(record.get('relationship_properties', {}))
                 } 
                 for record in relationships_result
             ]
@@ -71,54 +73,69 @@ class Neo4jGraphVisualizer:
         Returns:
             str: Base64 encoded image of the graph
         """
-        # Fetch graph data
-        graph_data = self.fetch_graph_data()
-        
-        # Create NetworkX graph
-        G = nx.DiGraph()
-        
-        # Add nodes
-        for node in graph_data['nodes']:
-            G.add_node(node['id'], 
-                       label=node['label'], 
-                       properties=node['properties'])
-        
-        # Add edges
-        for rel in graph_data['relationships']:
-            G.add_edge(rel['source'], rel['target'], type=rel['type'])
-        
-        # Visualization
-        plt.figure(figsize=(12, 8))
-        pos = nx.spring_layout(G, k=0.5)  # positions for all nodes
-        
-        # Draw nodes
-        nx.draw_networkx_nodes(G, pos, 
-                                node_color='lightblue', 
-                                node_size=300, 
-                                alpha=0.8)
-        
-        # Draw edges
-        nx.draw_networkx_edges(G, pos, 
-                                edge_color='gray', 
-                                arrows=True)
-        
-        # Draw labels
-        nx.draw_networkx_labels(G, pos, 
-                                 labels={node: f"{node}" for node in G.nodes()})
-        
-        plt.title("Neo4j Graph Visualization")
-        plt.axis('off')
-        
-        # Save to buffer
-        buffer = io.BytesIO()
-        plt.savefig(buffer, format='png')
-        buffer.seek(0)
-        image_png = buffer.getvalue()
-        buffer.close()
+        try:
+            # Fetch graph data
+            graph_data = self.fetch_graph_data()
+            
+            # Create NetworkX graph
+            G = nx.DiGraph()
+            
+            # Add nodes
+            for node in graph_data['nodes']:
+                # Use node's label or properties for display
+                label = node.get('properties', {}).get('name', str(node['id']))
+                G.add_node(node['id'], 
+                           label=label, 
+                           properties=node['properties'])
+            
+            # Add edges
+            for rel in graph_data['relationships']:
+                G.add_edge(rel['source'], rel['target'], 
+                           type=rel['type'], 
+                           properties=rel['properties'])
+            
+            # Visualization
+            plt.figure(figsize=(16, 12))
+            pos = nx.spring_layout(G, k=0.9, iterations=50)  # Improved layout
+            
+            # Draw nodes with color and size based on properties
+            node_sizes = [300 + len(str(G.nodes[node]['properties'])) * 10 for node in G.nodes()]
+            node_colors = ['lightblue' if idx % 2 == 0 else 'lightgreen' for idx in range(len(G.nodes()))]
+            
+            nx.draw_networkx_nodes(G, pos, 
+                                    node_color=node_colors, 
+                                    node_size=node_sizes, 
+                                    alpha=0.8)
+            
+            # Draw edges
+            nx.draw_networkx_edges(G, pos, 
+                                    edge_color='gray', 
+                                    arrows=True,
+                                    width=1.5)
+            
+            # Draw labels
+            nx.draw_networkx_labels(G, pos, 
+                                     labels={node: G.nodes[node]['label'] for node in G.nodes()},
+                                     font_size=8)
+            
+            plt.title("Neo4j Graph Visualization")
+            plt.axis('off')
+            
+            # Save to buffer
+            buffer = io.BytesIO()
+            plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
+            buffer.seek(0)
+            image_png = buffer.getvalue()
+            buffer.close()
+            plt.close()  # Close the plot to free memory
+            
+            # Encode
+            graphic = base64.b64encode(image_png).decode('utf-8')
+            return f"data:image/png;base64,{graphic}"
         
-        # Encode
-        graphic = base64.b64encode(image_png).decode('utf-8')
-        return f"data:image/png;base64,{graphic}"
+        except Exception as e:
+            print(f"Error in graph visualization: {e}")
+            return f"Error visualizing graph: {e}"
 
     def close(self):
         """Close the Neo4j driver connection"""
@@ -170,7 +187,7 @@ def main():
     
     try:
         # Launch the interface
-        interface.launch(share=True)
+        interface.launch(server_name='0.0.0.0', server_port=7860)
     except Exception as e:
         print(f"Error launching interface: {e}")
     finally: