mcp-vector-search 1.0.3__py3-none-any.whl → 1.1.22__py3-none-any.whl

mcp_vector_search/cli/commands/visualize/server.py

@@ -2,6 +2,7 @@
 
 This module handles running the local HTTP server to serve the
 D3.js visualization interface with chunked transfer for large JSON files.
+Uses orjson for 5-10x faster JSON serialization.
 """
 
 import asyncio
@@ -10,6 +11,7 @@ import webbrowser
 from collections.abc import AsyncGenerator
 from pathlib import Path
 
+import orjson
 import uvicorn
 from fastapi import FastAPI, Response
 from fastapi.responses import FileResponse, StreamingResponse
@@ -17,15 +19,17 @@ from fastapi.staticfiles import StaticFiles
 from rich.console import Console
 from rich.panel import Panel
 
+from mcp_vector_search import __version__
+
 console = Console()
 
 
-def find_free_port(start_port: int = 8080, end_port: int = 8099) -> int:
+def find_free_port(start_port: int = 8501, end_port: int = 8599) -> int:
     """Find a free port in the given range.
 
     Args:
-        start_port: Starting port number to check
-        end_port: Ending port number to check
+        start_port: Starting port number to check (default: 8501)
+        end_port: Ending port number to check (default: 8599)
 
     Returns:
         First available port in the range
@@ -80,6 +84,39 @@ def create_app(viz_dir: Path) -> FastAPI:
     """
     app = FastAPI(title="MCP Vector Search Visualization")
 
+    @app.get("/api/graph-status")
+    async def graph_status() -> Response:
+        """Get graph data generation status.
+
+        Returns:
+            JSON response with ready flag and file size
+        """
+        graph_file = viz_dir / "chunk-graph.json"
+
+        if not graph_file.exists():
+            return Response(
+                content='{"ready": false, "size": 0}',
+                media_type="application/json",
+                headers={"Cache-Control": "no-cache"},
+            )
+
+        try:
+            size = graph_file.stat().st_size
+            # Consider graph ready if file exists and has content (>100 bytes)
+            is_ready = size > 100
+            return Response(
+                content=f'{{"ready": {str(is_ready).lower()}, "size": {size}}}',
+                media_type="application/json",
+                headers={"Cache-Control": "no-cache"},
+            )
+        except Exception as e:
+            console.print(f"[red]Error checking graph status: {e}[/red]")
+            return Response(
+                content='{"ready": false, "size": 0}',
+                media_type="application/json",
+                headers={"Cache-Control": "no-cache"},
+            )
+
     @app.get("/api/graph")
     async def get_graph_data() -> Response:
         """Get graph data for D3 tree visualization.
@@ -97,14 +134,12 @@ def create_app(viz_dir: Path) -> FastAPI:
             )
 
         try:
-            import json
-
-            with open(graph_file) as f:
-                data = json.load(f)
+            with open(graph_file, "rb") as f:
+                data = orjson.loads(f.read())
 
