mcp-use 0.1.0__py3-none-any.whl

mcp_use/connectors/websocket.py

@@ -0,0 +1,142 @@
+"""
+WebSocket connector for MCP implementations.
+
+This module provides a connector for communicating with MCP implementations
+through WebSocket connections.
+"""
+
+import asyncio
+import json
+import uuid
+from typing import Any
+
+import websockets
+from websockets.client import WebSocketClientProtocol
+
+from .base import BaseConnector
+
+
+class WebSocketConnector(BaseConnector):
+    """Connector for MCP implementations using WebSocket transport.
+
+    This connector uses WebSockets to communicate with remote MCP implementations.
+    """
+
+    def __init__(
+        self, url: str, auth_token: str | None = None, headers: dict[str, str] | None = None
+    ):
+        """Initialize a new WebSocket connector.
+
+        Args:
+            url: The WebSocket URL to connect to.
+            auth_token: Optional authentication token.
+            headers: Optional additional headers.
+        """
+        self.url = url
+        self.auth_token = auth_token
+        self.headers = headers or {}
+        if auth_token:
+            self.headers["Authorization"] = f"Bearer {auth_token}"
+        self.ws: WebSocketClientProtocol | None = None
+        self.pending_requests: dict[str, asyncio.Future] = {}
+
+    async def connect(self) -> None:
+        """Establish a connection to the MCP implementation."""
+        self.ws = await websockets.connect(self.url, extra_headers=self.headers)
+
+        # Start the message receiver task
+        self._receiver_task = asyncio.create_task(self._receive_messages())
+
+    async def _receive_messages(self) -> None:
+        """Continuously receive and process messages from the WebSocket."""
+        if not self.ws:
+            raise RuntimeError("WebSocket is not connected")
+
+        try:
+            async for message in self.ws:
+                # Parse the message
+                data = json.loads(message)
+
+                # Check if this is a response to a pending request
+                request_id = data.get("id")
+                if request_id and request_id in self.pending_requests:
+                    future = self.pending_requests.pop(request_id)
+                    if "result" in data:
+                        future.set_result(data["result"])
+                    elif "error" in data:
+                        future.set_exception(Exception(data["error"]))
+        except Exception as e:
+            # If the websocket connection was closed or errored,
+            # reject all pending requests
+            for future in self.pending_requests.values():
+                if not future.done():
+                    future.set_exception(e)
+
+    async def disconnect(self) -> None:
+        """Close the connection to the MCP implementation."""
+        if self._receiver_task:
+            self._receiver_task.cancel()
+            try:
+                await self._receiver_task
+            except asyncio.CancelledError:
+                pass
+
+        if self.ws:
+            await self.ws.close()
+            self.ws = None
+
+        # Reject any pending requests
+        for future in self.pending_requests.values():
+            if not future.done():
+                future.set_exception(ConnectionError("WebSocket disconnected"))
+        self.pending_requests.clear()
+
+    async def _send_request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """Send a request and wait for a response."""
+        if not self.ws:
+            raise RuntimeError("WebSocket is not connected")
+
+        # Create a request ID
+        request_id = str(uuid.uuid4())
+
+        # Create a future to receive the response
+        future = asyncio.Future()
+        self.pending_requests[request_id] = future
+
+        # Send the request
+        await self.ws.send(json.dumps({"id": request_id, "method": method, "params": params or {}}))
+
+        # Wait for the response
+        try:
+            return await future
+        except Exception as e:
+            # Remove the request from pending requests
+            self.pending_requests.pop(request_id, None)
+            raise e
+
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and return session information."""
+        return await self._send_request("initialize")
+
+    async def list_tools(self) -> list[dict[str, Any]]:
+        """List all available tools from the MCP implementation."""
+        result = await self._send_request("tools/list")
+        return result.get("tools", [])
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
+        """Call an MCP tool with the given arguments."""
+        return await self._send_request("tools/call", {"name": name, "arguments": arguments})
+
+    async def list_resources(self) -> list[dict[str, Any]]:
+        """List all available resources from the MCP implementation."""
+        result = await self._send_request("resources/list")
+        return result
+
+    async def read_resource(self, uri: str) -> tuple[bytes, str]:
+        """Read a resource by URI."""
+        result = await self._send_request("resources/read", {"uri": uri})
+        return result.get("content", b""), result.get("mimeType", "")
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """Send a raw request to the MCP implementation."""
+        return await self._send_request(method, params)

mcp_use/logging.py

@@ -0,0 +1,96 @@
+"""
+Logger module for mcp_use.
+
+This module provides a centralized logging configuration for the mcp_use library,
+with customizable log levels and formatters.
+"""
+
+import logging
+import os
+import sys
+
+
+class Logger:
+    """Centralized logger for mcp_use.
+
+    This class provides logging functionality with configurable levels,
+    formatters, and handlers.
+    """
+
+    # Default log format
+    DEFAULT_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+
+    # Module-specific loggers
+    _loggers = {}
+
+    @classmethod
+    def get_logger(cls, name: str = "mcp_use") -> logging.Logger:
+        """Get a logger instance for the specified name.
+
+        Args:
+            name: Logger name, usually the module name (defaults to 'mcp_use')
+
+        Returns:
+            Configured logger instance
+        """
+        if name in cls._loggers:
+            return cls._loggers[name]
+
+        # Create new logger
+        logger = logging.getLogger(name)
+        cls._loggers[name] = logger
+
+        return logger
+
+    @classmethod
+    def configure(
+        cls,
+        level: int | str = logging.INFO,
+        format_str: str | None = None,
+        log_to_console: bool = True,
+        log_to_file: str | None = None,
+    ) -> None:
+        """Configure the root mcp_use logger.
+
+        Args:
+            level: Log level (default: INFO)
+            format_str: Log format string (default: DEFAULT_FORMAT)
+            log_to_console: Whether to log to console (default: True)
+            log_to_file: Path to log file (default: None)
+        """
+        root_logger = cls.get_logger()
+
+        # Set level
+        if isinstance(level, str):
+            level = getattr(logging, level.upper())
+        root_logger.setLevel(level)
+
+        # Clear existing handlers
+        for handler in root_logger.handlers[:]:
+            root_logger.removeHandler(handler)
+
+        # Set formatter
+        formatter = logging.Formatter(format_str or cls.DEFAULT_FORMAT)
+
+        # Add console handler if requested
+        if log_to_console:
+            console_handler = logging.StreamHandler(sys.stdout)
+            console_handler.setFormatter(formatter)
+            root_logger.addHandler(console_handler)
+
+        # Add file handler if requested
+        if log_to_file:
+            # Ensure directory exists
+            log_dir = os.path.dirname(log_to_file)
+            if log_dir and not os.path.exists(log_dir):
+                os.makedirs(log_dir)
+
+            file_handler = logging.FileHandler(log_to_file)
+            file_handler.setFormatter(formatter)
+            root_logger.addHandler(file_handler)
+
+
+# Configure default logger
+Logger.configure()
+
+logger = Logger.get_logger()

mcp_use/session.py

@@ -0,0 +1,168 @@
+"""
+Session manager for MCP connections.
+
+This module provides a session manager for MCP connections,
+which handles authentication, initialization, and tool discovery.
+"""
+
+from typing import Any
+
+from .connectors.base import BaseConnector
+from .tools.converter import ModelProvider, ToolConverter
+
+
+class MCPSession:
+    """Session manager for MCP connections.
+
+    This class manages the lifecycle of an MCP connection, including
+    authentication, initialization, and tool discovery.
+    """
+
+    def __init__(
+        self,
+        connector: BaseConnector,
+        model_provider: str | ModelProvider,
+        auto_connect: bool = True,
+    ) -> None:
+        """Initialize a new MCP session.
+
+        Args:
+            connector: The connector to use for communicating with the MCP implementation.
+            model_provider: The model provider to convert tools for.
+            auto_connect: Whether to automatically connect to the MCP implementation.
+        """
+        self.connector = connector
+        self.tool_converter = ToolConverter(model_provider)
+        self.session_info: dict[str, Any] | None = None
+        self.tools: list[dict[str, Any]] = []
+        self.auto_connect = auto_connect
+
+    async def __aenter__(self) -> "MCPSession":
+        """Enter the async context manager.
+
+        Returns:
+            The session instance.
+        """
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async context manager.
+
+        Args:
+            exc_type: The exception type, if an exception was raised.
+            exc_val: The exception value, if an exception was raised.
+            exc_tb: The exception traceback, if an exception was raised.
+        """
+        await self.disconnect()
+
+    async def connect(self) -> None:
+        """Connect to the MCP implementation."""
+        await self.connector.connect()
+
+    async def disconnect(self) -> None:
+        """Disconnect from the MCP implementation."""
+        await self.connector.disconnect()
+
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and discover available tools.
+
+        Returns:
+            The session information returned by the MCP implementation.
+        """
+        # Make sure we're connected
+        if not self.is_connected and self.auto_connect:
+            await self.connect()
+
+        # Initialize the session
+        self.session_info = await self.connector.initialize()
+
+        # Discover available tools
+        await self.discover_tools()
+
+        return self.session_info
+
+    @property
+    def is_connected(self) -> bool:
+        """Check if the connector is connected.
+
+        Returns:
+            True if the connector is connected, False otherwise.
+        """
+        return hasattr(self.connector, "client") and self.connector.client is not None
+
+    async def discover_tools(self) -> list[dict[str, Any]]:
+        """Discover available tools from the MCP implementation.
+
+        Returns:
+            The list of available tools in MCP format.
+        """
+        self.tools = self.connector.tools
+        return self.tools
+
+    def get_tools_for_llm(self) -> list[dict[str, Any]]:
+        """Get the tools in the format required by the LLM.
+
+        Returns:
+            The list of tools in the LLM-specific format.
+        """
+        return self.tool_converter.convert_tools(self.tools)
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
+        """Call an MCP tool with the given arguments.
+
+        Args:
+            name: The name of the tool to call.
+            arguments: The arguments to pass to the tool.
+
+        Returns:
+            The result of the tool call.
+        """
+        # Make sure we're connected
+        if not self.is_connected and self.auto_connect:
+            await self.connect()
+
+        return await self.connector.call_tool(name, arguments)
+
+    async def process_tool_calls(self, response: dict[str, Any]) -> list[dict[str, Any]]:
+        """Process tool calls from an LLM response.
+
+        This method parses tool calls from the LLM response, executes them,
+        and returns the results.
+
+        Args:
+            response: The response from the LLM.
+
+        Returns:
+            A list of tool call results, each containing 'name', 'arguments', and 'result' keys.
+        """
+        # Parse tool calls from the response
+        tool_calls = self.tool_converter.parse_tool_calls(response)
+
+        # Execute each tool call
+        results = []
+        for tool_call in tool_calls:
+            name = tool_call["name"]
+            arguments = tool_call["arguments"]
+
+            try:
+                result = await self.call_tool(name, arguments)
+                results.append(
+                    {
+                        "name": name,
+                        "arguments": arguments,
+                        "result": result,
+                        "error": None,
+                    }
+                )
+            except Exception as e:
+                results.append(
+                    {
+                        "name": name,
+                        "arguments": arguments,
+                        "result": None,
+                        "error": str(e),
+                    }
+                )
+
+        return results

mcp_use/tools/__init__.py

@@ -0,0 +1,11 @@
+"""
+Tool conversion utilities.
+
+This module provides utilities for converting between MCP tool schemas
+and LLM-specific tool formats.
+"""
+
+from .converter import ToolConverter
+from .formats import AnthropicToolFormat, OpenAIToolFormat, ToolFormat
+
+__all__ = ["ToolConverter", "ToolFormat", "OpenAIToolFormat", "AnthropicToolFormat"]

mcp_use/tools/converter.py

@@ -0,0 +1,108 @@
+"""
+Tool converter for different LLM providers.
+
+This module provides utilities for converting between MCP tool schemas
+and LLM-specific formats.
+"""
+
+from enum import Enum, auto
+from typing import Any
+
+from .formats import AnthropicToolFormat, OpenAIToolFormat, ToolFormat
+
+
+class ModelProvider(Enum):
+    """Enum for supported model providers."""
+
+    OPENAI = auto()
+    ANTHROPIC = auto()
+
+    @classmethod
+    def from_string(cls, value: str) -> "ModelProvider":
+        """Convert a string to a ModelProvider enum.
+
+        Args:
+            value: The string to convert.
+
+        Returns:
+            The corresponding ModelProvider enum value.
+
+        Raises:
+            ValueError: If the string is not a valid model provider.
+        """
+        value = value.lower()
+        if value in ("openai", "open_ai", "open-ai"):
+            return cls.OPENAI
+        elif value in ("anthropic", "claude"):
+            return cls.ANTHROPIC
+        else:
+            raise ValueError(f"Unsupported model provider: {value}")
+
+
+class ToolConverter:
+    """Converter for MCP tools to different LLM formats.
+
+    This class provides utilities for converting between MCP tool schemas
+    and LLM-specific formats.
+    """
+
+    _format_classes: dict[ModelProvider, type[ToolFormat]] = {
+        ModelProvider.OPENAI: OpenAIToolFormat,
+        ModelProvider.ANTHROPIC: AnthropicToolFormat,
+    }
+
+    def __init__(self, provider: str | ModelProvider) -> None:
+        """Initialize a new tool converter.
+
+        Args:
+            provider: The model provider to convert tools for.
+                Can be a string or a ModelProvider enum.
+
+        Raises:
+            ValueError: If the provider is not supported.
+        """
+        if isinstance(provider, str):
+            self.provider = ModelProvider.from_string(provider)
+        else:
+            self.provider = provider
+
+        # Create an instance of the appropriate format class
+        format_class = self._format_classes.get(self.provider)
+        if not format_class:
+            raise ValueError(f"Unsupported model provider: {provider}")
+
+        self._format = format_class()
+
+    def convert_tools(self, tools: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        """Convert a list of MCP tools to the LLM-specific format.
+
+        Args:
+            tools: The list of MCP tools to convert.
+
+        Returns:
+            The converted tools in the LLM-specific format.
+        """
+        return [self._format.convert_tool(tool) for tool in tools]
+
+    def convert_tool_call(self, name: str, arguments: dict[str, Any]) -> dict[str, Any]:
+        """Convert a tool call to the MCP format.
+
+        Args:
+            name: The name of the tool being called.
+            arguments: The arguments for the tool call.
+
+        Returns:
+            The converted tool call in the MCP format.
+        """
+        return self._format.convert_tool_call(name, arguments)
+
+    def parse_tool_calls(self, response: dict[str, Any]) -> list[dict[str, Any]]:
+        """Parse tool calls from an LLM response.
+
+        Args:
+            response: The response from the LLM.
+
+        Returns:
+            A list of parsed tool calls, each containing 'name' and 'arguments' keys.
+        """
+        return self._format.parse_tool_call(response)