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

mcp_vector_search/cli/commands/visualize/__init__.py

@@ -0,0 +1,39 @@
+"""Modular visualization package for code graphs.
+
+This package provides D3.js-based interactive visualization of code relationships,
+organized into modular components for maintainability.
+
+Structure:
+    - graph_builder.py: Graph data construction logic
+    - server.py: HTTP server for serving visualization
+    - templates/: HTML, CSS, and JavaScript generation
+    - exporters/: Export functionality (JSON, HTML)
+    - cli.py: Typer CLI commands (imported for backwards compatibility)
+"""
+
+from .cli import app
+from .exporters import export_to_html, export_to_json
+from .graph_builder import (
+    build_graph_data,
+    get_subproject_color,
+    parse_project_dependencies,
+)
+from .server import find_free_port, start_visualization_server
+from .templates import generate_html_template
+
+__all__ = [
+    # CLI
+    "app",
+    # Graph building
+    "build_graph_data",
+    "get_subproject_color",
+    "parse_project_dependencies",
+    # Server
+    "find_free_port",
+    "start_visualization_server",
+    # Templates
+    "generate_html_template",
+    # Exporters
+    "export_to_html",
+    "export_to_json",
+]

mcp_vector_search/cli/commands/visualize/cli.py

@@ -0,0 +1,276 @@
+"""Visualization commands for MCP Vector Search.
+
+This module provides a backwards-compatible interface to the refactored
+modular visualization components.
+"""
+
+import asyncio
+import shutil
+from fnmatch import fnmatch
+from pathlib import Path
+
+import typer
+from loguru import logger
+from rich.console import Console
+from rich.panel import Panel
+
+from ....core.database import ChromaVectorDatabase
+from ....core.embeddings import create_embedding_function
+from ....core.project import ProjectManager
+
+# Import from refactored modules (same directory)
+from .exporters import export_to_html, export_to_json
+from .graph_builder import build_graph_data
+from .server import find_free_port, start_visualization_server
+
+app = typer.Typer(
+    help="Visualize code chunk relationships",
+)
+console = Console()
+
+
+@app.callback(invoke_without_command=True)
+def visualize_callback(ctx: typer.Context) -> None:
+    """Visualize code chunk relationships.
+
+    When called without a subcommand, automatically starts the visualization server.
+    """
+    # If no subcommand was invoked, run serve with defaults
+    if ctx.invoked_subcommand is None:
+        # Call serve directly with default parameters
+        serve(port=8080, graph_file=Path("chunk-graph.json"), code_only=False)
+
+
+@app.command()
+def export(
+    output: Path = typer.Option(
+        Path("chunk-graph.json"),
+        "--output",
+        "-o",
+        help="Output file for chunk relationship data",
+    ),
+    file_path: str | None = typer.Option(
+        None,
+        "--file",
+        "-f",
+        help="Export only chunks from specific file (supports wildcards)",
+    ),
+    code_only: bool = typer.Option(
+        False,
+        "--code-only",
+        help="Exclude documentation chunks (text, comment, docstring)",
+    ),
+) -> None:
+    """Export chunk relationships as JSON for D3.js visualization.
+
+    Examples:
+        # Export all chunks
+        mcp-vector-search visualize export
+
+        # Export from specific file
+        mcp-vector-search visualize export --file src/main.py
+
+        # Custom output location
+        mcp-vector-search visualize export -o graph.json
+
+        # Export only code chunks (exclude documentation)
+        mcp-vector-search visualize export --code-only
+    """
+    asyncio.run(_export_chunks(output, file_path, code_only))
+
+
+async def _export_chunks(
+    output: Path, file_filter: str | None, code_only: bool = False
+) -> None:
+    """Export chunk relationship data.
+
+    Args:
+        output: Path to output JSON file
+        file_filter: Optional file pattern to filter chunks
+        code_only: If True, exclude documentation chunks (text, comment, docstring)
+    """
+    try:
+        # Load project
+        project_manager = ProjectManager(Path.cwd())
+
+        if not project_manager.is_initialized():
+            console.print(
+                "[red]Project not initialized. Run 'mcp-vector-search init' first.[/red]"
+            )
+            raise typer.Exit(1)
+
+        config = project_manager.load_config()
+
+        # Get database
+        embedding_function, _ = create_embedding_function(config.embedding_model)
+        database = ChromaVectorDatabase(
+            persist_directory=config.index_path,
+            embedding_function=embedding_function,
+        )
+        await database.initialize()
+
+        # Get all chunks with metadata
+        console.print("[cyan]Fetching chunks from database...[/cyan]")
+        chunks = await database.get_all_chunks()
+
+        if len(chunks) == 0:
+            console.print(
+                "[yellow]No chunks found in index. Run 'mcp-vector-search index' first.[/yellow]"
+            )
+            raise typer.Exit(1)
+
+        console.print(f"[green]✓[/green] Retrieved {len(chunks)} chunks")
+
+        # Apply file filter if specified
+        if file_filter:
+            chunks = [c for c in chunks if fnmatch(str(c.file_path), file_filter)]
+            console.print(
+                f"[cyan]Filtered to {len(chunks)} chunks matching '{file_filter}'[/cyan]"
+            )
+
+        # Apply code-only filter if requested
+        if code_only:
+            original_count = len(chunks)
+            chunks = [
+                c
+                for c in chunks
+                if c.chunk_type not in ["text", "comment", "docstring"]
+            ]
+            filtered_count = len(chunks)
+            console.print(
+                f"[dim]Filtered out {original_count - filtered_count} documentation chunks "
+                f"({original_count} → {filtered_count} chunks)[/dim]"
+            )
+
+        # Build graph data using refactored module
+        graph_data = await build_graph_data(
+            chunks=chunks,
+            database=database,
+            project_manager=project_manager,
+            code_only=code_only,
+        )
+
+        # Export to JSON using refactored module
+        export_to_json(graph_data, output)
+
+        await database.close()
+
+        console.print()
+        # Count cycles from graph_data links
+        cycles = [link for link in graph_data["links"] if link.get("is_cycle", False)]
+        cycle_warning = f"[yellow]Cycles: {len(cycles)} ⚠️[/yellow]\n" if cycles else ""
+
+        # Count subprojects
+        subprojects_count = len(graph_data["metadata"].get("subprojects", []))
+
+        console.print(
+            Panel.fit(
+                f"[green]✓[/green] Exported graph data to [cyan]{output}[/cyan]\n\n"
+                f"Nodes: {len(graph_data['nodes'])}\n"
+                f"Links: {len(graph_data['links'])}\n"
+                f"{cycle_warning}"
+                f"{'Subprojects: ' + str(subprojects_count) if subprojects_count else ''}\n\n"
+                f"[dim]Next: Run 'mcp-vector-search visualize serve' to view[/dim]",
+                title="Export Complete",
+                border_style="green",
+            )
+        )
+
+    except Exception as e:
+        logger.error(f"Export failed: {e}")
+        console.print(f"[red]✗ Export failed: {e}[/red]")
+        raise typer.Exit(1)
+
+
+@app.command()
+def serve(
+    port: int = typer.Option(
+        8080, "--port", "-p", help="Port for visualization server"
+    ),
+    graph_file: Path = typer.Option(
+        Path("chunk-graph.json"),
+        "--graph",
+        "-g",
+        help="Graph JSON file to visualize",
+    ),
+    code_only: bool = typer.Option(
+        False,
+        "--code-only",
+        help="Exclude documentation chunks (text, comment, docstring)",
+    ),
+) -> None:
+    """Start local HTTP server for D3.js visualization.
+
+    Examples:
+        # Start server on default port 8080
+        mcp-vector-search visualize serve
+
+        # Custom port
+        mcp-vector-search visualize serve --port 3000
+
+        # Custom graph file
+        mcp-vector-search visualize serve --graph my-graph.json
+
+        # Serve with code-only filter
+        mcp-vector-search visualize serve --code-only
+    """
+    # Use specified port or find free one
+    if port == 8080:  # Default port, try to find free one
+        try:
+            port = find_free_port(8080, 8099)
+        except OSError as e:
+            console.print(f"[red]✗ {e}[/red]")
+            raise typer.Exit(1)
+
+    # Get visualization directory - use project-local storage
+    project_manager = ProjectManager(Path.cwd())
+    if not project_manager.is_initialized():
+        console.print(
+            "[red]Project not initialized. Run 'mcp-vector-search init' first.[/red]"
+        )
+        raise typer.Exit(1)
+
+    viz_dir = project_manager.project_root / ".mcp-vector-search" / "visualization"
+
+    if not viz_dir.exists():
+        console.print(
+            f"[yellow]Visualization directory not found. Creating at {viz_dir}...[/yellow]"
+        )
+        viz_dir.mkdir(parents=True, exist_ok=True)
+
+    # Always ensure index.html exists (regenerate if missing)
+    html_file = viz_dir / "index.html"
+    if not html_file.exists():
+        console.print("[yellow]Creating visualization HTML file...[/yellow]")
+        export_to_html(html_file)
+
+    # Check if we need to regenerate the graph file
+    needs_regeneration = not graph_file.exists() or code_only
+
+    if graph_file.exists() and not needs_regeneration:
+        # Use existing unfiltered file
+        dest = viz_dir / "chunk-graph.json"
+        shutil.copy(graph_file, dest)
+        console.print(f"[green]✓[/green] Copied graph data to {dest}")
+    else:
+        # Generate new file (with filter if requested)
+        if graph_file.exists() and code_only:
+            console.print(
+                "[yellow]Regenerating filtered graph data (--code-only)...[/yellow]"
+            )
+        elif not graph_file.exists():
+            console.print(
+                f"[yellow]Graph file {graph_file} not found. Generating it now...[/yellow]"
+            )
+
+        asyncio.run(_export_chunks(graph_file, None, code_only))
+        console.print()
+
+        # Copy the newly generated graph to visualization directory
+        if graph_file.exists():
+            dest = viz_dir / "chunk-graph.json"
+            shutil.copy(graph_file, dest)
+            console.print(f"[green]✓[/green] Copied graph data to {dest}")
+
+    # Start server using refactored module
+    start_visualization_server(port, viz_dir, auto_open=True)

