graph_visualizer.py

#!/usr/bin/env python3
"""
Interactive Graph Visualization Tool for Code Relationships
Web interface to explore the code graph generated by the RAG system
"""

import os
import sys
import json
import pickle
from typing import Dict, List, Any
from pathlib import Path

# Configuration
GRAPH_PATH = "/Users/dmitryminchuk/Projects/ai/mcp/ollama-rag/code_graph.pkl"

# Auto-activate virtual environment if not already active
def ensure_venv():
    """Ensure we're running in the virtual environment"""
    script_dir = os.path.dirname(os.path.abspath(__file__))
    venv_path = os.path.join(script_dir, "venv")
    venv_python = os.path.join(venv_path, "bin", "python")
    
    if hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix):
        return  # Already in venv
    
    if os.path.exists(venv_python):
        os.execv(venv_python, [venv_python, __file__] + sys.argv[1:])
    else:
        print(f"Error: Virtual environment not found at {venv_path}", file=sys.stderr)
        print("Please run: pip install flask", file=sys.stderr)
        sys.exit(1)

# Ensure virtual environment before importing dependencies
ensure_venv()

try:
    from flask import Flask, render_template, jsonify, request
    import networkx as nx
except ImportError:
    print("Missing dependencies. Please install:")
    print("pip install flask")
    sys.exit(1)

class GraphVisualizer:
    """Graph visualization and analysis"""
    
    def __init__(self, graph_path: str):
        self.graph_path = graph_path
        self.graph = None
        self.file_to_classes = {}
        self.class_to_file = {}
        self.load_graph()
    
    def load_graph(self):
        """Load the code graph from pickle file"""
        if not os.path.exists(self.graph_path):
            print(f"Graph file not found: {self.graph_path}")
            print("Please run the MCP server first to generate the graph")
            return
        
        try:
            with open(self.graph_path, 'rb') as f:
                data = pickle.load(f)
                self.graph = data['graph']
                self.file_to_classes = data['file_to_classes']
                self.class_to_file = data['class_to_file']
            
            print(f"Loaded graph with {len(self.graph.nodes)} nodes and {len(self.graph.edges)} edges")
        except Exception as e:
            print(f"Error loading graph: {e}")
            self.graph = nx.DiGraph()
    
    def get_graph_stats(self) -> Dict[str, Any]:
        """Get basic statistics about the graph"""
        if not self.graph:
            return {"error": "Graph not loaded"}
        
        # Count relationship types
        relationship_counts = {}
        for _, _, data in self.graph.edges(data=True):
            rel_type = data.get('type', 'unknown')
            relationship_counts[rel_type] = relationship_counts.get(rel_type, 0) + 1
        
        # Get file types
        file_types = {}
        for node in self.graph.nodes():
            if isinstance(node, str) and '.' in node:
                ext = node.split('.')[-1]
                file_types[ext] = file_types.get(ext, 0) + 1
        
        return {
            "nodes": len(self.graph.nodes),
            "edges": len(self.graph.edges),
            "relationship_types": relationship_counts,
            "file_types": file_types,
            "classes": len(self.class_to_file),
            "files_with_classes": len(self.file_to_classes)
        }
    
    def get_subgraph(self, center_node: str, radius: int = 2) -> Dict[str, Any]:
        """Get a subgraph around a center node"""
        if not self.graph or center_node not in self.graph:
            return {"nodes": [], "links": []}
        
        # Find nodes within radius
        nodes_in_radius = set([center_node])
        current_nodes = set([center_node])
        
        for _ in range(radius):
            next_nodes = set()
            for node in current_nodes:
                # Add neighbors (both directions)
                next_nodes.update(self.graph.neighbors(node))
                next_nodes.update(self.graph.predecessors(node))
            
            current_nodes = next_nodes - nodes_in_radius
            nodes_in_radius.update(current_nodes)
        
        # Create subgraph
        subgraph = self.graph.subgraph(nodes_in_radius)
        
        # Convert to D3.js format
        nodes = []
        for node in subgraph.nodes():
            node_type = "file" if isinstance(node, str) and "/" in node else "class"
            size = 10 if node == center_node else 5
            
            nodes.append({
                "id": node,
                "name": os.path.basename(node) if "/" in node else node,
                "type": node_type,
                "size": size,
                "color": self._get_node_color(node_type, node == center_node)
            })
        
        links = []
        for source, target, data in subgraph.edges(data=True):
            links.append({
                "source": source,
                "target": target,
                "type": data.get('type', 'unknown'),
                "context": data.get('context', ''),
                "line_number": data.get('line_number', 0)
            })
        
        return {
            "nodes": nodes,
            "links": links,
            "center": center_node
        }
    
    def search_nodes(self, query: str, limit: int = 20) -> List[Dict[str, Any]]:
        """Search for nodes matching the query"""
        if not self.graph:
            return []
        
        query_lower = query.lower()
        matches = []
        
        for node in self.graph.nodes():
            if query_lower in str(node).lower():
                node_type = "file" if isinstance(node, str) and "/" in node else "class"
                matches.append({
                    "id": node,
                    "name": os.path.basename(node) if "/" in node else node,
                    "type": node_type,
                    "full_path": node
                })
        
        return matches[:limit]
    
    def get_node_details(self, node_id: str) -> Dict[str, Any]:
        """Get detailed information about a node"""
        if not self.graph or node_id not in self.graph:
            return {"error": "Node not found"}
        
        # Get incoming and outgoing relationships
        incoming = []
        for source, target, data in self.graph.in_edges(node_id, data=True):
            incoming.append({
                "source": source,
                "type": data.get('type', 'unknown'),
                "context": data.get('context', ''),
                "line_number": data.get('line_number', 0)
            })
        
        outgoing = []
        for source, target, data in self.graph.out_edges(node_id, data=True):
            outgoing.append({
                "target": target,
                "type": data.get('type', 'unknown'),
                "context": data.get('context', ''),
                "line_number": data.get('line_number', 0)
            })
        
        return {
            "id": node_id,
            "name": os.path.basename(node_id) if "/" in node_id else node_id,
            "type": "file" if "/" in node_id else "class",
            "incoming": incoming,
            "outgoing": outgoing,
            "degree": len(incoming) + len(outgoing)
        }
    
    def _get_node_color(self, node_type: str, is_center: bool = False) -> str:
        """Get color for node based on type"""
        if is_center:
            return "#ff6b6b"  # Red for center node
        elif node_type == "file":
            return "#4ecdc4"  # Teal for files
        else:
            return "#45b7d1"  # Blue for classes


