mcp-use 0.0.3__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,113 @@
+"""
+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
+
+
+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,
+        auto_connect: bool = True,
+    ) -> None:
+        """Initialize a new MCP session.
+
+        Args:
+            connector: The connector to use for communicating with the MCP implementation.
+            auto_connect: Whether to automatically connect to the MCP implementation.
+        """
+        self.connector = connector
+        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
+
+    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)