-            # Return nodes and links
+            # Return nodes and links using orjson for fast serialization
             return Response(
-                content=json.dumps(
+                content=orjson.dumps(
                     {"nodes": data.get("nodes", []), "links": data.get("links", [])}
                 ),
                 media_type="application/json",
@@ -118,6 +153,262 @@ def create_app(viz_dir: Path) -> FastAPI:
                 media_type="application/json",
             )
 
+    @app.get("/api/relationships/{chunk_id}")
+    async def get_chunk_relationships(chunk_id: str) -> Response:
+        """Get all relationships for a chunk (semantic + callers) on-demand.
+
+        Lazy loads relationships when user expands a node, avoiding expensive
+        upfront computation. Results are cached in-memory for the session.
+
+        Args:
+            chunk_id: The chunk ID to find relationships for
+
+        Returns:
+            JSON response with semantic neighbors and callers
+        """
+        graph_file = viz_dir / "chunk-graph.json"
+
+        if not graph_file.exists():
+            return Response(
+                content='{"error": "Graph data not found"}',
+                status_code=404,
+                media_type="application/json",
+            )
+
+        try:
+            import ast
+
+            with open(graph_file, "rb") as f:
+                data = orjson.loads(f.read())
+
+            # Find the target chunk
+            target_node = None
+            for node in data.get("nodes", []):
+                if node.get("id") == chunk_id:
+                    target_node = node
+                    break
+
+            if not target_node:
+                return Response(
+                    content='{"error": "Chunk not found"}',
+                    status_code=404,
+                    media_type="application/json",
+                )
+
+            function_name = target_node.get("function_name") or target_node.get(
+                "class_name"
+            )
+            target_file = target_node.get("file_path", "")
+            target_content = target_node.get("content", "")
+
+            # Compute callers (who calls this function)
+            callers = []
+
+            def extract_calls(code: str) -> set[str]:
+                calls = set()
+                try:
+                    tree = ast.parse(code)
+                    for node in ast.walk(tree):
+                        if isinstance(node, ast.Call):
+                            if isinstance(node.func, ast.Name):
+                                calls.add(node.func.id)
+                            elif isinstance(node.func, ast.Attribute):
+                                calls.add(node.func.attr)
+                except SyntaxError:
+                    pass
+                return calls
+
+            if function_name:
+                for node in data.get("nodes", []):
+                    if node.get("type") != "chunk":
+                        continue
+                    node_file = node.get("file_path", "")
+                    if node_file == target_file:
+                        continue
+                    content = node.get("content", "")
+                    if function_name in extract_calls(content):
+                        caller_name = node.get("function_name") or node.get(
+                            "class_name"
+                        )
+                        if caller_name == "__init__":
+                            continue
+                        callers.append(
+                            {
+                                "id": node.get("id"),
+                                "name": caller_name
+                                or f"chunk_{node.get('start_line', 0)}",
+                                "file": node_file,
+                                "type": node.get("chunk_type", "code"),
+                            }
+                        )
+
+            # Compute semantic neighbors (similar code)
+            # Simple approach: find chunks with similar function names or content overlap
+            semantic = []
+            target_words = set(target_content.lower().split())
+
+            for node in data.get("nodes", []):
+                if node.get("type") != "chunk" or node.get("id") == chunk_id:
+                    continue
+                content = node.get("content", "")
+                node_words = set(content.lower().split())
+                # Jaccard similarity
+                if target_words and node_words:
+                    intersection = len(target_words & node_words)
+                    union = len(target_words | node_words)
+                    similarity = intersection / union if union > 0 else 0
+                    if similarity > 0.3:  # 30% threshold
+                        semantic.append(
+                            {
+                                "id": node.get("id"),
+                                "name": node.get("function_name")
+                                or node.get("class_name")
+                                or "chunk",
+                                "file": node.get("file_path", ""),
+                                "similarity": round(similarity, 2),
+                            }
+                        )
+
+            # Sort by similarity and limit
+            semantic.sort(key=lambda x: x["similarity"], reverse=True)
+            semantic = semantic[:10]
+
+            return Response(
+                content=orjson.dumps(
+                    {
+                        "chunk_id": chunk_id,
+                        "callers": callers,
+                        "caller_count": len(callers),
+                        "semantic": semantic,
+                        "semantic_count": len(semantic),
+                    }
+                ),
+                media_type="application/json",
+                headers={"Cache-Control": "max-age=300"},
+            )
+        except Exception as e:
+            console.print(f"[red]Error computing relationships: {e}[/red]")
+            return Response(
+                content='{"error": "Failed to compute relationships"}',
+                status_code=500,
+                media_type="application/json",
+            )
+
+    @app.get("/api/callers/{chunk_id}")
+    async def get_chunk_callers(chunk_id: str) -> Response:
+        """Get callers for a specific code chunk (lazy loaded on-demand).
+
+        This computes callers for a single chunk instantly instead of
+        pre-computing all relationships (which takes 20+ minutes).
+
+        Args:
+            chunk_id: The chunk ID to find callers for
+
+        Returns:
+            JSON response with callers array
+        """
+        graph_file = viz_dir / "chunk-graph.json"
+
+        if not graph_file.exists():
+            return Response(
+                content='{"error": "Graph data not found", "callers": []}',
+                status_code=404,
+                media_type="application/json",
+            )
+
+        try:
+            import ast
+
+            with open(graph_file, "rb") as f:
+                data = orjson.loads(f.read())
+
+            # Find the target chunk
+            target_node = None
+            for node in data.get("nodes", []):
+                if node.get("id") == chunk_id:
+                    target_node = node
+                    break
+
+            if not target_node:
+                return Response(
+                    content='{"error": "Chunk not found", "callers": []}',
+                    status_code=404,
+                    media_type="application/json",
+                )
+
+            # Get the function/class name from the target
+            function_name = target_node.get("function_name") or target_node.get(
+                "class_name"
+            )
+            if not function_name:
+                return Response(
+                    content=orjson.dumps({"callers": [], "function_name": None}),
+                    media_type="application/json",
+                )
+
+            target_file = target_node.get("file_path", "")
+
+            # Find callers by scanning other chunks
+            callers = []
+
+            def extract_calls(code: str) -> set[str]:
+                """Extract function calls from code using AST."""
+                calls = set()
+                try:
+                    tree = ast.parse(code)
+                    for node in ast.walk(tree):
+                        if isinstance(node, ast.Call):
+                            if isinstance(node.func, ast.Name):
+                                calls.add(node.func.id)
+                            elif isinstance(node.func, ast.Attribute):
+                                calls.add(node.func.attr)
+                except SyntaxError:
+                    pass
+                return calls
+
+            for node in data.get("nodes", []):
+                # Skip non-code chunks and same-file chunks
+                if node.get("type") != "chunk":
+                    continue
+                node_file = node.get("file_path", "")
+                if node_file == target_file:
+                    continue
+
+                # Check if this chunk calls our target function
+                content = node.get("content", "")
+                if function_name in extract_calls(content):
+                    caller_name = node.get("function_name") or node.get("class_name")
+                    if caller_name == "__init__":
+                        continue  # Skip noise
+
+                    callers.append(
+                        {
+                            "id": node.get("id"),
+                            "name": caller_name or f"chunk_{node.get('start_line', 0)}",
+                            "file": node_file,
+                            "type": node.get("chunk_type", "code"),
+                        }
+                    )
+
+            return Response(
+                content=orjson.dumps(
+                    {
+                        "callers": callers,
+                        "function_name": function_name,
+                        "count": len(callers),
+                    }
+                ),
+                media_type="application/json",
+                headers={"Cache-Control": "max-age=300"},  # Cache for 5 minutes
+            )
+        except Exception as e:
+            console.print(f"[red]Error computing callers: {e}[/red]")
+            return Response(
+                content='{"error": "Failed to compute callers", "callers": []}',
+                status_code=500,
+                media_type="application/json",
+            )
+
     @app.get("/api/chunks")
     async def get_file_chunks(file_id: str) -> Response:
         """Get code chunks for a specific file.
@@ -138,10 +429,8 @@ def create_app(viz_dir: Path) -> FastAPI:
             )
 
         try:
-            import json
-
-            with open(graph_file) as f:
-                data = json.load(f)
+            with open(graph_file, "rb") as f:
+                data = orjson.loads(f.read())
 
             # Find chunks associated with this file
             # Look for nodes that have this file as parent via containment links
@@ -159,7 +448,7 @@ def create_app(viz_dir: Path) -> FastAPI:
                     )
 
             return Response(
-                content=json.dumps({"chunks": chunks}),
+                content=orjson.dumps({"chunks": chunks}),
                 media_type="application/json",
                 headers={"Cache-Control": "no-cache"},
             )
@@ -277,7 +566,7 @@ def start_visualization_server(
                 f"URL: [cyan]{url}[/cyan]\n"
                 f"Directory: [dim]{viz_dir}[/dim]\n\n"
                 f"[dim]Press Ctrl+C to stop[/dim]",
-                title="Server Started",
+                title=f"Server Started v{__version__}",
                 border_style="green",
             )
         )

mcp_vector_search/cli/commands/visualize/templates/base.py

@@ -6,6 +6,8 @@ to generate the complete HTML page for the D3.js visualization.
 
 import time
 
+from mcp_vector_search import __build__, __version__
+
 from .scripts import get_all_scripts
 from .styles import get_all_styles
 
@@ -43,14 +45,10 @@ def generate_html_template() -> str:
 <body>
     <div id="controls">
         <h1>🔍 Code Tree</h1>
-
-        <div class="search-container">
-            <input type="text" id="search-input" placeholder="Search nodes..." oninput="handleSearchInput(event)" onkeydown="handleSearchKeydown(event)">
-            <div id="search-results" class="search-results"></div>
-        </div>
+        <div class="version-badge">v{__version__} (build {__build__})</div>
 
         <div class="control-group">
-            <label style="color: #c9d1d9; margin-bottom: 8px;">Layout Mode</label>
+            <label style="color: var(--text-primary); margin-bottom: 8px;">Layout Mode</label>
             <div class="toggle-switch-container">
                 <span class="toggle-label">Linear</span>
                 <label class="toggle-switch">
@@ -61,6 +59,15 @@ def generate_html_template() -> str:
             </div>
         </div>
 
+        <div class="control-group">
+            <label style="color: var(--text-primary); margin-bottom: 8px;">Show Files</label>
+            <div class="filter-buttons">
+                <button class="filter-btn active" data-filter="all" onclick="setFileFilter('all')">All</button>
+                <button class="filter-btn" data-filter="code" onclick="setFileFilter('code')">Code</button>
+                <button class="filter-btn" data-filter="docs" onclick="setFileFilter('docs')">Docs</button>
+            </div>
+        </div>
+
         <h3>Legend</h3>
         <div class="legend">
             <div class="legend-category">
@@ -127,6 +134,53 @@ def generate_html_template() -> str:
             </div>
         </div>
 
+        <!-- Search Section -->
+        <h3>🔎 Search</h3>
+        <div class="search-container">
+            <input type="text" id="search-input" placeholder="Search nodes..." oninput="handleSearchInput(event)" onkeydown="handleSearchKeydown(event)">
+            <div id="search-results" class="search-results"></div>
+        </div>
+
+        <!-- Options Section -->
+        <h3>📋 Reports</h3>
+        <div class="legend" style="margin-top: 8px;">
+            <div class="legend-category" style="border-bottom: none;">
+                <div class="legend-item report-btn" onclick="showComplexityReport()">
+                    <span class="report-icon">📊</span>
+                    <span>Complexity</span>
+                </div>
+                <div class="legend-item report-btn" onclick="showCodeSmells()">
+                    <span class="report-icon">🔍</span>
+                    <span>Code Smells</span>
+                </div>
+                <div class="legend-item report-btn" onclick="showDependencies()">
+                    <span class="report-icon">🔗</span>
+                    <span>Dependencies</span>
+                </div>
+                <div class="legend-item report-btn" onclick="showTrends()">
+                    <span class="report-icon">📈</span>
+                    <span>Trends</span>
+                </div>
+                <div class="legend-item report-btn" onclick="generateRemediationReport()">
+                    <span class="report-icon">📋</span>
+                    <span>Remediation</span>
+                </div>
+            </div>
+        </div>
+
+        <h3 style="margin-top: 16px;">Options</h3>
+        <div class="legend" style="margin-top: 8px;">
+            <div class="legend-category" style="border-bottom: none;">
+                <!-- Theme Toggle -->
+                <div class="legend-item" style="margin-bottom: 12px; padding: 0;">
+                    <button class="theme-toggle-icon-btn" onclick="toggleTheme()" title="Toggle dark/light theme">
+                        <span class="theme-icon" id="theme-icon">🌙</span>
+                    </button>
+                    <span style="margin-left: 8px; color: var(--text-secondary); font-size: 12px;">Theme</span>
+                </div>
+            </div>
+        </div>
+
         <div class="stats" id="stats"></div>
     </div>