# Flask application
app = Flask(__name__)
visualizer = GraphVisualizer(GRAPH_PATH)

@app.route('/')
def index():
    """Main page with graph visualization"""
    return render_template('graph.html')

@app.route('/api/stats')
def get_stats():
    """Get graph statistics"""
    return jsonify(visualizer.get_graph_stats())

@app.route('/api/subgraph')
def get_subgraph():
    """Get subgraph around a center node"""
    center = request.args.get('center')
    radius = int(request.args.get('radius', 2))
    
    if not center:
        return jsonify({"error": "Center node required"}), 400
    
    return jsonify(visualizer.get_subgraph(center, radius))

@app.route('/api/search')
def search_nodes():
    """Search for nodes"""
    query = request.args.get('q', '')
    limit = int(request.args.get('limit', 20))
    
    if not query:
        return jsonify([])
    
    return jsonify(visualizer.search_nodes(query, limit))

@app.route('/api/node/<path:node_id>')
def get_node_details(node_id):
    """Get detailed information about a node"""
    return jsonify(visualizer.get_node_details(node_id))


if __name__ == '__main__':
    # Create templates directory if it doesn't exist
    os.makedirs('templates', exist_ok=True)
    
    print("Starting Graph Visualizer...")
    print("Open http://localhost:5001 in your browser")
    print("Press Ctrl+C to stop")
    
    app.run(debug=True, host='0.0.0.0', port=5001)