fleet-memory/visualizations/interactive_graph.py
2025-10-30 12:53:12 +01:00

297 lines
9.9 KiB
Python

"""
Interactive HTML graph visualization of memory system.
Uses pyvis to create a smooth, interactive network graph that can be
explored in the browser. Shows all memory units and their links with weights.
"""
import psycopg2
from dotenv import load_dotenv
import os
from pyvis.network import Network
import networkx as nx
load_dotenv()
def create_interactive_graph():
"""Create an interactive HTML graph visualization."""
# Connect to database
conn = psycopg2.connect(os.getenv('DATABASE_URL'))
cursor = conn.cursor()
# Get all memory units (no agent_id filter)
cursor.execute("""
SELECT id, text, event_date, context
FROM memory_units
ORDER BY event_date
""")
units = cursor.fetchall()
# Get all links with weights (no agent_id filter)
cursor.execute("""
SELECT
ml.from_unit_id,
ml.to_unit_id,
ml.link_type,
ml.weight,
e.canonical_name as entity_name
FROM memory_links ml
LEFT JOIN entities e ON ml.entity_id = e.id
ORDER BY ml.link_type, ml.weight DESC
""")
links = cursor.fetchall()
# Get entity information (no agent_id filter)
cursor.execute("""
SELECT ue.unit_id, e.canonical_name, e.entity_type
FROM unit_entities ue
JOIN entities e ON ue.entity_id = e.id
ORDER BY ue.unit_id
""")
unit_entities = cursor.fetchall()
cursor.close()
conn.close()
# Build entity mapping
entity_map = {}
for unit_id, entity_name, entity_type in unit_entities:
if unit_id not in entity_map:
entity_map[unit_id] = []
entity_map[unit_id].append(f"{entity_name} ({entity_type})")
# Create pyvis network
net = Network(
height="900px",
width="100%",
bgcolor="#ffffff",
font_color="#000000",
heading="Entity-Aware Memory Graph - Interactive Visualization"
)
# Configure physics for smooth layout with performance optimizations
net.set_options("""
{
"nodes": {
"font": {
"size": 14,
"face": "Tahoma"
},
"borderWidth": 2,
"borderWidthSelected": 3
},
"edges": {
"smooth": {
"enabled": false
},
"font": {
"size": 10,
"align": "middle"
}
},
"physics": {
"enabled": true,
"stabilization": {
"enabled": true,
"iterations": 100,
"updateInterval": 10
},
"barnesHut": {
"gravitationalConstant": -12000,
"centralGravity": 0.2,
"springLength": 350,
"springConstant": 0.02,
"damping": 0.09,
"avoidOverlap": 0.8
},
"solver": "barnesHut",
"timestep": 0.5,
"adaptiveTimestep": true
},
"interaction": {
"hover": true,
"tooltipDelay": 100,
"navigationButtons": true,
"keyboard": true
}
}
""")
# Add nodes
for unit_id, text, event_date, context in units:
# Truncate text for display
display_text = text[:50] + "..." if len(text) > 50 else text
# Get entities
entities = entity_map.get(unit_id, [])
entity_str = "\\n".join(entities) if entities else "No entities"
# Build node label and title (hover)
label = display_text
title = f"""
<b>Text:</b> {text}<br>
<b>Date:</b> {event_date.date()}<br>
<b>Context:</b> {context}<br>
<b>Entities:</b> {entity_str}
"""
# Color by entity count
if len(entities) == 0:
color = "#e0e0e0" # Gray
size = 20
elif len(entities) == 1:
color = "#90caf9" # Light blue
size = 25
else:
color = "#42a5f5" # Dark blue
size = 30
net.add_node(
str(unit_id),
label=label,
title=title,
color=color,
size=size,
shape="box",
font={"color": "#000000"}
)
# Add edges with colors and weights
for from_id, to_id, link_type, weight, entity_name in links:
# Set color and style based on link type
if link_type == 'temporal':
color = "#00bcd4" # Cyan
dashes = [5, 5]
width = 0.5
label = f"T: {weight:.2f}"
elif link_type == 'semantic':
color = "#ff69b4" # Pink
dashes = False
width = 0.5
label = f"S: {weight:.2f}"
elif link_type == 'entity':
color = "#ffd700" # Gold
dashes = False
width = 0.8
label = f"{entity_name}: {weight:.2f}"
else:
color = "#999999"
dashes = False
width = 0.5
label = f"{weight:.2f}"
net.add_edge(
str(from_id),
str(to_id),
value=weight * 1, # Scale for visual thickness
color=color,
dashes=dashes,
width=width,
label=label,
title=f"{link_type.upper()}: {weight:.3f}" + (f" (Entity: {entity_name})" if entity_name else "")
)
# Add legend as HTML
legend_html = """
<div style="position: absolute; top: 80px; left: 10px; background: white;
padding: 15px; border: 2px solid #333; border-radius: 8px;
font-family: Tahoma; box-shadow: 2px 2px 8px rgba(0,0,0,0.3); z-index: 1000;">
<h3 style="margin-top: 0; border-bottom: 2px solid #333; padding-bottom: 5px;">Legend</h3>
<h4 style="margin-bottom: 5px;">Link Types:</h4>
<div style="margin-left: 10px;">
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 40px; height: 1px;
background: #00bcd4; border-top: 1px dashed #00bcd4;
vertical-align: middle;"></span>
<span style="margin-left: 10px;"><b>Temporal</b> - Time-based (cyan, dashed)</span>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 40px; height: 1px;
background: #ff69b4; vertical-align: middle;"></span>
<span style="margin-left: 10px;"><b>Semantic</b> - Meaning-based (pink, solid)</span>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 40px; height: 1.5px;
background: #ffd700; vertical-align: middle;"></span>
<span style="margin-left: 10px;"><b>Entity</b> - Same entity (gold)</span>
</div>
</div>
<h4 style="margin-bottom: 5px; margin-top: 15px;">Node Colors:</h4>
<div style="margin-left: 10px;">
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 20px; height: 20px;
background: #e0e0e0; border: 1px solid #999;
vertical-align: middle;"></span>
<span style="margin-left: 10px;">Gray - No entities</span>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 20px; height: 20px;
background: #90caf9; border: 1px solid #999;
vertical-align: middle;"></span>
<span style="margin-left: 10px;">Light Blue - 1 entity</span>
</div>
<div style="margin: 5px 0;">
<span style="display: inline-block; width: 20px; height: 20px;
background: #42a5f5; border: 1px solid #999;
vertical-align: middle;"></span>
<span style="margin-left: 10px;">Dark Blue - 2+ entities</span>
</div>
</div>
<div style="margin-top: 15px; padding-top: 10px; border-top: 1px solid #ccc;
font-size: 11px; color: #666;">
<b>Tip:</b> Hover over nodes/edges for details<br>
<b>Controls:</b> Drag to move, scroll to zoom
</div>
</div>
"""
# Generate the HTML
output_file = "memory_graph_interactive.html"
net.save_graph(output_file)
# Read the generated HTML and inject our legend
with open(output_file, 'r') as f:
html_content = f.read()
# Inject legend after the opening body tag
html_content = html_content.replace('<body>', '<body>' + legend_html)
# Add script to disable physics after stabilization for better performance
physics_script = """
<script type="text/javascript">
// Disable physics after initial stabilization for better performance
network.on("stabilizationIterationsDone", function () {
network.setOptions({ physics: false });
console.log("Physics disabled - graph should be much more responsive now!");
});
</script>
"""
html_content = html_content.replace('</body>', physics_script + '</body>')
# Write back
with open(output_file, 'w') as f:
f.write(html_content)
print(f"\n{'='*80}")
print("INTERACTIVE GRAPH GENERATED")
print(f"{'='*80}")
print(f"\nFile: {output_file}")
print(f"Units: {len(units)}")
print(f"Links: {len(links)}")
print("\nFeatures:")
print(" • Smooth, physics-based layout")
print(" • Interactive - drag nodes, zoom, pan")
print(" • Hover for details on nodes and edges")
print(" • Color-coded by link type and entity count")
print(" • Built-in navigation controls")
print(f"\n{'='*80}")
print(f"✓ Open {output_file} in your browser to explore!")
print(f"{'='*80}\n")
if __name__ == "__main__":
create_interactive_graph()