mcp-use 0.0.3__py3-none-any.whl

mcp_use/client.py

@@ -0,0 +1,222 @@
+"""
+Client for managing MCP servers and sessions.
+
+This module provides a high-level client that manages MCP servers, connectors,
+and sessions from configuration.
+"""
+
+import json
+from typing import Any
+
+from .config import create_connector_from_config, load_config_file
+from .session import MCPSession
+
+
+class MCPClient:
+    """Client for managing MCP servers and sessions.
+
+    This class provides a unified interface for working with MCP servers,
+    handling configuration, connector creation, and session management.
+    """
+
+    def __init__(
+        self,
+        config: str | dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize a new MCP client.
+
+        Args:
+            config: Either a dict containing configuration or a path to a JSON config file.
+                   If None, an empty configuration is used.
+        """
+        self.config: dict[str, Any] = {}
+        self.sessions: dict[str, MCPSession] = {}
+        self.active_session: str | None = None
+
+        # Load configuration if provided
+        if config is not None:
+            if isinstance(config, str):
+                self.config = load_config_file(config)
+            else:
+                self.config = config
+
+    @classmethod
+    def from_dict(cls, config: dict[str, Any]) -> "MCPClient":
+        """Create a MCPClient from a dictionary.
+
+        Args:
+            config: The configuration dictionary.
+        """
+        return cls(config=config)
+
+    @classmethod
+    def from_config_file(cls, filepath: str) -> "MCPClient":
+        """Create a MCPClient from a configuration file.
+
+        Args:
+            filepath: The path to the configuration file.
+        """
+        return cls(config=load_config_file(filepath))
+
+    def add_server(
+        self,
+        name: str,
+        server_config: dict[str, Any],
+    ) -> None:
+        """Add a server configuration.
+
+        Args:
+            name: The name to identify this server.
+            server_config: The server configuration.
+        """
+        if "mcpServers" not in self.config:
+            self.config["mcpServers"] = {}
+
+        self.config["mcpServers"][name] = server_config
+
+    def remove_server(self, name: str) -> None:
+        """Remove a server configuration.
+
+        Args:
+            name: The name of the server to remove.
+        """
+        if "mcpServers" in self.config and name in self.config["mcpServers"]:
+            del self.config["mcpServers"][name]
+
+            # If we removed the active session, set active_session to None
+            if name == self.active_session:
+                self.active_session = None
+
+    def get_server_names(self) -> list[str]:
+        """Get the list of configured server names.
+
+        Returns:
+            List of server names.
+        """
+        return list(self.config.get("mcpServers", {}).keys())
+
+    def save_config(self, filepath: str) -> None:
+        """Save the current configuration to a file.
+
+        Args:
+            filepath: The path to save the configuration to.
+        """
+        with open(filepath, "w") as f:
+            json.dump(self.config, f, indent=2)
+
+    async def create_session(
+        self,
+        server_name: str | None = None,
+        auto_initialize: bool = True,
+    ) -> MCPSession:
+        """Create a session for the specified server.
+
+        Args:
+            server_name: The name of the server to create a session for.
+                        If None, uses the first available server.
+            auto_initialize: Whether to automatically initialize the session.
+
+        Returns:
+            The created MCPSession.
+
+        Raises:
+            ValueError: If no servers are configured or the specified server doesn't exist.
+        """
+        # Get server config
+        servers = self.config.get("mcpServers", {})
+        if not servers:
+            raise ValueError("No MCP servers defined in config")
+
+        # If server_name not specified, use the first one
+        if not server_name:
+            server_name = next(iter(servers.keys()))
+
+        if server_name not in servers:
+            raise ValueError(f"Server '{server_name}' not found in config")
+
+        server_config = servers[server_name]
+        connector = create_connector_from_config(server_config)
+
+        # Create the session
+        session = MCPSession(connector)
+        self.sessions[server_name] = session
+
+        # Make this the active session
+        self.active_session = server_name
+
+        # Initialize if requested
+        if auto_initialize:
+            await session.initialize()
+
+        return session
+
+    def get_session(self, server_name: str | None = None) -> MCPSession:
+        """Get an existing session.
+
+        Args:
+            server_name: The name of the server to get the session for.
+                        If None, uses the active session.
+
+        Returns:
+            The MCPSession for the specified server.
+
+        Raises:
+            ValueError: If no active session exists or the specified session doesn't exist.
+        """
+        if server_name is None:
+            if self.active_session is None:
+                raise ValueError("No active session")
+            server_name = self.active_session
+
+        if server_name not in self.sessions:
+            raise ValueError(f"No session exists for server '{server_name}'")
+
+        return self.sessions[server_name]
+
+    async def close_session(self, server_name: str | None = None) -> None:
+        """Close a session.
+
+        Args:
+            server_name: The name of the server to close the session for.
+                        If None, uses the active session.
+
+        Raises:
+            ValueError: If no active session exists or the specified session doesn't exist.
+        """
+        session = self.get_session(server_name)
+        await session.disconnect()
+
+        # Remove the session
+        if server_name is None:
+            server_name = self.active_session
+
+        if server_name in self.sessions:
+            del self.sessions[server_name]
+
+        # If we closed the active session, set active_session to None
+        if server_name == self.active_session:
+            self.active_session = None
+
+    async def close_all_sessions(self) -> None:
+        """Close all active sessions."""
+        for server_name in list(self.sessions.keys()):
+            await self.close_session(server_name)
+
+    async def __aenter__(self) -> "MCPClient":
+        """Enter the async context manager.
+
+        Creates a session for the first available server if no sessions exist.
+
+        Returns:
+            The client instance.
+        """
+        if not self.sessions and self.config.get("mcpServers"):
+            await self.create_session()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async context manager.
+
+        Closes all active sessions.
+        """
+        await self.close_all_sessions()

mcp_use/config.py

@@ -0,0 +1,93 @@
+"""
+Configuration loader for MCP session.
+
+This module provides functionality to load MCP configuration from JSON files.
+"""
+
+import json
+from typing import Any
+
+from .connectors import BaseConnector, HttpConnector, StdioConnector, WebSocketConnector
+from .session import MCPSession
+
+
+def load_config_file(filepath: str) -> dict[str, Any]:
+    """Load a configuration file.
+
+    Args:
+        filepath: Path to the configuration file
+
+    Returns:
+        The parsed configuration
+    """
+    with open(filepath) as f:
+        return json.load(f)
+
+
+def create_connector_from_config(server_config: dict[str, Any]) -> BaseConnector:
+    """Create a connector based on server configuration.
+
+    Args:
+        server_config: The server configuration section
+
+    Returns:
+        A configured connector instance
+    """
+    # Stdio connector (command-based)
+    if "command" in server_config and "args" in server_config:
+        return StdioConnector(
+            command=server_config["command"],
+            args=server_config["args"],
+            env=server_config.get("env", None),
+        )
+
+    # HTTP connector
+    elif "url" in server_config:
+        return HttpConnector(
+            url=server_config["url"],
+            headers=server_config.get("headers", None),
+            auth=server_config.get("auth", None),
+        )
+
+    # WebSocket connector
+    elif "ws_url" in server_config:
+        return WebSocketConnector(
+            url=server_config["ws_url"],
+            headers=server_config.get("headers", None),
+            auth=server_config.get("auth", None),
+        )
+
+    raise ValueError("Cannot determine connector type from config")
+
+
+def create_session_from_config(
+    filepath: str,
+    server_name: str | None = None,
+) -> MCPSession:
+    """Create an MCPSession from a configuration file.
+
+    Args:
+        filepath: Path to the configuration file
+        server_name: Name of the server to use from config, uses first if None
+
+    Returns:
+        Configured MCPSession instance
+    """
+    config = load_config_file(filepath)
+
+    # Get server config
+    servers = config.get("mcpServers", {})
+    if not servers:
+        raise ValueError("No MCP servers defined in config")
+
+    # If server_name not specified, use the first one
+    if not server_name:
+        server_name = next(iter(servers.keys()))
+
+    if server_name not in servers:
+        raise ValueError(f"Server '{server_name}' not found in config")
+
+    server_config = servers[server_name]
+    connector = create_connector_from_config(server_config)
+
+    return MCPSession(connector)

mcp_use/connectors/__init__.py

@@ -0,0 +1,13 @@
+"""
+Connectors for various MCP transports.
+
+This module provides interfaces for connecting to MCP implementations
+through different transport mechanisms.
+"""
+
+from .base import BaseConnector
+from .http import HttpConnector
+from .stdio import StdioConnector
+from .websocket import WebSocketConnector
+
+__all__ = ["BaseConnector", "StdioConnector", "WebSocketConnector", "HttpConnector"]

mcp_use/connectors/base.py

@@ -0,0 +1,59 @@
+"""
+Base connector for MCP implementations.
+
+This module provides the base connector interface that all MCP connectors
+must implement.
+"""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from mcp.types import CallToolResult, Tool
+
+
+class BaseConnector(ABC):
+    """Base class for MCP connectors.
+
+    This class defines the interface that all MCP connectors must implement.
+    """
+
+    @abstractmethod
+    async def connect(self) -> None:
+        """Establish a connection to the MCP implementation."""
+        pass
+
+    @abstractmethod
+    async def disconnect(self) -> None:
+        """Close the connection to the MCP implementation."""
+        pass
+
+    @abstractmethod
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and return session information."""
+        pass
+
+    @property
+    @abstractmethod
+    def tools(self) -> list[Tool]:
+        """Get the list of available tools."""
+        pass
+
+    @abstractmethod
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> CallToolResult:
+        """Call an MCP tool with the given arguments."""
+        pass
+
+    @abstractmethod
+    async def list_resources(self) -> list[dict[str, Any]]:
+        """List all available resources from the MCP implementation."""
+        pass
+
+    @abstractmethod
+    async def read_resource(self, uri: str) -> tuple[bytes, str]:
+        """Read a resource by URI."""
+        pass
+
+    @abstractmethod
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """Send a raw request to the MCP implementation."""
+        pass

mcp_use/connectors/http.py

@@ -0,0 +1,126 @@
+"""
+HTTP connector for MCP implementations.
+
+This module provides a connector for communicating with MCP implementations
+through HTTP APIs.
+"""
+
+from typing import Any
+
+import aiohttp
+
+from .base import BaseConnector
+
+
+class HttpConnector(BaseConnector):
+    """Connector for MCP implementations using HTTP transport.
+
+    This connector uses HTTP requests to communicate with remote MCP implementations.
+    """
+
+    def __init__(
+        self, base_url: str, auth_token: str | None = None, headers: dict[str, str] | None = None
+    ):
+        """Initialize a new HTTP connector.
+
+        Args:
+            base_url: The base URL of the MCP HTTP API.
+            auth_token: Optional authentication token.
+            headers: Optional additional headers.
+        """
+        self.base_url = base_url.rstrip("/")
+        self.auth_token = auth_token
+        self.headers = headers or {}
+        if auth_token:
+            self.headers["Authorization"] = f"Bearer {auth_token}"
+        self.session: aiohttp.ClientSession | None = None
+
+    async def connect(self) -> None:
+        """Establish a connection to the MCP implementation."""
+        self.session = aiohttp.ClientSession(headers=self.headers)
+
+    async def disconnect(self) -> None:
+        """Close the connection to the MCP implementation."""
+        if self.session:
+            await self.session.close()
+            self.session = None
+
+    async def _request(self, method: str, endpoint: str, data: dict[str, Any] | None = None) -> Any:
+        """Send an HTTP request to the MCP API.
+
+        Args:
+            method: The HTTP method (GET, POST, etc.).
+            endpoint: The API endpoint path.
+            data: Optional request data.
+
+        Returns:
+            The parsed JSON response.
+        """
+        if not self.session:
+            raise RuntimeError("HTTP session is not connected")
+
+        url = f"{self.base_url}/{endpoint.lstrip('/')}"
+
+        if method.upper() == "GET" and data:
+            # For GET requests, convert data to query parameters
+            async with self.session.get(url, params=data) as response:
+                response.raise_for_status()
+                return await response.json()
+        else:
+            # For other methods, send data as JSON body
+            async with self.session.request(method, url, json=data) as response:
+                response.raise_for_status()
+                return await response.json()
+
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and return session information."""
+        return await self._request("POST", "initialize")
+
+    async def list_tools(self) -> list[dict[str, Any]]:
+        """List all available tools from the MCP implementation."""
+        result = await self._request("GET", "tools")
+        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._request("POST", f"tools/{name}", arguments)
+
+    async def list_resources(self) -> list[dict[str, Any]]:
+        """List all available resources from the MCP implementation."""
+        result = await self._request("GET", "resources")
+        return result
+
+    async def read_resource(self, uri: str) -> tuple[bytes, str]:
+        """Read a resource by URI."""
+        # For resources, we may need to handle binary data
+        if not self.session:
+            raise RuntimeError("HTTP session is not connected")
+
+        url = f"{self.base_url}/resources/read"
+
+        async with self.session.get(url, params={"uri": uri}) as response:
+            response.raise_for_status()
+
+            # Check if this is a JSON response or binary data
+            content_type = response.headers.get("Content-Type", "")
+            if "application/json" in content_type:
+                data = await response.json()
+                content = data.get("content", b"")
+                mime_type = data.get("mimeType", "")
+
+                # If content is base64 encoded, decode it
+                if isinstance(content, str):
+                    import base64
+
+                    content = base64.b64decode(content)
+
+                return content, mime_type
+            else:
+                # Assume binary response
+                content = await response.read()
+                return content, content_type
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """Send a raw request to the MCP implementation."""
+        # For custom methods, we'll use the RPC-style endpoint
+        return await self._request("POST", "rpc", {"method": method, "params": params or {}})

mcp_use/connectors/stdio.py

@@ -0,0 +1,124 @@
+"""
+StdIO connector for MCP implementations.
+
+This module provides a connector for communicating with MCP implementations
+through the standard input/output streams.
+"""
+
+from typing import Any
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+from mcp.types import Tool
+
+from ..logging import logger
+from .base import BaseConnector
+
+
+class StdioConnector(BaseConnector):
+    """Connector for MCP implementations using stdio transport.
+
+    This connector uses the stdio transport to communicate with MCP implementations
+    that are executed as child processes.
+    """
+
+    def __init__(
+        self,
+        command: str = "npx",
+        args: list[str] | None = None,
+        env: dict[str, str] | None = None,
+    ):
+        """Initialize a new stdio connector.
+
+        Args:
+            command: The command to execute.
+            args: Optional command line arguments.
+            env: Optional environment variables.
+        """
+        self.command = command
+        self.args = args
+        self.env = env
+        self.client: ClientSession | None = None
+        self._stdio_ctx = None
+        self._tools: list[Tool] | None = None
+
+    async def connect(self) -> None:
+        """Establish a connection to the MCP implementation."""
+        server_params = StdioServerParameters(command=self.command, args=self.args, env=self.env)
+        self._stdio_ctx = stdio_client(server_params)
+        read_stream, write_stream = await self._stdio_ctx.__aenter__()
+        self.client = ClientSession(read_stream, write_stream, sampling_callback=None)
+        await self.client.__aenter__()
+
+    async def disconnect(self) -> None:
+        """Close the connection to the MCP implementation."""
+        try:
+            if self.client:
+                await self.client.__aexit__(None, None, None)
+                self.client = None
+            if self._stdio_ctx:
+                await self._stdio_ctx.__aexit__(None, None, None)
+                self._stdio_ctx = None
+        except Exception as e:
+            logger.warning(f"Warning: Error during stdio connector disconnect: {e}")
+        finally:
+            # Always clean up references even if there were errors
+            self.client = None
+            self._stdio_ctx = None
+            self._tools = None
+
+    async def __aenter__(self) -> "StdioConnector":
+        """Enter the async context manager."""
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the async context manager."""
+        await self.disconnect()
+
+    async def initialize(self) -> dict[str, Any]:
+        """Initialize the MCP session and return session information."""
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+
+        # Initialize the session
+        result = await self.client.initialize()
+
+        # Get available tools
+        tools_result = await self.client.list_tools()
+        self._tools = tools_result.tools
+
+        return result
+
+    @property
+    def tools(self) -> list[Tool]:
+        """Get the list of available tools."""
+        if not self._tools:
+            raise RuntimeError("MCP client is not initialized")
+        return self._tools
+
+    async def call_tool(self, name: str, arguments: dict[str, Any]) -> Any:
+        """Call an MCP tool with the given arguments."""
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+        return await self.client.call_tool(name, arguments)
+
+    async def list_resources(self) -> list[dict[str, Any]]:
+        """List all available resources from the MCP implementation."""
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+        resources = await self.client.list_resources()
+        return resources
+
+    async def read_resource(self, uri: str) -> tuple[bytes, str]:
+        """Read a resource by URI."""
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+        resource = await self.client.read_resource(uri)
+        return resource.content, resource.mimeType
+
+    async def request(self, method: str, params: dict[str, Any] | None = None) -> Any:
+        """Send a raw request to the MCP implementation."""
+        if not self.client:
+            raise RuntimeError("MCP client is not connected")
+        return await self.client.request({"method": method, "params": params or {}})