394 lines
13 KiB
Python
394 lines
13 KiB
Python
"""
|
|
Memory visualization module.
|
|
|
|
Provides visual representations of memory networks and search paths.
|
|
"""
|
|
import time
|
|
from typing import List, Dict, Any, Optional, Tuple
|
|
import networkx as nx
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.patches import FancyBboxPatch
|
|
from rich.console import Console
|
|
from rich.table import Table
|
|
from rich.panel import Panel
|
|
from rich.layout import Layout
|
|
from rich.live import Live
|
|
from rich.text import Text
|
|
from rich import box
|
|
|
|
|
|
class MemoryVisualizer:
|
|
"""
|
|
Visualizes memory networks and search paths.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the visualizer."""
|
|
self.console = Console()
|
|
|
|
def visualize_memory_graph(
|
|
self,
|
|
units: List[Dict[str, Any]],
|
|
links: List[Dict[str, Any]],
|
|
output_file: str = "memory_graph.png",
|
|
highlight_nodes: Optional[List[str]] = None,
|
|
):
|
|
"""
|
|
Create a visual representation of the memory graph.
|
|
|
|
Args:
|
|
units: List of memory units (id, text, context, etc.)
|
|
links: List of links (from_unit_id, to_unit_id, link_type, weight)
|
|
output_file: Output file path for the visualization
|
|
highlight_nodes: Optional list of node IDs to highlight
|
|
"""
|
|
# Create directed graph
|
|
G = nx.DiGraph()
|
|
|
|
# Add nodes
|
|
node_labels = {}
|
|
for unit in units:
|
|
unit_id = str(unit['id'])
|
|
# Truncate text for display
|
|
label = unit['text'][:40] + "..." if len(unit['text']) > 40 else unit['text']
|
|
G.add_node(unit_id)
|
|
node_labels[unit_id] = label
|
|
|
|
# Add edges
|
|
temporal_edges = []
|
|
semantic_edges = []
|
|
for link in links:
|
|
from_id = str(link['from_unit_id'])
|
|
to_id = str(link['to_unit_id'])
|
|
weight = link['weight']
|
|
link_type = link['link_type']
|
|
|
|
if link_type == 'temporal':
|
|
temporal_edges.append((from_id, to_id, weight))
|
|
else: # semantic
|
|
semantic_edges.append((from_id, to_id, weight))
|
|
|
|
G.add_edge(from_id, to_id, weight=weight, type=link_type)
|
|
|
|
# Create figure
|
|
fig, ax = plt.subplots(figsize=(20, 14))
|
|
ax.set_facecolor('#1a1a2e')
|
|
fig.patch.set_facecolor('#0f0f1e')
|
|
|
|
# Use spring layout for better visualization
|
|
pos = nx.spring_layout(G, k=2, iterations=50, seed=42)
|
|
|
|
# Draw temporal edges (blue)
|
|
if temporal_edges:
|
|
nx.draw_networkx_edges(
|
|
G, pos,
|
|
edgelist=[(e[0], e[1]) for e in temporal_edges],
|
|
edge_color='#4ecdc4',
|
|
alpha=0.6,
|
|
width=2,
|
|
arrows=True,
|
|
arrowsize=15,
|
|
arrowstyle='->',
|
|
connectionstyle='arc3,rad=0.1',
|
|
ax=ax
|
|
)
|
|
|
|
# Draw semantic edges (purple)
|
|
if semantic_edges:
|
|
nx.draw_networkx_edges(
|
|
G, pos,
|
|
edgelist=[(e[0], e[1]) for e in semantic_edges],
|
|
edge_color='#ff6b9d',
|
|
alpha=0.6,
|
|
width=2,
|
|
arrows=True,
|
|
arrowsize=15,
|
|
arrowstyle='->',
|
|
connectionstyle='arc3,rad=0.1',
|
|
ax=ax
|
|
)
|
|
|
|
# Determine node colors
|
|
node_colors = []
|
|
for node in G.nodes():
|
|
if highlight_nodes and node in highlight_nodes:
|
|
node_colors.append('#ffd93d') # Yellow for highlighted
|
|
else:
|
|
node_colors.append('#6c63ff') # Purple for normal
|
|
|
|
# Draw nodes
|
|
nx.draw_networkx_nodes(
|
|
G, pos,
|
|
node_color=node_colors,
|
|
node_size=3000,
|
|
alpha=0.9,
|
|
ax=ax
|
|
)
|
|
|
|
# Draw labels
|
|
nx.draw_networkx_labels(
|
|
G, pos,
|
|
node_labels,
|
|
font_size=8,
|
|
font_color='white',
|
|
font_weight='bold',
|
|
ax=ax
|
|
)
|
|
|
|
# Add legend
|
|
legend_elements = [
|
|
plt.Line2D([0], [0], color='#4ecdc4', lw=2, label='Temporal Links'),
|
|
plt.Line2D([0], [0], color='#ff6b9d', lw=2, label='Semantic Links'),
|
|
plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#6c63ff',
|
|
markersize=10, label='Memory Unit', linestyle=''),
|
|
]
|
|
if highlight_nodes:
|
|
legend_elements.append(
|
|
plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#ffd93d',
|
|
markersize=10, label='Highlighted', linestyle='')
|
|
)
|
|
|
|
ax.legend(handles=legend_elements, loc='upper left', facecolor='#2d2d44',
|
|
edgecolor='white', fontsize=10, labelcolor='white')
|
|
|
|
# Title
|
|
ax.set_title('Memory Network Graph\nTemporal + Semantic Architecture',
|
|
color='white', fontsize=16, fontweight='bold', pad=20)
|
|
|
|
ax.axis('off')
|
|
plt.tight_layout()
|
|
plt.savefig(output_file, dpi=150, facecolor='#0f0f1e')
|
|
plt.close()
|
|
|
|
self.console.print(f"[green]✓[/green] Memory graph saved to [cyan]{output_file}[/cyan]")
|
|
|
|
|
|
class LiveSearchTracer:
|
|
"""
|
|
Live tracer for search operations showing spreading activation in real-time.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the live tracer."""
|
|
self.console = Console()
|
|
self.visited_nodes = []
|
|
self.current_node = None
|
|
self.search_results = []
|
|
self.query = ""
|
|
self.budget_used = 0
|
|
self.budget_total = 0
|
|
|
|
def start_search(self, query: str, budget: int):
|
|
"""
|
|
Start a new search trace.
|
|
|
|
Args:
|
|
query: Search query
|
|
budget: Thinking budget
|
|
"""
|
|
self.query = query
|
|
self.budget_total = budget
|
|
self.budget_used = 0
|
|
self.visited_nodes = []
|
|
self.current_node = None
|
|
self.search_results = []
|
|
|
|
def visit_node(
|
|
self,
|
|
node_id: str,
|
|
text: str,
|
|
activation: float,
|
|
recency: float,
|
|
frequency: float,
|
|
weight: float,
|
|
is_entry_point: bool = False,
|
|
):
|
|
"""
|
|
Record a node visit.
|
|
|
|
Args:
|
|
node_id: Node ID
|
|
text: Node text
|
|
activation: Activation strength
|
|
recency: Recency weight
|
|
frequency: Frequency weight
|
|
weight: Combined weight
|
|
is_entry_point: Whether this is an entry point
|
|
"""
|
|
self.current_node = {
|
|
'id': node_id,
|
|
'text': text,
|
|
'activation': activation,
|
|
'recency': recency,
|
|
'frequency': frequency,
|
|
'weight': weight,
|
|
'is_entry_point': is_entry_point,
|
|
}
|
|
self.visited_nodes.append(self.current_node)
|
|
self.budget_used += 1
|
|
|
|
def add_result(
|
|
self,
|
|
text: str,
|
|
weight: float,
|
|
activation: float,
|
|
recency: float,
|
|
frequency: float,
|
|
):
|
|
"""
|
|
Add a search result.
|
|
|
|
Args:
|
|
text: Result text
|
|
weight: Combined weight
|
|
activation: Activation strength
|
|
recency: Recency weight
|
|
frequency: Frequency weight
|
|
"""
|
|
self.search_results.append({
|
|
'text': text,
|
|
'weight': weight,
|
|
'activation': activation,
|
|
'recency': recency,
|
|
'frequency': frequency,
|
|
})
|
|
|
|
def render_live(self) -> Layout:
|
|
"""
|
|
Render the current state.
|
|
|
|
Returns:
|
|
Rich Layout with current state
|
|
"""
|
|
layout = Layout()
|
|
layout.split_column(
|
|
Layout(name="header", size=3),
|
|
Layout(name="body"),
|
|
Layout(name="footer", size=5)
|
|
)
|
|
|
|
# Header
|
|
header_text = Text()
|
|
header_text.append("🔍 ", style="bold cyan")
|
|
header_text.append(f"Query: ", style="bold white")
|
|
header_text.append(f"{self.query}", style="bold yellow")
|
|
layout["header"].update(Panel(header_text, style="cyan"))
|
|
|
|
# Body - split into current node and visited
|
|
layout["body"].split_row(
|
|
Layout(name="current", ratio=1),
|
|
Layout(name="path", ratio=1),
|
|
)
|
|
|
|
# Current node
|
|
if self.current_node:
|
|
current_table = Table(
|
|
title="Current Node",
|
|
show_header=False,
|
|
box=box.ROUNDED,
|
|
style="green"
|
|
)
|
|
current_table.add_column("Key", style="cyan")
|
|
current_table.add_column("Value", style="white")
|
|
|
|
status = "🎯 ENTRY POINT" if self.current_node['is_entry_point'] else "🔄 EXPLORING"
|
|
current_table.add_row("Status", status)
|
|
current_table.add_row("Text", self.current_node['text'][:50] + "...")
|
|
current_table.add_row(
|
|
"Weights",
|
|
f"A:{self.current_node['activation']:.2f} "
|
|
f"R:{self.current_node['recency']:.2f} "
|
|
f"F:{self.current_node['frequency']:.2f}"
|
|
)
|
|
current_table.add_row(
|
|
"Combined",
|
|
f"[bold yellow]{self.current_node['weight']:.3f}[/bold yellow]"
|
|
)
|
|
|
|
layout["current"].update(Panel(current_table, border_style="green"))
|
|
else:
|
|
layout["current"].update(Panel("Initializing...", border_style="dim"))
|
|
|
|
# Visited path
|
|
path_table = Table(
|
|
title=f"Visited Nodes ({len(self.visited_nodes)})",
|
|
box=box.SIMPLE,
|
|
show_header=True,
|
|
style="blue"
|
|
)
|
|
path_table.add_column("#", style="dim", width=4)
|
|
path_table.add_column("Text", style="white", width=35)
|
|
path_table.add_column("Weight", justify="right", style="yellow", width=8)
|
|
path_table.add_column("Type", style="cyan", width=8)
|
|
|
|
for i, node in enumerate(reversed(self.visited_nodes[-10:])): # Last 10
|
|
node_type = "ENTRY" if node['is_entry_point'] else "SPREAD"
|
|
path_table.add_row(
|
|
str(len(self.visited_nodes) - i),
|
|
node['text'][:32] + "...",
|
|
f"{node['weight']:.3f}",
|
|
node_type
|
|
)
|
|
|
|
layout["path"].update(Panel(path_table, border_style="blue"))
|
|
|
|
# Footer - progress bar
|
|
progress = self.budget_used / self.budget_total if self.budget_total > 0 else 0
|
|
bar_width = 50
|
|
filled = int(bar_width * progress)
|
|
bar = "█" * filled + "░" * (bar_width - filled)
|
|
|
|
footer_text = Text()
|
|
footer_text.append(f"Progress: ", style="bold white")
|
|
footer_text.append(bar, style="yellow")
|
|
footer_text.append(f" {self.budget_used}/{self.budget_total}", style="bold cyan")
|
|
footer_text.append(f" ({progress*100:.1f}%)", style="dim")
|
|
|
|
layout["footer"].update(Panel(footer_text, style="yellow"))
|
|
|
|
return layout
|
|
|
|
def show_final_results(self):
|
|
"""
|
|
Show final search results in a nice table.
|
|
"""
|
|
self.console.print("\n")
|
|
results_table = Table(
|
|
title="🎯 Search Results",
|
|
show_header=True,
|
|
header_style="bold magenta",
|
|
box=box.DOUBLE_EDGE,
|
|
title_style="bold white"
|
|
)
|
|
|
|
results_table.add_column("Rank", style="cyan", justify="center", width=6)
|
|
results_table.add_column("Text", style="white", width=50)
|
|
results_table.add_column("Weight", justify="right", style="yellow", width=8)
|
|
results_table.add_column("A", justify="right", style="green", width=6)
|
|
results_table.add_column("R", justify="right", style="blue", width=6)
|
|
results_table.add_column("F", justify="right", style="magenta", width=6)
|
|
|
|
for i, result in enumerate(self.search_results, 1):
|
|
rank_style = "bold yellow" if i <= 3 else "cyan"
|
|
results_table.add_row(
|
|
f"#{i}",
|
|
result['text'][:47] + "...",
|
|
f"{result['weight']:.3f}",
|
|
f"{result['activation']:.2f}",
|
|
f"{result['recency']:.2f}",
|
|
f"{result['frequency']:.2f}",
|
|
style=rank_style if i <= 3 else None
|
|
)
|
|
|
|
self.console.print(results_table)
|
|
|
|
# Summary stats
|
|
summary = Table.grid(padding=(0, 2))
|
|
summary.add_column(style="bold cyan")
|
|
summary.add_column(style="white")
|
|
|
|
summary.add_row("Total nodes visited:", f"{len(self.visited_nodes)}")
|
|
summary.add_row("Budget used:", f"{self.budget_used}/{self.budget_total}")
|
|
summary.add_row("Results found:", f"{len(self.search_results)}")
|
|
|
|
self.console.print(Panel(summary, title="Summary", border_style="green", padding=(1, 2)))
|