hanzo-mcp 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

hanzo_mcp/tools/agent/lmstudio_agent.py

@@ -0,0 +1,385 @@
+"""LM Studio agent tool for parallel model execution.
+
+This module provides a tool for running tasks in parallel across multiple LM Studio models.
+"""
+
+import asyncio
+import json
+import logging
+import time
+from typing import Any, Dict, List, Optional, final, override
+
+from mcp.server.fastmcp import Context as MCPContext
+from mcp.server.fastmcp import FastMCP
+
+from hanzo_mcp.tools.common.base import BaseTool
+from hanzo_mcp.tools.common.context import DocumentContext, create_tool_context
+from hanzo_mcp.tools.common.permissions import PermissionManager
+from hanzo_mcp.tools.agent.lmstudio_provider import LMStudioProvider
+
+logger = logging.getLogger(__name__)
+
+
+@final
+class LMStudioAgentTool(BaseTool):
+    """Tool for parallel execution of tasks across multiple LM Studio models."""
+
+    @property
+    @override
+    def name(self) -> str:
+        """Get the tool name.
+        
+        Returns:
+            Tool name
+        """
+        return "lmstudio_dispatch"
+        
+    @property
+    @override
+    def description(self) -> str:
+        """Get the tool description.
+        
+        Returns:
+            Tool description
+        """
+        return """Run tasks in parallel across multiple LM Studio models.
+
+This tool allows you to dispatch the same task or different tasks to multiple locally available
+LM Studio models and execute them in parallel. This is useful for comparing model responses,
+leveraging different model strengths, or simply speeding up processing by distributing tasks.
+
+The task prompts can be the same for all models or different per model.
+
+Args:
+    model_tasks: A list of configurations, each with a 'model' name, and a 'prompt'.
+                Optionally can include 'system_prompt', 'temperature', 'max_tokens' and 'top_p'
+
+Returns:
+    Combined results from all model executions with performance metrics
+"""
+
+    @property
+    @override
+    def parameters(self) -> dict[str, Any]:
+        """Get the parameter specifications for the tool.
+        
+        Returns:
+            Parameter specifications
+        """
+        return {
+            "properties": {
+                "model_tasks": {
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "model": {
+                                "type": "string",
+                                "description": "Name of the LM Studio model to use"
+                            },
+                            "identifier": {
+                                "type": "string",
+                                "description": "Optional identifier for the model instance"
+                            },
+                            "prompt": {
+                                "type": "string",
+                                "description": "Task prompt for the model"
+                            },
+                            "system_prompt": {
+                                "type": "string",
+                                "description": "Optional system prompt for the model"
+                            },
+                            "temperature": {
+                                "type": "number",
+                                "description": "Sampling temperature (defaults to 0.7)"
+                            },
+                            "max_tokens": {
+                                "type": "integer",
+                                "description": "Maximum tokens to generate (defaults to 2048)"
+                            },
+                            "top_p": {
+                                "type": "number",
+                                "description": "Top-p sampling parameter (defaults to 0.95)"
+                            }
+                        },
+                        "required": ["model", "prompt"]
+                    },
+                    "description": "List of model task configurations to execute in parallel"
+                }
+            },
+            "required": ["model_tasks"],
+            "type": "object"
+        }
+        
+    @property
+    @override
+    def required(self) -> list[str]:
+        """Get the list of required parameter names.
+        
+        Returns:
+            List of required parameter names
+        """
+        return ["model_tasks"]
+        
+    def __init__(self, document_context: DocumentContext, permission_manager: PermissionManager) -> None:
+        """Initialize the LM Studio agent tool.
+        
+        Args:
+            document_context: Document context for tracking file contents
+            permission_manager: Permission manager for access control
+        """
+        self.document_context = document_context
+        self.permission_manager = permission_manager
+        self.provider = LMStudioProvider()
+        
+    @override
+    async def call(self, ctx: MCPContext, **params: Any) -> str:
+        """Execute the tool with the given parameters.
+        
+        Args:
+            ctx: MCP context
+            **params: Tool parameters
+            
+        Returns:
+            Tool execution result
+        """
+        start_time = time.time()
+        
+        # Create tool context
+        tool_ctx = create_tool_context(ctx)
+        tool_ctx.set_tool_info(self.name)
+        
+        # Extract parameters
+        model_tasks = params.get("model_tasks")
+        if not model_tasks:
+            await tool_ctx.error("Parameter 'model_tasks' is required but was not provided")
+            return "Error: Parameter 'model_tasks' is required but was not provided"
+            
+        if not isinstance(model_tasks, list):
+            await tool_ctx.error("Parameter 'model_tasks' must be an array")
+            return "Error: Parameter 'model_tasks' must be an array"
+            
+        if not model_tasks:
+            await tool_ctx.error("At least one model task must be provided")
+            return "Error: At least one model task must be provided"
+            
+        # Validate each model task
+        for i, task in enumerate(model_tasks):
+            if not isinstance(task, dict):
+                await tool_ctx.error(f"Model task at index {i} must be an object")
+                return f"Error: Model task at index {i} must be an object"
+                
+            if "model" not in task:
+                await tool_ctx.error(f"Model task at index {i} must have a 'model' property")
+                return f"Error: Model task at index {i} must have a 'model' property"
+                
+            if "prompt" not in task:
+                await tool_ctx.error(f"Model task at index {i} must have a 'prompt' property")
+                return f"Error: Model task at index {i} must have a 'prompt' property"
+        
+        # Initialize the provider if needed
+        await self.provider.initialize()
+        
+        # Execute the tasks in parallel
+        await tool_ctx.info(f"Executing {len(model_tasks)} tasks across LM Studio models")
+        result = await self._execute_parallel_tasks(model_tasks, tool_ctx)
+        
+        # Calculate execution time
+        execution_time = time.time() - start_time
+        
+        # Format the result
+        formatted_result = self._format_result(result, execution_time)
+        
+        # Log completion
+        await tool_ctx.info(f"LM Studio model execution completed in {execution_time:.2f}s")
+        
+        return formatted_result
+    
+    async def _execute_parallel_tasks(self, model_tasks: List[Dict[str, Any]], tool_ctx: Any) -> List[Dict[str, Any]]:
+        """Execute multiple model tasks in parallel.
+        
+        Args:
+            model_tasks: List of model task configurations
+            tool_ctx: Tool context for logging
+            
+        Returns:
+            List of task results
+        """
+        # Create tasks for loading models
+        load_tasks = []
+        
+        for task in model_tasks:
+            model_name = task["model"]
+            identifier = task.get("identifier")
+            
+            await tool_ctx.info(f"Loading model: {model_name}" + (f" as {identifier}" if identifier else ""))
+            load_tasks.append(self.provider.load_model(model_name, identifier))
+            
+        # Wait for all models to load
+        try:
+            model_ids = await asyncio.gather(*load_tasks)
+        except Exception as e:
+            await tool_ctx.error(f"Failed to load models: {str(e)}")
+            return [{"error": f"Failed to load models: {str(e)}"}]
+            
+        # Create tasks for generating responses
+        generation_tasks = []
+        
+        for i, (task, model_id) in enumerate(zip(model_tasks, model_ids)):
+            prompt = task["prompt"]
+            system_prompt = task.get("system_prompt")
+            max_tokens = task.get("max_tokens", 2048)
+            temperature = task.get("temperature", 0.7)
+            top_p = task.get("top_p", 0.95)
+            
+            await tool_ctx.info(f"Generating with model {model_id}")
+            generation_tasks.append(
+                self._execute_single_task(
+                    model_id=model_id,
+                    prompt=prompt,
+                    system_prompt=system_prompt,
+                    max_tokens=max_tokens,
+                    temperature=temperature,
+                    top_p=top_p,
+                    task_index=i,
+                    tool_ctx=tool_ctx,
+                    original_task=task
+                )
+            )
+            
+        # Wait for all generation tasks to complete
+        results = await asyncio.gather(*generation_tasks, return_exceptions=True)
+        
+        # Process results, handling any exceptions
+        processed_results = []
+        for i, result in enumerate(results):
+            if isinstance(result, Exception):
+                processed_results.append({
+                    "model": model_tasks[i]["model"],
+                    "error": str(result),
+                    "success": False
+                })
+            else:
+                processed_results.append(result)
+                
+        return processed_results
+    
+    async def _execute_single_task(
+        self,
+        model_id: str,
+        prompt: str,
+        system_prompt: Optional[str],
+        max_tokens: int,
+        temperature: float,
+        top_p: float,
+        task_index: int,
+        tool_ctx: Any,
+        original_task: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """Execute a single model task.
+        
+        Args:
+            model_id: Model identifier
+            prompt: Prompt for the model
+            system_prompt: Optional system prompt
+            max_tokens: Maximum tokens to generate
+            temperature: Sampling temperature
+            top_p: Top-p sampling parameter
+            task_index: Index of the task
+            tool_ctx: Tool context for logging
+            original_task: Original task configuration
+            
+        Returns:
+            Task result
+        """
+        task_start_time = time.time()
+        
+        try:
+            # Generate response
+            generated_text, metadata = await self.provider.generate(
+                model_id=model_id,
+                prompt=prompt,
+                system_prompt=system_prompt,
+                max_tokens=max_tokens,
+                temperature=temperature,
+                top_p=top_p
+            )
+            
+            # Calculate execution time
+            task_execution_time = time.time() - task_start_time
+            
+            await tool_ctx.info(f"Task {task_index} completed in {task_execution_time:.2f}s")
+            
+            # Return result
+            return {
+                "model": original_task["model"],
+                "identifier": model_id,
+                "result": generated_text,
+                "execution_time": task_execution_time,
+                "success": True,
+                "metadata": metadata
+            }
+        except Exception as e:
+            await tool_ctx.error(f"Error executing task {task_index}: {str(e)}")
+            return {
+                "model": original_task["model"],
+                "identifier": model_id,
+                "error": str(e),
+                "execution_time": time.time() - task_start_time,
+                "success": False
+            }
+    
+    def _format_result(self, results: List[Dict[str, Any]], total_execution_time: float) -> str:
+        """Format the task results.
+        
+        Args:
+            results: List of task results
+            total_execution_time: Total execution time
+            
+        Returns:
+            Formatted results
+        """
+        # Calculate summary statistics
+        successful = [r for r in results if r.get("success", False)]
+        failed = [r for r in results if not r.get("success", False)]
+        
+        # Create the result string
+        output = [f"### LM Studio Dispatch Results\n"]
+        output.append(f"**Total execution time:** {total_execution_time:.2f}s")
+        output.append(f"**Models used:** {len(results)}")
+        output.append(f"**Successful:** {len(successful)}")
+        output.append(f"**Failed:** {len(failed)}\n")
+        
+        # Add the results for each model
+        for i, result in enumerate(results):
+            model_name = result.get("model", "Unknown model")
+            model_id = result.get("identifier", model_name)
+            
+            output.append(f"## Model {i+1}: {model_name}")
+            
+            if result.get("success", False):
+                exec_time = result.get("execution_time", 0)
+                output.append(f"**Execution time:** {exec_time:.2f}s")
+                
+                # Add the result
+                output.append("\n**Result:**\n")
+                output.append(result.get("result", "No result"))
+            else:
+                output.append(f"**Error:** {result.get('error', 'Unknown error')}")
+                
+            output.append("\n" + "-" * 40 + "\n")
+            
+        return "\n".join(output)
+    
+    @override
+    def register(self, mcp_server: FastMCP) -> None:
+        """Register this tool with the MCP server.
+        
+        Args:
+            mcp_server: The FastMCP server instance
+        """
+        tool_self = self  # Create a reference to self for use in the closure
+        
+        @mcp_server.tool(name=self.name, description=self.mcp_description)
+        async def lmstudio_dispatch(ctx: MCPContext, model_tasks: List[Dict[str, Any]]) -> str:
+            return await tool_self.call(ctx, model_tasks=model_tasks)

hanzo_mcp/tools/agent/lmstudio_provider.py

@@ -0,0 +1,219 @@
+"""LM Studio provider for agent delegation.
+
+Enables the use of local LLMs via LM Studio's Python SDK.
+"""
+
+import logging
+import asyncio
+import functools
+from typing import Any, Dict, List, Optional, Tuple
+
+from hanzo_mcp.tools.agent.base_provider import BaseModelProvider
+
+logger = logging.getLogger(__name__)
+
+
+class LMStudioProvider(BaseModelProvider):
+    """Provider for local models via LM Studio Python SDK."""
+
+    def __init__(self):
+        """Initialize the LM Studio provider."""
+        self.models = {}
+        self.initialized = False
+        
+    async def initialize(self) -> None:
+        """Initialize the LM Studio provider.
+        
+        Import is done here to avoid dependency issues if LM Studio SDK is not installed.
+        """
+        if self.initialized:
+            return
+            
+        try:
+            # Dynamic import to avoid dependency issues if LM Studio is not installed
+            from importlib.util import find_spec
+            if find_spec("lmstudio") is None:
+                logger.warning("LM Studio Python SDK not installed. Install with 'pip install lmstudio'")
+                return
+                
+            # Import the LM Studio module
+            import lmstudio as lms
+            self.lms = lms
+            self.initialized = True
+            logger.info("LM Studio provider initialized successfully")
+        except ImportError as e:
+            logger.error(f"Failed to import LM Studio Python SDK: {str(e)}")
+            logger.error("Install LM Studio Python SDK with 'pip install lmstudio'")
+        except Exception as e:
+            logger.error(f"Failed to initialize LM Studio provider: {str(e)}")
+            
+    async def load_model(self, model_name: str, identifier: Optional[str] = None) -> str:
+        """Load a model from LM Studio.
+        
+        Args:
+            model_name: The name of the model to load
+            identifier: Optional identifier for the model instance
+            
+        Returns:
+            The identifier for the loaded model
+        """
+        if not self.initialized:
+            await self.initialize()
+            if not self.initialized:
+                raise RuntimeError("LM Studio provider failed to initialize")
+        
+        model_id = identifier or model_name
+        
+        try:
+            if model_id in self.models:
+                logger.info(f"Model {model_id} already loaded")
+                return model_id
+                
+            logger.info(f"Loading model {model_name}")
+            
+            # Use the thread pool to run the blocking operation
+            loop = asyncio.get_event_loop()
+            model = await loop.run_in_executor(
+                None, 
+                functools.partial(self.lms.llm, model_name)
+            )
+            
+            # Store the model with its identifier
+            self.models[model_id] = model
+            logger.info(f"Model {model_name} loaded successfully as {model_id}")
+            return model_id
+            
+        except Exception as e:
+            logger.error(f"Failed to load model {model_name}: {str(e)}")
+            raise RuntimeError(f"Failed to load model {model_name}: {str(e)}")
+    
+    async def generate(
+        self,
+        model_id: str,
+        prompt: str,
+        system_prompt: Optional[str] = None,
+        max_tokens: int = 4096,
+        temperature: float = 0.7,
+        top_p: float = 0.95,
+        stop_sequences: Optional[List[str]] = None,
+    ) -> Tuple[str, Dict[str, Any]]:
+        """Generate a response from the model.
+        
+        Args:
+            model_id: The identifier of the model to use
+            prompt: The prompt to send to the model
+            system_prompt: Optional system prompt to send to the model
+            max_tokens: Maximum number of tokens to generate
+            temperature: Sampling temperature
+            top_p: Top-p sampling parameter
+            stop_sequences: Optional list of strings that will stop generation
+            
+        Returns:
+            A tuple of (generated text, metadata)
+        """
+        if not self.initialized:
+            await self.initialize()
+            if not self.initialized:
+                raise RuntimeError("LM Studio provider failed to initialize")
+        
+        if model_id not in self.models:
+            raise ValueError(f"Model {model_id} not loaded")
+            
+        model = self.models[model_id]
+        
+        try:
+            logger.debug(f"Generating with model {model_id}")
+            
+            # Prepare generation parameters
+            params = {
+                "max_tokens": max_tokens,
+                "temperature": temperature,
+                "top_p": top_p
+            }
+            
+            if stop_sequences:
+                params["stop"] = stop_sequences
+                
+            # Generate response
+            loop = asyncio.get_event_loop()
+            
+            if system_prompt:
+                # For chat models with system prompt
+                response_future = loop.run_in_executor(
+                    None,
+                    functools.partial(
+                        model.chat,
+                        system=system_prompt, 
+                        message=prompt,
+                        **params
+                    )
+                )
+            else:
+                # For completion models without system prompt
+                response_future = loop.run_in_executor(
+                    None,
+                    functools.partial(
+                        model.respond,
+                        prompt,
+                        **params
+                    )
+                )
+                
+            response = await response_future
+            
+            # Extract the generated text
+            if isinstance(response, dict) and "text" in response:
+                generated_text = response["text"]
+            elif isinstance(response, str):
+                generated_text = response
+            else:
+                generated_text = str(response)
+                
+            # Metadata
+            metadata = {
+                "model": model_id,
+                "usage": {
+                    "prompt_tokens": -1,  # LM Studio Python SDK doesn't provide token counts
+                    "completion_tokens": -1,
+                    "total_tokens": -1
+                }
+            }
+            
+            logger.debug(f"Generated {len(generated_text)} chars with model {model_id}")
+            return generated_text, metadata
+        except Exception as e:
+            logger.error(f"Failed to generate with model {model_id}: {str(e)}")
+            raise RuntimeError(f"Failed to generate with model {model_id}: {str(e)}")
+            
+    async def unload_model(self, model_id: str) -> None:
+        """Unload a model from LM Studio.
+        
+        Args:
+            model_id: The identifier of the model to unload
+        """
+        if not self.initialized:
+            return
+            
+        if model_id not in self.models:
+            logger.warning(f"Model {model_id} not loaded")
+            return
+            
+        try:
+            # Just remove the model reference, Python garbage collection will handle it
+            del self.models[model_id]
+            logger.info(f"Model {model_id} unloaded")
+        except Exception as e:
+            logger.error(f"Failed to unload model {model_id}: {str(e)}")
+            
+    async def shutdown(self) -> None:
+        """Shutdown the LM Studio provider."""
+        if not self.initialized:
+            return
+            
+        try:
+            # Clear all model references
+            self.models = {}
+            self.initialized = False
+            logger.info("LM Studio provider shut down")
+        except Exception as e:
+            logger.error(f"Failed to shut down LM Studio provider: {str(e)}")