CultriX/Generate-Knowledge-Graphs
3
1import matplotlib2matplotlib.use('Agg') # Use non-interactive backend to avoid GUI issues3import matplotlib.pyplot as plt4import networkx as nx5import numpy as np6from typing import Dict, List, Any, Tuple, Optional7import json8import io9import base6410import tempfile11import os12import plotly.graph_objects as go13import plotly.express as px14from pyvis.network import Network15 16class GraphVisualizer:17 def __init__(self):18 self.color_map = {19 'PERSON': '#FF6B6B',20 'ORGANIZATION': '#4ECDC4', 21 'LOCATION': '#45B7D1',22 'CONCEPT': '#96CEB4',23 'EVENT': '#FFEAA7',24 'OBJECT': '#DDA0DD',25 'UNKNOWN': '#95A5A6'26 }27 28 def visualize_graph(self, 29 graph: nx.DiGraph, 30 layout_type: str = "spring",31 show_labels: bool = True,32 show_edge_labels: bool = False,33 node_size_factor: float = 1.0,34 figsize: Tuple[int, int] = (12, 8)) -> str:35 """Create a matplotlib visualization of the graph and return file path."""36 37 if not graph.nodes():38 return self._create_empty_graph_image()39 40 # Create figure41 plt.figure(figsize=figsize)42 plt.clf()43 44 # Calculate layout45 pos = self._calculate_layout(graph, layout_type)46 47 # Get node properties48 node_colors = [self.color_map.get(graph.nodes[node].get('type', 'UNKNOWN'), '#95A5A6') 49 for node in graph.nodes()]50 node_sizes = [graph.nodes[node].get('size', 20) * node_size_factor * 10 51 for node in graph.nodes()]52 53 # Draw nodes54 nx.draw_networkx_nodes(graph, pos, 55 node_color=node_colors,56 node_size=node_sizes,57 alpha=0.8)58 59 # Draw edges60 nx.draw_networkx_edges(graph, pos,61 edge_color='gray',62 arrows=True,63 arrowsize=20,64 alpha=0.6,65 width=1.5)66 67 # Draw labels68 if show_labels:69 # Create labels with importance scores70 labels = {}71 for node in graph.nodes():72 importance = graph.nodes[node].get('importance', 0.0)73 labels[node] = f"{node}\n({importance:.2f})"74 75 nx.draw_networkx_labels(graph, pos, labels, font_size=8)76 77 # Draw edge labels78 if show_edge_labels:79 edge_labels = {(u, v): data.get('relationship', '') 80 for u, v, data in graph.edges(data=True)}81 nx.draw_networkx_edge_labels(graph, pos, edge_labels, font_size=6)82 83 plt.title("Knowledge Graph", fontsize=16, fontweight='bold')84 plt.axis('off')85 plt.tight_layout()86 87 # Save to temporary file88 temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')89 plt.savefig(temp_file.name, format='png', dpi=150, bbox_inches='tight')90 plt.close()91 92 return temp_file.name93 94 def _calculate_layout(self, graph: nx.DiGraph, layout_type: str) -> Dict[str, Tuple[float, float]]:95 """Calculate node positions using specified layout algorithm."""96 try:97 if layout_type == "spring":98 return nx.spring_layout(graph, k=1, iterations=50)99 elif layout_type == "circular":100 return nx.circular_layout(graph)101 elif layout_type == "shell":102 return nx.shell_layout(graph)103 elif layout_type == "kamada_kawai":104 return nx.kamada_kawai_layout(graph)105 elif layout_type == "random":106 return nx.random_layout(graph)107 else:108 return nx.spring_layout(graph, k=1, iterations=50)109 except:110 # Fallback to simple layout if algorithm fails111 return nx.spring_layout(graph, k=1, iterations=50)112 113 def _create_empty_graph_image(self) -> str:114 """Create an image for empty graph."""115 plt.figure(figsize=(8, 6))116 plt.text(0.5, 0.5, 'No graph data to display', 117 horizontalalignment='center', verticalalignment='center',118 fontsize=16, transform=plt.gca().transAxes)119 plt.axis('off')120 121 # Save to temporary file122 temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')123 plt.savefig(temp_file.name, format='png', dpi=150, bbox_inches='tight')124 plt.close()125 126 return temp_file.name127 128 def create_interactive_html(self, graph: nx.DiGraph) -> str:129 """Create an interactive HTML visualization using vis.js."""130 if not graph.nodes():131 return "<div>No graph data to display</div>"132 133 # Convert graph to vis.js format134 nodes = []135 edges = []136 137 for node, data in graph.nodes(data=True):138 nodes.append({139 "id": node,140 "label": node,141 "color": self.color_map.get(data.get('type', 'UNKNOWN'), '#95A5A6'),142 "size": data.get('size', 20),143 "title": f"Type: {data.get('type', 'UNKNOWN')}<br>"144 f"Importance: {data.get('importance', 0.0):.2f}<br>"145 f"Description: {data.get('description', 'N/A')}"146 })147 148 for u, v, data in graph.edges(data=True):149 edges.append({150 "from": u,151 "to": v,152 "label": data.get('relationship', ''),153 "title": data.get('description', ''),154 "arrows": {"to": {"enabled": True}}155 })156 157 html_template = f"""158 <!DOCTYPE html>159 <html>160 <head>161 <script src="https://unpkg.com/vis-network/standalone/umd/vis-network.min.js"></script>162 <style>163 #mynetworkid {{164 width: 100%;165 height: 600px;166 border: 1px solid lightgray;167 }}168 </style>169 </head>170 <body>171 <div id="mynetworkid"></div>172 173 <script>174 var nodes = new vis.DataSet({json.dumps(nodes)});175 var edges = new vis.DataSet({json.dumps(edges)});176 var container = document.getElementById('mynetworkid');177 178 var data = {{179 nodes: nodes,180 edges: edges181 }};182 183 var options = {{184 nodes: {{185 shape: 'dot',186 scaling: {{187 min: 10,188 max: 30189 }},190 font: {{191 size: 12,192 face: 'Tahoma'193 }}194 }},195 edges: {{196 font: {{align: 'middle'}},197 color: {{color:'gray'}},198 arrows: {{to: {{enabled: true, scaleFactor: 1}}}}199 }},200 physics: {{201 enabled: true,202 stabilization: {{enabled: true, iterations: 200}}203 }},204 interaction: {{205 hover: true,206 tooltipDelay: 200207 }}208 }};209 210 var network = new vis.Network(container, data, options);211 </script>212 </body>213 </html>214 """215 216 return html_template217 218 def create_statistics_summary(self, graph: nx.DiGraph, stats: Dict[str, Any]) -> str:219 """Create a formatted statistics summary."""220 if not graph.nodes():221 return "No graph statistics available."222 223 # Entity type distribution224 type_counts = {}225 for node, data in graph.nodes(data=True):226 node_type = data.get('type', 'UNKNOWN')227 type_counts[node_type] = type_counts.get(node_type, 0) + 1228 229 # Relationship type distribution230 rel_counts = {}231 for u, v, data in graph.edges(data=True):232 rel_type = data.get('relationship', 'unknown')233 rel_counts[rel_type] = rel_counts.get(rel_type, 0) + 1234 235 summary = f"""236 ## Graph Statistics237 238 **Basic Metrics:**239 - Nodes: {stats['num_nodes']}240 - Edges: {stats['num_edges']}241 - Density: {stats['density']:.3f}242 - Connected: {'Yes' if stats['is_connected'] else 'No'}243 - Components: {stats['num_components']}244 - Average Degree: {stats['avg_degree']:.2f}245 246 **Entity Types:**247 """248 249 for entity_type, count in sorted(type_counts.items()):250 summary += f"\n- {entity_type}: {count}"251 252 summary += "\n\n**Relationship Types:**"253 for rel_type, count in sorted(rel_counts.items()):254 summary += f"\n- {rel_type}: {count}"255 256 return summary257 258 def create_entity_list(self, graph: nx.DiGraph, sort_by: str = "importance") -> str:259 """Create a formatted list of entities."""260 if not graph.nodes():261 return "No entities found."262 263 entities = []264 for node, data in graph.nodes(data=True):265 entities.append({266 'name': node,267 'type': data.get('type', 'UNKNOWN'),268 'importance': data.get('importance', 0.0),269 'description': data.get('description', 'N/A'),270 'connections': graph.degree(node)271 })272 273 # Sort entities274 if sort_by == "importance":275 entities.sort(key=lambda x: x['importance'], reverse=True)276 elif sort_by == "connections":277 entities.sort(key=lambda x: x['connections'], reverse=True)278 elif sort_by == "name":279 entities.sort(key=lambda x: x['name'])280 281 entity_list = "## Entities\n\n"282 for entity in entities:283 entity_list += f"""284**{entity['name']}** ({entity['type']})285- Importance: {entity['importance']:.2f}286- Connections: {entity['connections']}287- Description: {entity['description']}288 289"""290 291 return entity_list292 293 def get_layout_options(self) -> List[str]:294 """Get available layout options."""295 return ["spring", "circular", "shell", "kamada_kawai", "random"]296 297 def get_entity_types(self, graph: nx.DiGraph) -> List[str]:298 """Get unique entity types from the graph."""299 types = set()300 for node, data in graph.nodes(data=True):301 types.add(data.get('type', 'UNKNOWN'))302 return sorted(list(types))303 304 def create_plotly_interactive(self, graph: nx.DiGraph, layout_type: str = "spring") -> go.Figure:305 """Create an interactive Plotly visualization of the graph."""306 if not graph.nodes():307 # Return empty figure308 fig = go.Figure()309 fig.add_annotation(310 text="No graph data to display",311 xref="paper", yref="paper",312 x=0.5, y=0.5, xanchor='center', yanchor='middle',313 showarrow=False, font=dict(size=16)314 )315 return fig316 317 # Calculate layout318 pos = self._calculate_layout(graph, layout_type)319 320 # Prepare node data321 node_x = []322 node_y = []323 node_text = []324 node_info = []325 node_colors = []326 node_sizes = []327 328 for node in graph.nodes():329 x, y = pos[node]330 node_x.append(x)331 node_y.append(y)332 333 data = graph.nodes[node]334 node_type = data.get('type', 'UNKNOWN')335 importance = data.get('importance', 0.0)336 description = data.get('description', 'N/A')337 connections = graph.degree(node)338 339 node_text.append(node)340 node_info.append(341 f"<b>{node}</b><br>"342 f"Type: {node_type}<br>"343 f"Importance: {importance:.2f}<br>"344 f"Connections: {connections}<br>"345 f"Description: {description}"346 )347 node_colors.append(self.color_map.get(node_type, '#95A5A6'))348 node_sizes.append(max(10, data.get('size', 20)))349 350 # Prepare edge data351 edge_x = []352 edge_y = []353 edge_info = []354 355 for edge in graph.edges():356 x0, y0 = pos[edge[0]]357 x1, y1 = pos[edge[1]]358 edge_x.extend([x0, x1, None])359 edge_y.extend([y0, y1, None])360 361 edge_data = graph.edges[edge]362 relationship = edge_data.get('relationship', 'connected')363 edge_info.append(f"{edge[0]} → {edge[1]}<br>Relationship: {relationship}")364 365 # Create edge trace366 edge_trace = go.Scatter(367 x=edge_x, y=edge_y,368 line=dict(width=2, color='gray'),369 hoverinfo='none',370 mode='lines'371 )372 373 # Create node trace374 node_trace = go.Scatter(375 x=node_x, y=node_y,376 mode='markers+text',377 hoverinfo='text',378 text=node_text,379 hovertext=node_info,380 textposition="middle center",381 marker=dict(382 size=node_sizes,383 color=node_colors,384 line=dict(width=2, color='white')385 )386 )387 388 # Create figure389 fig = go.Figure(data=[edge_trace, node_trace],390 layout=go.Layout(391 title='Interactive Knowledge Graph',392 titlefont_size=16,393 showlegend=False,394 hovermode='closest',395 margin=dict(b=20,l=5,r=5,t=40),396 annotations=[ dict(397 text="Hover over nodes for details. Drag to pan, scroll to zoom.",398 showarrow=False,399 xref="paper", yref="paper",400 x=0.005, y=-0.002,401 xanchor='left', yanchor='bottom',402 font=dict(color="gray", size=12)403 )],404 xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),405 yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),406 plot_bgcolor='white'407 ))408 409 return fig410 411 def create_pyvis_interactive(self, graph: nx.DiGraph, layout_type: str = "spring") -> str:412 """Create an interactive pyvis visualization and return HTML file path."""413 if not graph.nodes():414 return self._create_empty_pyvis_graph()415 416 # Create pyvis network417 net = Network(height="600px", width="100%", bgcolor="#ffffff", font_color="black")418 419 # Configure physics420 net.set_options("""421 {422 "physics": {423 "enabled": true,424 "stabilization": {"enabled": true, "iterations": 200},425 "barnesHut": {426 "gravitationalConstant": -2000,427 "centralGravity": 0.3,428 "springLength": 95,429 "springConstant": 0.04,430 "damping": 0.09431 }432 },433 "interaction": {434 "hover": true,435 "tooltipDelay": 200,436 "hideEdgesOnDrag": false437 }438 }439 """)440 441 # Add nodes442 for node, data in graph.nodes(data=True):443 node_type = data.get('type', 'UNKNOWN')444 importance = data.get('importance', 0.0)445 description = data.get('description', 'N/A')446 connections = graph.degree(node)447 448 # Node properties449 color = self.color_map.get(node_type, '#95A5A6')450 size = max(10, data.get('size', 20))451 452 # Tooltip text453 title = f"""454 <b>{node}</b><br>455 Type: {node_type}<br>456 Importance: {importance:.2f}<br>457 Connections: {connections}<br>458 Description: {description}459 """460 461 net.add_node(node, label=node, title=title, color=color, size=size)462 463 # Add edges464 for u, v, data in graph.edges(data=True):465 relationship = data.get('relationship', 'connected')466 title = f"{u} → {v}<br>Relationship: {relationship}"467 468 net.add_edge(u, v, title=title, arrows="to", color="gray")469 470 # Save to temporary file471 temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.html', mode='w')472 net.save_graph(temp_file.name)473 temp_file.close()474 475 return temp_file.name476 477 def _create_empty_pyvis_graph(self) -> str:478 """Create an empty pyvis graph."""479 net = Network(height="600px", width="100%", bgcolor="#ffffff", font_color="black")480 net.add_node(1, label="No graph data", color="#cccccc")481 482 temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.html', mode='w')483 net.save_graph(temp_file.name)484 temp_file.close()485 486 return temp_file.name487 488 def get_visualization_options(self) -> List[str]:489 """Get available visualization types."""490 return ["matplotlib", "plotly", "pyvis", "vis.js"]491 492 def get_relationship_types(self, graph: nx.DiGraph) -> List[str]:493 """Get unique relationship types from the graph."""494 types = set()495 for u, v, data in graph.edges(data=True):496 types.add(data.get('relationship', 'unknown'))497 return sorted(list(types))498