mcp_utils.py

import mcp
from mcp.client.streamable_http import streamablehttp_client
import base64
import asyncio
import json
import uuid
import traceback

async def connect_to_mcp(mcp_server_url):
    """Connect to MCP server and return available tools"""
    config = {
        "type": "object",
        "default": {},
        "description": "Empty configuration for MCP Domino's Pizza server"
    }
    config_b64 = base64.b64encode(json.dumps(config).encode())
    url = mcp_server_url + f"&config={config_b64}"
    
    try:
        async with streamablehttp_client(url) as (read_stream, write_stream, _):
            async with mcp.ClientSession(read_stream, write_stream) as session:
                await session.initialize()
                tools_result = await session.list_tools()
                return [t.name for t in tools_result.tools]
    except Exception as e:
        print(f"Error connecting to MCP server: {e}")
        return None

async def execute_mcp_tool(tool_name: str, mcp_server_url: str, params: dict = None):
    """Execute a specific MCP tool with optional parameters"""
    config = {
        "type": "object",
        "default": {},
        "description": "Empty configuration for MCP Domino's Pizza server"
    }
    config_b64 = base64.b64encode(json.dumps(config).encode())
    url = mcp_server_url + f"&config={config_b64}"
    
    try:
        async with streamablehttp_client(url) as (read_stream, write_stream, _):
            async with mcp.ClientSession(read_stream, write_stream) as session:
                await session.initialize()
                try:
                    # First verify the tool exists
                    tools_result = await session.list_tools()
                    available_tools = [t.name for t in tools_result.tools]
                    
                    if tool_name not in available_tools:
                        return f"Tool '{tool_name}' not found. Available tools: {', '.join(available_tools)}"
                    
                    # Execute the tool using the correct method name with parameters
                    result = await session.call_tool(tool_name, params or {})
                    print(result.content[0].text)
                    if result:
                        try:
                            # Get the text from the first content item
                            result_text = result.content[0].text
                            # Parse the text as JSON
                            parsed_result = json.loads(result_text)
                            # Return the raw JSON
                            return json.dumps(parsed_result)
                        except json.JSONDecodeError:
                            # If not JSON, return as is
                            return result_text
                    else:
                        return f"Tool '{tool_name}' executed but returned no result"
                except Exception as e:
                    print(f"Error during tool execution: {str(e)}")
                    return f"Error executing tool '{tool_name}': {str(e)}"
    except Exception as e:
        print(f"Error connecting to MCP server: {str(e)}")
        return f"Failed to connect to MCP server: {str(e)}"

def handle_mcp_command(command_text, conversation_id, current_path):
    """Handle MCP-related commands"""
    try:
        # Parse the command
        parts = command_text.split()
        if len(parts) < 2:
            # Just /mcp - list tools
            tools = asyncio.run(connect_to_mcp())
            if tools:
                return f"Connected to MCP server. Available tools: {', '.join(tools)}"
            else:
                return "Failed to connect to MCP server"
        
        # Check for /mcp run <tool_name> [params]
        if parts[1] == "run" and len(parts) >= 3:
            tool_name = parts[2]
            # Parse parameters if provided
            params = {}
            if len(parts) > 3:
                try:
                    # Join the remaining parts in case the JSON contains spaces
                    params_str = " ".join(parts[3:])
                    params = json.loads(params_str)
                except json.JSONDecodeError as e:
                    return f"Invalid parameters format. Please provide valid JSON. Error: {str(e)}"
            
            # Execute the tool with parameters
            result = asyncio.run(execute_mcp_tool(tool_name, params))
            if result:
                return result
            else:
                return f"Failed to execute tool {tool_name}"
        else:
            return "Invalid MCP command. Use '/mcp' to list tools or '/mcp run <tool_name> [parameters]' to execute a tool"
            
    except Exception as e:
        print(f"Error in handle_mcp_command: {str(e)}")
        return f"Error handling MCP command: {str(e)}"

async def process_natural_language_command(command_text: str, mcp_server_url: str) -> tuple:
    """Process natural language command and map to MCP tool"""
    try:
        print(f"\n[DEBUG] Processing natural language command: {command_text}")
        
        # First get available tools
        print("[DEBUG] Attempting to connect to MCP server...")
        tools = await connect_to_mcp(mcp_server_url)
        if not tools:
            print("[ERROR] Failed to get tools from MCP server")
            return None, "Failed to connect to MCP server"
        print(f"[DEBUG] Available tools: {tools}")

        # Use Claude to map the command to a tool and extract parameters
        prompt = f"""Given the following command and available tools, determine which tool to use and extract parameters.
Command: {command_text}
Available tools: {', '.join(tools)}

Return the response in JSON format:
{{
    "tool": "tool_name",
    "parameters": {{"param1": "value1"}}
}}

Example:
Command: "find stores near 85054"
Response: {{
    "tool": "findNearbyStores",
    "parameters": {{"address": "85054"}}
}}"""

        print("[DEBUG] Sending command to Claude for analysis...")
        # Call Claude to analyze the command
        response = call_claude(prompt, "", str(uuid.uuid4()), "mcp", 
            "You are an AI assistant that helps map natural language commands to specific tools. Return only valid JSON.")
        
        if response:
            print(f"[DEBUG] Claude response: {response}")
            try:
                result = json.loads(response)
                tool_name = result.get("tool")
                params = result.get("parameters", {})
                print(f"[DEBUG] Parsed tool: {tool_name}")
                print(f"[DEBUG] Parsed parameters: {params}")
                return tool_name, params
            except json.JSONDecodeError as e:
                print(f"[ERROR] Failed to parse Claude's response as JSON: {e}")
                print(f"[ERROR] Raw response: {response}")
                return None, "Failed to parse tool mapping"
        else:
            print("[ERROR] Claude returned no response")
            return None, "Failed to analyze command"
    except Exception as e:
        print(f"[ERROR] Unexpected error in process_natural_language_command: {str(e)}")
        print(f"[ERROR] Full traceback: {traceback.format_exc()}")
        return None, str(e)