#!/usr/bin/env python3
"""
Enhanced Interactive Knowledge Graph Visualizer Module for QAFD_RAG
Creates an HTML file with interactive graph showing all nodes and relationships
with improved color schemes for different components
This module is designed to be imported by the main QAFD_RAG script.
"""
import networkx as nx
from pyvis.network import Network
import json
from pathlib import Path
def create_interactive_kg_visualization(graphml_path, output_html="knowledge_graph_interactive.html"):
"""
Create an interactive HTML visualization of the knowledge graph with enhanced colors
Args:
graphml_path (str): Path to the GraphML file
output_html (str): Output HTML file path
Returns:
bool: True if successful, False otherwise
"""
print(f"Loading graph from: {graphml_path}")
# Load the GraphML file
try:
G = nx.read_graphml(graphml_path)
print(f"ā
Graph loaded successfully!")
print(f" Nodes: {G.number_of_nodes()}")
print(f" Edges: {G.number_of_edges()}")
# ADDED: Debug edges
if G.number_of_edges() > 0:
print("Sample edges:")
for i, (source, target, data) in enumerate(G.edges(data=True)):
if i < 3: # Show first 3 edges
print(f" {source} -> {target}: {data}")
else:
print("WARNING: No edges found in GraphML file!")
except Exception as e:
print(f"ā Error loading graph: {e}")
return False
# Create Pyvis network with enhanced styling
net = Network(
height="900px",
width="100%",
bgcolor="#0a0a0a", # Darker background for better contrast
font_color="white",
directed=False,
notebook=False
)
# Enhanced color schemes
node_colors = {
'entity': {
'color': '#FF4757', # Bright red for entities
'border': '#FF3742',
'highlight': '#FF6B7A'
},
'chunk': {
'color': '#2ED573', # Bright green for chunks
'border': '#20BF6B',
'highlight': '#54E091'
},
'table': {
'color': '#3742FA', # Electric blue for tables
'border': '#2F3542',
'highlight': '#5352ED'
},
'column': {
'color': '#FFA726', # Orange for columns
'border': '#FF9800',
'highlight': '#FFB74D'
},
'complete_table': { # ADDED: Handle complete_table type
'color': '#3742FA', # Same as table
'border': '#2F3542',
'highlight': '#5352ED'
},
'default': {
'color': '#A4B0BE', # Light gray for others
'border': '#747D8C',
'highlight': '#DDD6FE'
}
}
# Edge color schemes
edge_colors = {
'primary_key': '#FF6B35', # Orange-red for primary key edges
'foreign_key': '#4834D4', # Purple for foreign key edges
'contains': '#00D2D3', # Cyan for containment relationships
'belongs_to': '#FF9FF3', # Pink for belongs_to relationships
'references': '#54A0FF', # Light blue for references
'has': '#5F27CD', # Deep purple for has relationships
'table_structure': '#00D2D3', # Same as contains
'default': '#57606F' # Gray for standard edges
}
print("Adding nodes with enhanced styling...")
# First, let's analyze what node types we actually have
actual_node_types = set()
entity_types = set()
for node_id, node_data in G.nodes(data=True):
node_type = node_data.get('type', 'unknown')
entity_type = node_data.get('entity_type', 'unknown')
actual_node_types.add(node_type)
entity_types.add(entity_type)
print(f"Detected 'type' attributes: {actual_node_types}")
print(f"Detected 'entity_type' attributes: {entity_types}")
for node_id, node_data in G.nodes(data=True):
# Read the entity_type directly from the data
entity_type = node_data.get('entity_type', '').lower()
node_type = node_data.get('type', '').lower()
# Use entity_type as the primary classifier
if entity_type in ['table', 'complete_table']:
detected_type = 'table'
elif entity_type == 'column':
detected_type = 'column'
elif entity_type in ['entity', 'chunk']:
detected_type = entity_type
elif node_type in ['table', 'column', 'entity', 'chunk']:
detected_type = node_type
else:
detected_type = 'default'
color_scheme = node_colors.get(detected_type, node_colors['default'])
# Create detailed hover info
title = f"Node ID: {node_id}
"
# title += f"Type: {detected_type.title()}
"
# Show entity_type if it exists
if 'entity_type' in node_data:
title += f"Entity Type: {node_data['entity_type']}
"
# Add all other attributes
for key, value in node_data.items():
if key not in ['type', 'entity_type']:
# Truncate long values for readability
display_value = str(value)[:100] + "..." if len(str(value)) > 100 else str(value)
title += f"{key.title()}: {display_value}
"
# Determine node size based on detected type
size_mapping = {
'table': 35, # Largest for tables
'column': 25, # Medium for columns
'entity': 30, # Large for entities
'chunk': 20, # Smaller for chunks
'default': 18 # Smallest for others
}
# Add node with enhanced styling
net.add_node(
str(node_id),
label=str(node_id)[:25] + ("..." if len(str(node_id)) > 25 else ""),
title=title,
color={
'background': color_scheme['color'],
'border': color_scheme['border'],
'highlight': {
'background': color_scheme['highlight'],
'border': color_scheme['border']
}
},
size=size_mapping.get(detected_type, size_mapping['default']),
font={'size': 14, 'color': 'white', 'face': 'arial'},
borderWidth=3,
shadow={'enabled': True, 'color': 'rgba(0,0,0,0.5)', 'size': 10}
)
print("Adding edges with relationship-based coloring...")
# First, analyze actual edge types
actual_edge_types = set()
for source, target, edge_data in G.edges(data=True):
# Try multiple ways to get relationship info
relationship = (
edge_data.get('relationship') or
edge_data.get('description') or
edge_data.get('keywords') or
edge_data.get('label') or
'default'
)
actual_edge_types.add(str(relationship).lower())
print(f"Detected edge types: {actual_edge_types}")
edges_added = 0
for source, target, edge_data in G.edges(data=True):
# Try to get meaningful relationship info
relationship = (
edge_data.get('relationship') or
edge_data.get('description') or
edge_data.get('keywords') or
edge_data.get('label') or
'default'
)
# Convert to string and normalize
relationship_str = str(relationship).lower()
# Determine edge color and width based on relationship type
edge_color = edge_colors['default']
edge_width = 2
# Enhanced relationship detection with keyword matching
if 'foreign_key' in relationship_str or 'references' in relationship_str:
edge_color = edge_colors['foreign_key']
edge_width = 3
elif 'primary_key' in relationship_str:
edge_color = edge_colors['primary_key']
edge_width = 4
elif 'contains' in relationship_str or 'table_structure' in relationship_str:
edge_color = edge_colors['contains']
edge_width = 3
elif 'belongs_to' in relationship_str:
edge_color = edge_colors['belongs_to']
edge_width = 2
elif 'has' in relationship_str:
edge_color = edge_colors['has']
edge_width = 2
else:
# For any other specific relationships
edge_color = '#FFD93D' # Bright yellow for other relationships
edge_width = 2
# Create detailed edge title
title = f"Relationship: {relationship}
"
title += f"From: {source}
"
title += f"To: {target}
"
# Add other edge attributes
for key, value in edge_data.items():
if key not in ['relationship', 'label', 'description', 'keywords']:
title += f"{key.title()}: {str(value)}
"
# Add edge with enhanced styling
net.add_edge(
str(source),
str(target),
label=str(relationship)[:25] if relationship != 'default' else '', # Show label only if meaningful
title=title,
color={
'color': edge_color,
'highlight': edge_color,
'hover': edge_color,
'opacity': 0.8
},
width=edge_width,
font={'size': 11, 'color': 'white'}
)
edges_added += 1
print(f"Successfully added {edges_added} edges to visualization")
# Enhanced physics and layout configuration
net.set_options("""
var options = {
"physics": {
"enabled": true,
"stabilization": {
"iterations": 300,
"updateInterval": 25
},
"barnesHut": {
"gravitationalConstant": -2000,
"centralGravity": 0.2,
"springLength": 120,
"springConstant": 0.05,
"damping": 0.15,
"avoidOverlap": 0.2
}
},
"nodes": {
"font": {
"size": 14,
"color": "white",
"face": "arial",
"strokeWidth": 2,
"strokeColor": "black"
},
"borderWidth": 3,
"shadow": {
"enabled": true,
"color": "rgba(0,0,0,0.5)",
"size": 10
},
"chosen": {
"node": true,
"label": true
}
},
"edges": {
"font": {
"size": 11,
"color": "white",
"face": "arial",
"strokeWidth": 1,
"strokeColor": "black"
},
"smooth": {
"type": "continuous",
"roundness": 0.5
},
"shadow": {
"enabled": true,
"color": "rgba(0,0,0,0.3)"
},
"chosen": {
"edge": true,
"label": true
}
},
"interaction": {
"hover": true,
"tooltipDelay": 200,
"hideEdgesOnDrag": false,
"hideNodesOnDrag": false
},
"layout": {
"improvedLayout": true,
"hierarchical": false
}
}
""")
# Generate and save the HTML
try:
print("Generating enhanced HTML visualization...")
html_string = net.generate_html()
# Add custom CSS for better styling
custom_css = """
"""
# Add legend HTML with correct colors
legend_html = """