mcp_vector_search/cli/commands/visualize/exporters/__init__.py

@@ -0,0 +1,12 @@
+"""Export functionality for graph visualization data.
+
+This package contains modules for exporting graph data to various formats.
+"""
+
+from .html_exporter import export_to_html
+from .json_exporter import export_to_json
+
+__all__ = [
+    "export_to_html",
+    "export_to_json",
+]

mcp_vector_search/cli/commands/visualize/exporters/html_exporter.py

@@ -0,0 +1,33 @@
+"""HTML export functionality for visualization.
+
+This module handles creating standalone HTML files for the visualization.
+"""
+
+from pathlib import Path
+
+from rich.console import Console
+
+from ..templates.base import generate_html_template
+
+console = Console()
+
+
+def export_to_html(output_path: Path) -> None:
+    """Export visualization to standalone HTML file.
+
+    Args:
+        output_path: Path to output HTML file
+    """
+    # Generate HTML template
+    html_content = generate_html_template()
+
+    # Ensure output directory exists
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Write to file
+    with open(output_path, "w") as f:
+        f.write(html_content)
+
+    console.print(
+        f"[green]✓[/green] Created visualization HTML at [cyan]{output_path}[/cyan]"
+    )

mcp_vector_search/cli/commands/visualize/exporters/json_exporter.py

@@ -0,0 +1,29 @@
+"""JSON export functionality for graph data.
+
+This module handles exporting graph data to JSON format.
+"""
+
+import json
+from pathlib import Path
+from typing import Any
+
+from rich.console import Console
+
+console = Console()
+
+
+def export_to_json(graph_data: dict[str, Any], output_path: Path) -> None:
+    """Export graph data to JSON file.
+
+    Args:
+        graph_data: Graph data dictionary containing nodes, links, and metadata
+        output_path: Path to output JSON file
+    """
+    # Ensure output directory exists
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    # Write to file
+    with open(output_path, "w") as f:
+        json.dump(graph_data, f, indent=2)
+
+    console.print(f"[green]✓[/green] Exported graph data to [cyan]{output_path}[/cyan]")