hanzo-mcp 0.8.8__py3-none-any.whl → 0.9.0__py3-none-any.whl

hanzo_mcp/tools/vector/node_tool.py

@@ -0,0 +1,538 @@
+"""Node management tool for hanzo-node operations."""
+
+import asyncio
+import json
+import logging
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, List, Optional, final, override
+
+import httpx
+from mcp.server import FastMCP
+from mcp.server.fastmcp import Context as MCPContext
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.context import create_tool_context
+
+logger = logging.getLogger(__name__)
+
+# Default hanzo-node configuration
+DEFAULT_NODE_CONFIG = {
+    "host": "localhost",
+    "port": 3690,
+    "db_path": "~/.hanzo/lancedb",
+    "models_path": "~/.hanzo/models",
+    "log_level": "info",
+    "embedding_model": "text-embedding-3-small",
+    "embedding_dimensions": 1536
+}
+
+
+@final
+class NodeTool(BaseTool):
+    """Tool for managing local hanzo-node instance.
+    
+    This tool provides management capabilities for the local hanzo-node:
+    - Download and install hanzo-node
+    - Configure node settings via ~/.hanzo
+    - Start/stop/restart node
+    - Check node status and health
+    - Manage vector store (LanceDB)
+    - Download and load models
+    - View logs and diagnostics
+    """
+    
+    def __init__(self):
+        """Initialize node management tool."""
+        self.hanzo_dir = Path.home() / ".hanzo"
+        self.config_file = self.hanzo_dir / "node_config.json"
+        self.models_dir = self.hanzo_dir / "models"
+        self.logs_dir = self.hanzo_dir / "logs"
+        self.lancedb_dir = self.hanzo_dir / "lancedb"
+        
+        # Ensure directories exist
+        for directory in [self.hanzo_dir, self.models_dir, self.logs_dir, self.lancedb_dir]:
+            directory.mkdir(parents=True, exist_ok=True)
+            
+        # Initialize config if it doesn't exist
+        self._ensure_config()
+        
+    def _ensure_config(self) -> None:
+        """Ensure node configuration exists."""
+        if not self.config_file.exists():
+            config = DEFAULT_NODE_CONFIG.copy()
+            # Expand paths
+            config["db_path"] = str(self.lancedb_dir)
+            config["models_path"] = str(self.models_dir)
+            self._save_config(config)
+            
+    def _load_config(self) -> Dict[str, Any]:
+        """Load node configuration."""
+        try:
+            with open(self.config_file, 'r') as f:
+                return json.load(f)
+        except Exception:
+            return DEFAULT_NODE_CONFIG.copy()
+            
+    def _save_config(self, config: Dict[str, Any]) -> None:
+        """Save node configuration."""
+        with open(self.config_file, 'w') as f:
+            json.dump(config, f, indent=2)
+            
+    async def _is_node_running(self, host: str = "localhost", port: int = 3690) -> bool:
+        """Check if hanzo-node is running."""
+        try:
+            async with httpx.AsyncClient(timeout=5.0) as client:
+                response = await client.get(f"http://{host}:{port}/health")
+                return response.status_code == 200
+        except Exception:
+            return False
+            
+    async def _find_node_executable(self) -> Optional[Path]:
+        """Find hanzo-node executable."""
+        # Check common locations
+        possible_paths = [
+            Path.home() / ".hanzo" / "bin" / "hanzo-node",
+            Path("/usr/local/bin/hanzo-node"),
+            Path("/opt/hanzo/bin/hanzo-node"),
+            # Development paths
+            Path.home() / "work" / "hanzo" / "node" / "apps" / "hanzo-node" / "target" / "release" / "hanzo-node",
+            Path.home() / "work" / "hanzo" / "target" / "release" / "hanzo-node",
+        ]
+        
+        # Also check PATH
+        if shutil.which("hanzo-node"):
+            possible_paths.append(Path(shutil.which("hanzo-node")))
+            
+        for path in possible_paths:
+            if path.exists() and path.is_file():
+                return path
+                
+        return None
+        
+    async def _download_node(self, tool_ctx) -> bool:
+        """Download hanzo-node if not available."""
+        await tool_ctx.info("Downloading hanzo-node...")
+        
+        # For now, we'll build from source if available
+        source_dir = Path.home() / "work" / "hanzo" / "node" / "apps" / "hanzo-node"
+        if source_dir.exists():
+            await tool_ctx.info("Building hanzo-node from source...")
+            try:
+                # Build the project
+                result = subprocess.run(
+                    ["cargo", "build", "--release"],
+                    cwd=source_dir,
+                    capture_output=True,
+                    text=True,
+                    timeout=300  # 5 minutes
+                )
+                
+                if result.returncode == 0:
+                    # Copy binary to hanzo directory
+                    source_binary = source_dir / "target" / "release" / "hanzo-node"
+                    target_binary = self.hanzo_dir / "bin" / "hanzo-node"
+                    target_binary.parent.mkdir(parents=True, exist_ok=True)
+                    
+                    if source_binary.exists():
+                        shutil.copy2(source_binary, target_binary)
+                        target_binary.chmod(0o755)
+                        await tool_ctx.info(f"Built and installed hanzo-node to {target_binary}")
+                        return True
+                    else:
+                        await tool_ctx.error("Build succeeded but binary not found")
+                        return False
+                else:
+                    await tool_ctx.error(f"Build failed: {result.stderr}")
+                    return False
+                    
+            except subprocess.TimeoutExpired:
+                await tool_ctx.error("Build timed out after 5 minutes")
+                return False
+            except Exception as e:
+                await tool_ctx.error(f"Build error: {e}")
+                return False
+        else:
+            await tool_ctx.warning("Source code not found. Manual installation required.")
+            await tool_ctx.info("To install hanzo-node manually:")
+            await tool_ctx.info("1. Download from releases or build from source")
+            await tool_ctx.info("2. Place binary at ~/.hanzo/bin/hanzo-node")
+            await tool_ctx.info("3. Run 'node(action=\"status\")' to verify")
+            return False
+
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name."""
+        return "node"
+        
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description."""
+        return """Manage local hanzo-node instance.
+
+This tool provides comprehensive management of your local hanzo-node:
+
+Configuration:
+- Configure node settings via ~/.hanzo/node_config.json
+- Set embedding models, database paths, ports
+- Manage authentication and security settings
+
+Lifecycle:
+- Install/download hanzo-node if not available
+- Start, stop, restart the node process
+- Check node status and health
+- View logs and diagnostics
+
+Data Management:
+- Manage vector store (LanceDB) location and settings
+- Download and configure ML models
+- Backup and restore node data
+- Clear caches and temporary files
+
+Examples:
+node(action="status")  # Check if node is running
+node(action="start")   # Start the node
+node(action="stop")    # Stop the node
+node(action="restart") # Restart the node
+node(action="config", key="port", value="3691")  # Update config
+node(action="logs")    # View recent logs
+node(action="models")  # List available models
+node(action="install") # Install/build hanzo-node
+"""
+
+    @override
+    async def call(
+        self,
+        ctx: MCPContext,
+        action: str,
+        key: Optional[str] = None,
+        value: Optional[str] = None,
+        force: bool = False,
+        **kwargs
+    ) -> str:
+        """Execute node management operations.
+        
+        Args:
+            ctx: MCP context
+            action: Action to perform (status, start, stop, restart, config, logs, models, install)
+            key: Configuration key (for config action)
+            value: Configuration value (for config action)
+            force: Force action without confirmation
+            **kwargs: Additional action-specific parameters
+            
+        Returns:
+            Operation result or status information
+        """
+        tool_ctx = create_tool_context(ctx)
+        await tool_ctx.set_tool_info(self.name)
+        
+        if action == "status":
+            return await self._handle_status(tool_ctx)
+        elif action == "start":
+            return await self._handle_start(tool_ctx, force)
+        elif action == "stop":
+            return await self._handle_stop(tool_ctx, force)
+        elif action == "restart":
+            return await self._handle_restart(tool_ctx, force)
+        elif action == "config":
+            return await self._handle_config(tool_ctx, key, value)
+        elif action == "logs":
+            return await self._handle_logs(tool_ctx, kwargs.get("lines", 50))
+        elif action == "models":
+            return await self._handle_models(tool_ctx)
+        elif action == "install":
+            return await self._handle_install(tool_ctx, force)
+        elif action == "clean":
+            return await self._handle_clean(tool_ctx, force)
+        else:
+            return f"Unknown action: {action}. Available actions: status, start, stop, restart, config, logs, models, install, clean"
+            
+    async def _handle_status(self, tool_ctx) -> str:
+        """Handle status check."""
+        config = self._load_config()
+        host = config.get("host", "localhost")
+        port = config.get("port", 3690)
+        
+        status_info = [f"Hanzo Node Status"]
+        status_info.append(f"Configuration: {self.config_file}")
+        status_info.append(f"Expected URL: http://{host}:{port}")
+        
+        # Check if node is running
+        is_running = await self._is_node_running(host, port)
+        status_info.append(f"Running: {is_running}")
+        
+        # Check for executable
+        executable = await self._find_node_executable()
+        if executable:
+            status_info.append(f"Executable: {executable}")
+        else:
+            status_info.append("Executable: Not found")
+            
+        # Check directories
+        status_info.append(f"Data directory: {self.hanzo_dir}")
+        status_info.append(f"Models directory: {self.models_dir} ({len(list(self.models_dir.glob('*')))} files)")
+        status_info.append(f"LanceDB directory: {self.lancedb_dir} ({len(list(self.lancedb_dir.glob('*')))} files)")
+        
+        # Check config
+        status_info.append(f"Configured port: {port}")
+        status_info.append(f"Configured host: {host}")
+        status_info.append(f"DB path: {config.get('db_path', 'Not set')}")
+        status_info.append(f"Models path: {config.get('models_path', 'Not set')}")
+        
+        return "\n".join(status_info)
+        
+    async def _handle_start(self, tool_ctx, force: bool) -> str:
+        """Handle node start."""
+        config = self._load_config()
+        host = config.get("host", "localhost")
+        port = config.get("port", 3690)
+        
+        # Check if already running
+        if await self._is_node_running(host, port):
+            return f"Hanzo node is already running on {host}:{port}"
+            
+        # Find executable
+        executable = await self._find_node_executable()
+        if not executable:
+            return "Hanzo node executable not found. Run node(action=\"install\") first."
+            
+        await tool_ctx.info(f"Starting hanzo-node on {host}:{port}...")
+        
+        try:
+            # Prepare environment and arguments
+            env = os.environ.copy()
+            env["HANZO_CONFIG"] = str(self.config_file)
+            
+            # Start the process in background
+            process = subprocess.Popen(
+                [str(executable), "--config", str(self.config_file)],
+                env=env,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                start_new_session=True
+            )
+            
+            # Wait a moment and check if it started successfully
+            await asyncio.sleep(2)
+            
+            if await self._is_node_running(host, port):
+                return f"Successfully started hanzo-node on {host}:{port} (PID: {process.pid})"
+            else:
+                # Process might have failed
+                return_code = process.poll()
+                if return_code is not None:
+                    _, stderr = process.communicate()
+                    return f"Failed to start hanzo-node (exit code {return_code}): {stderr.decode()}"
+                else:
+                    return f"Started hanzo-node process (PID: {process.pid}) but health check failed"
+                    
+        except Exception as e:
+            return f"Error starting hanzo-node: {e}"
+            
+    async def _handle_stop(self, tool_ctx, force: bool) -> str:
+        """Handle node stop."""
+        config = self._load_config()
+        host = config.get("host", "localhost")
+        port = config.get("port", 3690)
+        
+        # Check if running
+        if not await self._is_node_running(host, port):
+            return "Hanzo node is not running"
+            
+        await tool_ctx.info("Stopping hanzo-node...")
+        
+        try:
+            # Try graceful shutdown via API first
+            async with httpx.AsyncClient(timeout=10.0) as client:
+                try:
+                    response = await client.post(f"http://{host}:{port}/api/shutdown")
+                    if response.status_code == 200:
+                        # Wait for shutdown
+                        await asyncio.sleep(2)
+                        if not await self._is_node_running(host, port):
+                            return "Successfully stopped hanzo-node"
+                except Exception:
+                    pass
+                    
+            # If graceful shutdown failed, try finding and killing the process
+            if force:
+                try:
+                    # Find processes by name
+                    result = subprocess.run(
+                        ["pgrep", "-f", "hanzo-node"],
+                        capture_output=True,
+                        text=True
+                    )
+                    
+                    if result.returncode == 0 and result.stdout.strip():
+                        pids = result.stdout.strip().split('\n')
+                        for pid in pids:
+                            subprocess.run(["kill", "-TERM", pid])
+                            
+                        await asyncio.sleep(2)
+                        if not await self._is_node_running(host, port):
+                            return f"Force stopped hanzo-node (killed {len(pids)} processes)"
+                        else:
+                            # Try SIGKILL
+                            for pid in pids:
+                                subprocess.run(["kill", "-KILL", pid])
+                            return f"Force killed hanzo-node processes"
+                except Exception as e:
+                    return f"Error force stopping: {e}"
+            else:
+                return "Graceful shutdown failed. Use force=true for forceful shutdown."
+                
+        except Exception as e:
+            return f"Error stopping hanzo-node: {e}"
+            
+    async def _handle_restart(self, tool_ctx, force: bool) -> str:
+        """Handle node restart."""
+        stop_result = await self._handle_stop(tool_ctx, force)
+        await asyncio.sleep(1)
+        start_result = await self._handle_start(tool_ctx, force)
+        return f"Restart: {stop_result} -> {start_result}"
+        
+    async def _handle_config(self, tool_ctx, key: Optional[str], value: Optional[str]) -> str:
+        """Handle configuration management."""
+        config = self._load_config()
+        
+        if key is None:
+            # Show current config
+            formatted_config = json.dumps(config, indent=2)
+            return f"Current configuration:\n{formatted_config}"
+            
+        if value is None:
+            # Show specific key
+            if key in config:
+                return f"{key}: {config[key]}"
+            else:
+                return f"Configuration key '{key}' not found"
+                
+        # Update configuration
+        old_value = config.get(key, "Not set")
+        config[key] = value
+        
+        # Validate important settings
+        if key == "port":
+            try:
+                port = int(value)
+                if not (1024 <= port <= 65535):
+                    return "Error: Port must be between 1024 and 65535"
+            except ValueError:
+                return "Error: Port must be a number"
+                
+        # Save updated config
+        self._save_config(config)
+        
+        await tool_ctx.info(f"Updated {key}: {old_value} -> {value}")
+        return f"Configuration updated: {key} = {value}\nRestart required for changes to take effect."
+        
+    async def _handle_logs(self, tool_ctx, lines: int = 50) -> str:
+        """Handle log viewing."""
+        log_file = self.logs_dir / "hanzo-node.log"
+        
+        if not log_file.exists():
+            return "No log file found. Node may not have been started yet."
+            
+        try:
+            # Read last N lines
+            with open(log_file, 'r') as f:
+                all_lines = f.readlines()
+                recent_lines = all_lines[-lines:] if len(all_lines) > lines else all_lines
+                
+            if not recent_lines:
+                return "Log file is empty"
+                
+            log_content = ''.join(recent_lines)
+            return f"Last {len(recent_lines)} lines from hanzo-node.log:\n{log_content}"
+            
+        except Exception as e:
+            return f"Error reading log file: {e}"
+            
+    async def _handle_models(self, tool_ctx) -> str:
+        """Handle model management."""
+        models_info = [f"Models directory: {self.models_dir}"]
+        
+        # List model files
+        model_files = list(self.models_dir.glob("*"))
+        if model_files:
+            models_info.append(f"\nFound {len(model_files)} model files:")
+            for model_file in sorted(model_files):
+                size_mb = model_file.stat().st_size / (1024 * 1024)
+                models_info.append(f"  {model_file.name} ({size_mb:.1f} MB)")
+        else:
+            models_info.append("\nNo model files found")
+            
+        # Show config
+        config = self._load_config()
+        embedding_model = config.get("embedding_model", "Not set")
+        models_info.append(f"\nConfigured embedding model: {embedding_model}")
+        
+        return "\n".join(models_info)
+        
+    async def _handle_install(self, tool_ctx, force: bool) -> str:
+        """Handle node installation."""
+        # Check if already exists
+        existing = await self._find_node_executable()
+        if existing and not force:
+            return f"Hanzo node already installed at {existing}. Use force=true to reinstall."
+            
+        # Attempt download/build
+        success = await self._download_node(tool_ctx)
+        if success:
+            return "Successfully installed hanzo-node"
+        else:
+            return "Failed to install hanzo-node. See logs above for details."
+            
+    async def _handle_clean(self, tool_ctx, force: bool) -> str:
+        """Handle cleanup operations."""
+        if not force:
+            return "Clean operation requires force=true. This will remove logs and temporary files."
+            
+        await tool_ctx.info("Cleaning hanzo-node data...")
+        
+        cleaned = []
+        
+        # Clean logs
+        if self.logs_dir.exists():
+            for log_file in self.logs_dir.glob("*.log"):
+                log_file.unlink()
+                cleaned.append(f"Removed log: {log_file.name}")
+                
+        # Clean temporary files
+        temp_patterns = ["*.tmp", "*.lock", "*.pid"]
+        for pattern in temp_patterns:
+            for temp_file in self.hanzo_dir.glob(pattern):
+                temp_file.unlink()
+                cleaned.append(f"Removed temp file: {temp_file.name}")
+                
+        if cleaned:
+            return f"Cleaned {len(cleaned)} files:\n" + "\n".join(cleaned)
+        else:
+            return "No files to clean"
+
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register this tool with the MCP server."""
+        tool_self = self
+        
+        @mcp_server.tool(name=self.name, description=self.description)
+        async def node(
+            ctx: MCPContext,
+            action: str,
+            key: Optional[str] = None,
+            value: Optional[str] = None,
+            force: bool = False,
+        ) -> str:
+            return await tool_self.call(
+                ctx,
+                action=action,
+                key=key,
+                value=value,
+                force=force,
+            )

hanzo_mcp/tools/vector/project_manager.py

@@ -153,9 +153,7 @@ class ProjectVectorManager:
 
         return None
 
-    def get_vector_store(
-        self, project_info: Optional[ProjectInfo] = None
-    ) -> InfinityVectorStore:
+    def get_vector_store(self, project_info: Optional[ProjectInfo] = None) -> InfinityVectorStore:
         """Get vector store for a project or global store.
 
         Args:
@@ -178,9 +176,7 @@ class ProjectVectorManager:
 
         if project_key not in self.vector_stores:
             # Get index path based on configuration
-            index_path = self.index_config.get_index_path(
-                "vector", str(project_info.root_path)
-            )
+            index_path = self.index_config.get_index_path("vector", str(project_info.root_path))
             index_path.mkdir(parents=True, exist_ok=True)
 
             self.vector_stores[project_key] = InfinityVectorStore(
@@ -270,9 +266,7 @@ class ProjectVectorManager:
             search_tasks.append(
                 asyncio.get_event_loop().run_in_executor(
                     self.executor,
-                    lambda: global_store.search(
-                        query, limit_per_project, score_threshold
-                    ),
+                    lambda: global_store.search(query, limit_per_project, score_threshold),
                 )
             )
             project_names.append("global")
@@ -287,9 +281,7 @@ class ProjectVectorManager:
             search_tasks.append(
                 asyncio.get_event_loop().run_in_executor(
                     self.executor,
-                    lambda vs=vector_store: vs.search(
-                        query, limit_per_project, score_threshold
-                    ),
+                    lambda vs=vector_store: vs.search(query, limit_per_project, score_threshold),
                 )
             )
             project_names.append(project_info.name)