feat: support multiple MCPClient connections

tinyagent/mcp_client.py

@@ -1,162 +1,76 @@
-import asyncio
-import json
-import logging
-from typing import Dict, List, Optional, Any, Tuple, Callable
+#!/usr/bin/env python
+# coding=utf-8
+"""
+MCPClient implementation for TinyAgent based on HuggingFace SmolAgents
+"""
+from __future__ import annotations
 
-# Keep your MCPClient implementation unchanged
-import asyncio
-from contextlib import AsyncExitStack
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Union, List, Dict
 
-# MCP core imports
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
+from mcpadapt.core import MCPAdapt
+from mcpadapt.smolagents_adapter import SmolAgentsAdapter
+from smolagents.tools import Tool
 
-# Set up logging
-logger = logging.getLogger(__name__)
+__all__ = ["MCPClient"]
 
-class MCPClient:
-    def __init__(self, logger: Optional[logging.Logger] = None):
-        self.session = None
-        self.exit_stack = AsyncExitStack()
-        self.logger = logger or logging.getLogger(__name__)
-
-        # Simplified callback system
-        self.callbacks: List[callable] = []
-
-        self.logger.debug("MCPClient initialized")
+if TYPE_CHECKING:
+    from mcpadapt.core import StdioServerParameters
 
-    def add_callback(self, callback: callable) -> None:
-        """
-        Add a callback function to the client.
-
-        Args:
-            callback: A function that accepts (event_name, client, **kwargs)
-        """
-        self.callbacks.append(callback)
-
-    async def _run_callbacks(self, event_name: str, **kwargs) -> None:
-        """
-        Run all registered callbacks for an event.
-
-        Args:
-            event_name: The name of the event
-            **kwargs: Additional data for the event
-        """
-        for callback in self.callbacks:
-            try:
-                logger.debug(f"Running callback: {callback}")
-                if asyncio.iscoroutinefunction(callback):
-                    logger.debug(f"Callback is a coroutine function")
-                    await callback(event_name, self, **kwargs)
-                else:
-                    # Check if the callback is a class with an async __call__ method
-                    if hasattr(callback, '__call__') and asyncio.iscoroutinefunction(callback.__call__):
-                        logger.debug(f"Callback is a class with an async __call__ method")  
-                        await callback(event_name, self, **kwargs)
-                    else:
-                        logger.debug(f"Callback is a regular function")
-                        callback(event_name, self, **kwargs)
-            except Exception as e:
-                logger.error(f"Error in callback for {event_name}: {str(e)}")
+class MCPClient:
+    """Manage the connection to an MCP server and expose tools to TinyAgent."""
+    def __init__(
+        self,
+        server_parameters: Union[
+            StdioServerParameters,
+            Dict[str, Any],
+            List[Union[StdioServerParameters, Dict[str, Any]]],
+        ],
+    ):
+        try:
+            # Ensure MCPAdapt is available
+            import mcpadapt
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "Please install the 'mcp' extra to use MCPClient: `pip install 'tinyagent[mcp]'`"
+            )
+        self._adapter = MCPAdapt(server_parameters, SmolAgentsAdapter())
+        self._tools: List[Tool] | None = None
+        # Automatically connect on init
+        self.connect()
 
-    async def connect(self, command: str, args: list[str]):
-        """
-        Launches the MCP server subprocess and initializes the client session.
-        :param command: e.g. "python" or "node"
-        :param args: list of args to pass, e.g. ["my_server.py"] or ["build/index.js"]
-        """
-        # Prepare stdio transport parameters
-        params = StdioServerParameters(command=command, args=args)
-        # Open the stdio client transport
-        self.stdio, self.sock_write = await self.exit_stack.enter_async_context(
-            stdio_client(params)
-        )
-        # Create and initialize the MCP client session
-        self.session = await self.exit_stack.enter_async_context(
-            ClientSession(self.stdio, self.sock_write)
-        )
-        await self.session.initialize()
+    def connect(self) -> None:
+        """Connect to the MCP server and initialize tools."""
+        self._tools = self._adapter.__enter__()
 
-    async def list_tools(self):
-        resp = await self.session.list_tools()
-        print("Available tools:")
-        for tool in resp.tools:
-            print(f" • {tool.name}: {tool.description}")
+    def disconnect(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        exc_traceback: TracebackType | None = None,
+    ) -> None:
+        """Disconnect from the MCP server."""
+        self._adapter.__exit__(exc_type, exc_value, exc_traceback)
 
-    async def call_tool(self, name: str, arguments: dict):
-        """
-        Invokes a named tool and returns its raw content list.
-        """
-        # Notify tool start
-        await self._run_callbacks("tool_start", tool_name=name, arguments=arguments)
-
-        try:
-            resp = await self.session.call_tool(name, arguments)
-
-            # Notify tool end
-            await self._run_callbacks("tool_end", tool_name=name, arguments=arguments, 
-                                    result=resp.content, success=True)
-
-            return resp.content
-        except Exception as e:
-            # Notify tool end with error
-            await self._run_callbacks("tool_end", tool_name=name, arguments=arguments, 
-                                    error=str(e), success=False)
-            raise
+    def get_tools(self) -> List[Tool]:
+        """Return the list of tools from the MCP server."""
+        if self._tools is None:
+            raise ValueError(
+                "Could not retrieve tools, ensure `connect()` has been called."
+            )
+        return self._tools
 
-    async def close(self):
-        """Clean up subprocess and streams."""
-        if self.exit_stack:
-            try:
-                await self.exit_stack.aclose()
-            except (RuntimeError, asyncio.CancelledError) as e:
-                # Log the error but don't re-raise it
-                self.logger.error(f"Error during client cleanup: {e}")
-            finally:
-                # Always reset these regardless of success or failure
-                self.session = None
-                self.exit_stack = AsyncExitStack()
+    def __enter__(
+        self,
+    ) -> List[Tool]:
+        """Enter context, return tools."""
+        return self.get_tools()
 
-async def run_example():
-    """Example usage of MCPClient with proper logging."""
-    import sys
-    from tinyagent.hooks.logging_manager import LoggingManager
-
-    # Create and configure logging manager
-    log_manager = LoggingManager(default_level=logging.INFO)
-    log_manager.set_levels({
-        'tinyagent.mcp_client': logging.DEBUG,  # Debug for this module
-        'tinyagent.tiny_agent': logging.INFO,
-    })
-
-    # Configure a console handler
-    console_handler = logging.StreamHandler(sys.stdout)
-    log_manager.configure_handler(
-        console_handler,
-        format_string='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
-        level=logging.DEBUG
-    )
-
-    # Get module-specific logger
-    mcp_logger = log_manager.get_logger('tinyagent.mcp_client')
-
-    mcp_logger.debug("Starting MCPClient example")
-
-    # Create client with our logger
-    client = MCPClient(logger=mcp_logger)
-
-    try:
-        # Connect to a simple echo server
-        await client.connect("python", ["-m", "mcp.examples.echo_server"])
-
-        # List available tools
-        await client.list_tools()
-
-        # Call the echo tool
-        result = await client.call_tool("echo", {"message": "Hello, MCP!"})
-        mcp_logger.info(f"Echo result: {result}")
-
-    finally:
-        # Clean up
-        await client.close()
-        mcp_logger.debug("Example completed")
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        exc_traceback: TracebackType | None,
+    ) -> None:
+        """Exit context, disconnect."""
+        self.disconnect(exc_type, exc_value, exc_traceback)