agno 2.0.1__py3-none-any.whl → 2.3.0__py3-none-any.whl

agno/tools/mcp/multi_mcp.py

@@ -0,0 +1,347 @@
+import weakref
+from contextlib import AsyncExitStack
+from dataclasses import asdict
+from datetime import timedelta
+from types import TracebackType
+from typing import List, Literal, Optional, Union
+
+from agno.tools import Toolkit
+from agno.tools.function import Function
+from agno.tools.mcp.params import SSEClientParams, StreamableHTTPClientParams
+from agno.utils.log import log_debug, log_error, log_info, log_warning
+from agno.utils.mcp import get_entrypoint_for_tool, prepare_command
+
+try:
+    from mcp import ClientSession, StdioServerParameters
+    from mcp.client.sse import sse_client
+    from mcp.client.stdio import get_default_environment, stdio_client
+    from mcp.client.streamable_http import streamablehttp_client
+except (ImportError, ModuleNotFoundError):
+    raise ImportError("`mcp` not installed. Please install using `pip install mcp`")
+
+
+class MultiMCPTools(Toolkit):
+    """
+    A toolkit for integrating multiple Model Context Protocol (MCP) servers with Agno agents.
+    This allows agents to access tools, resources, and prompts exposed by MCP servers.
+
+    Can be used in three ways:
+    1. Direct initialization with a ClientSession
+    2. As an async context manager with StdioServerParameters
+    3. As an async context manager with SSE or Streamable HTTP endpoints
+    """
+
+    def __init__(
+        self,
+        commands: Optional[List[str]] = None,
+        urls: Optional[List[str]] = None,
+        urls_transports: Optional[List[Literal["sse", "streamable-http"]]] = None,
+        *,
+        env: Optional[dict[str, str]] = None,
+        server_params_list: Optional[
+            list[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]]
+        ] = None,
+        timeout_seconds: int = 10,
+        client=None,
+        include_tools: Optional[list[str]] = None,
+        exclude_tools: Optional[list[str]] = None,
+        refresh_connection: bool = False,
+        allow_partial_failure: bool = False,
+        **kwargs,
+    ):
+        """
+        Initialize the MCP toolkit.
+
+        Args:
+            commands: List of commands to run to start the servers. Should be used in conjunction with env.
+            urls: List of URLs for SSE and/or Streamable HTTP endpoints.
+            urls_transports: List of transports to use for the given URLs.
+            server_params_list: List of StdioServerParameters or SSEClientParams or StreamableHTTPClientParams for creating new sessions.
+            env: The environment variables to pass to the servers. Should be used in conjunction with commands.
+            client: The underlying MCP client (optional, used to prevent garbage collection).
+            timeout_seconds: Timeout in seconds for managing timeouts for Client Session if Agent or Tool doesn't respond.
+            include_tools: Optional list of tool names to include (if None, includes all).
+            exclude_tools: Optional list of tool names to exclude (if None, excludes none).
+            allow_partial_failure: If True, allows toolkit to initialize even if some MCP servers fail to connect. If False, any failure will raise an exception.
+            refresh_connection: If True, the connection and tools will be refreshed on each run
+        """
+        super().__init__(name="MultiMCPTools", **kwargs)
+
+        if urls_transports is not None:
+            if "sse" in urls_transports:
+                log_info("SSE as a standalone transport is deprecated. Please use Streamable HTTP instead.")
+
+        if urls is not None:
+            if urls_transports is None:
+                log_warning(
+                    "The default transport 'streamable-http' will be used. You can explicitly set the transports by providing the urls_transports parameter."
+                )
+            else:
+                if len(urls) != len(urls_transports):
+                    raise ValueError("urls and urls_transports must be of the same length")
+
+        # Set these after `__init__` to bypass the `_check_tools_filters`
+        # beacuse tools are not available until `initialize()` is called.
+        self.include_tools = include_tools
+        self.exclude_tools = exclude_tools
+        self.refresh_connection = refresh_connection
+
+        if server_params_list is None and commands is None and urls is None:
+            raise ValueError("Either server_params_list or commands or urls must be provided")
+
+        self.server_params_list: List[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]] = (
+            server_params_list or []
+        )
+        self.timeout_seconds = timeout_seconds
+        self.commands: Optional[List[str]] = commands
+        self.urls: Optional[List[str]] = urls
+        # Merge provided env with system env
+        if env is not None:
+            env = {
+                **get_default_environment(),
+                **env,
+            }
+        else:
+            env = get_default_environment()
+
+        if commands is not None:
+            for command in commands:
+                parts = prepare_command(command)
+                cmd = parts[0]
+                arguments = parts[1:] if len(parts) > 1 else []
+                self.server_params_list.append(StdioServerParameters(command=cmd, args=arguments, env=env))
+
+        if urls is not None:
+            if urls_transports is not None:
+                for url, transport in zip(urls, urls_transports):
+                    if transport == "streamable-http":
+                        self.server_params_list.append(StreamableHTTPClientParams(url=url))
+                    else:
+                        self.server_params_list.append(SSEClientParams(url=url))
+            else:
+                for url in urls:
+                    self.server_params_list.append(StreamableHTTPClientParams(url=url))
+
+        self._async_exit_stack = AsyncExitStack()
+
+        self._client = client
+
+        self._initialized = False
+        self._connection_task = None
+        self._successful_connections = 0
+        self._sessions: list[ClientSession] = []
+
+        self.allow_partial_failure = allow_partial_failure
+
+        def cleanup():
+            """Cancel active connections"""
+            if self._connection_task and not self._connection_task.done():
+                self._connection_task.cancel()
+
+        # Setup cleanup logic before the instance is garbage collected
+        self._cleanup_finalizer = weakref.finalize(self, cleanup)
+
+    @property
+    def initialized(self) -> bool:
+        return self._initialized
+
+    async def is_alive(self) -> bool:
+        try:
+            for session in self._sessions:
+                await session.send_ping()
+            return True
+        except (RuntimeError, BaseException):
+            return False
+
+    async def connect(self, force: bool = False):
+        """Initialize a MultiMCPTools instance and connect to the MCP servers"""
+
+        if force:
+            # Clean up the session and context so we force a new connection
+            self._sessions = []
+            self._successful_connections = 0
+            self._initialized = False
+            self._connection_task = None
+
+        if self._initialized:
+            return
+
+        try:
+            await self._connect()
+        except (RuntimeError, BaseException) as e:
+            log_error(f"Failed to connect to {str(self)}: {e}")
+
+    @classmethod
+    async def create_and_connect(
+        cls,
+        commands: Optional[List[str]] = None,
+        urls: Optional[List[str]] = None,
+        urls_transports: Optional[List[Literal["sse", "streamable-http"]]] = None,
+        *,
+        env: Optional[dict[str, str]] = None,
+        server_params_list: Optional[
+            List[Union[SSEClientParams, StdioServerParameters, StreamableHTTPClientParams]]
+        ] = None,
+        timeout_seconds: int = 5,
+        client=None,
+        include_tools: Optional[list[str]] = None,
+        exclude_tools: Optional[list[str]] = None,
+        refresh_connection: bool = False,
+        **kwargs,
+    ) -> "MultiMCPTools":
+        """Initialize a MultiMCPTools instance and connect to the MCP servers"""
+        instance = cls(
+            commands=commands,
+            urls=urls,
+            urls_transports=urls_transports,
+            env=env,
+            server_params_list=server_params_list,
+            timeout_seconds=timeout_seconds,
+            client=client,
+            include_tools=include_tools,
+            exclude_tools=exclude_tools,
+            refresh_connection=refresh_connection,
+            **kwargs,
+        )
+
+        await instance._connect()
+        return instance
+
+    async def _connect(self) -> None:
+        """Connects to the MCP servers and initializes the tools"""
+        if self._initialized:
+            return
+
+        server_connection_errors = []
+
+        for server_params in self.server_params_list:
+            try:
+                # Handle stdio connections
+                if isinstance(server_params, StdioServerParameters):
+                    stdio_transport = await self._async_exit_stack.enter_async_context(stdio_client(server_params))
+                    read, write = stdio_transport
+                    session = await self._async_exit_stack.enter_async_context(
+                        ClientSession(read, write, read_timeout_seconds=timedelta(seconds=self.timeout_seconds))
+                    )
+                    await self.initialize(session)
+                    self._successful_connections += 1
+
+                # Handle SSE connections
+                elif isinstance(server_params, SSEClientParams):
+                    client_connection = await self._async_exit_stack.enter_async_context(
+                        sse_client(**asdict(server_params))
+                    )
+                    read, write = client_connection
+                    session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
+                    await self.initialize(session)
+                    self._successful_connections += 1
+
+                # Handle Streamable HTTP connections
+                elif isinstance(server_params, StreamableHTTPClientParams):
+                    client_connection = await self._async_exit_stack.enter_async_context(
+                        streamablehttp_client(**asdict(server_params))
+                    )
+                    read, write = client_connection[0:2]
+                    session = await self._async_exit_stack.enter_async_context(ClientSession(read, write))
+                    await self.initialize(session)
+                    self._successful_connections += 1
+
+            except Exception as e:
+                if not self.allow_partial_failure:
+                    raise ValueError(f"MCP connection failed: {e}")
+
+                log_error(f"Failed to initialize MCP server with params {server_params}: {e}")
+                server_connection_errors.append(str(e))
+                continue
+
+        if self._successful_connections > 0:
+            await self.build_tools()
+
+        if self._successful_connections == 0 and server_connection_errors:
+            raise ValueError(f"All MCP connections failed: {server_connection_errors}")
+
+        if not self._initialized and self._successful_connections > 0:
+            self._initialized = True
+
+    async def close(self) -> None:
+        """Close the MCP connections and clean up resources"""
+        if not self._initialized:
+            return
+
+        try:
+            await self._async_exit_stack.aclose()
+            self._sessions = []
+            self._successful_connections = 0
+
+        except (RuntimeError, BaseException) as e:
+            log_error(f"Failed to close MCP connections: {e}")
+
+        self._initialized = False
+
+    async def __aenter__(self) -> "MultiMCPTools":
+        """Enter the async context manager."""
+        try:
+            await self._connect()
+        except (RuntimeError, BaseException) as e:
+            log_error(f"Failed to connect to {str(self)}: {e}")
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Union[type[BaseException], None],
+        exc_val: Union[BaseException, None],
+        exc_tb: Union[TracebackType, None],
+    ):
+        """Exit the async context manager."""
+        await self._async_exit_stack.aclose()
+        self._initialized = False
+        self._successful_connections = 0
+
+    async def build_tools(self) -> None:
+        for session in self._sessions:
+            # Get the list of tools from the MCP server
+            available_tools = await session.list_tools()
+
+            # Filter tools based on include/exclude lists
+            filtered_tools = []
+            for tool in available_tools.tools:
+                if self.exclude_tools and tool.name in self.exclude_tools:
+                    continue
+                if self.include_tools is None or tool.name in self.include_tools:
+                    filtered_tools.append(tool)
+
+            # Register the tools with the toolkit
+            for tool in filtered_tools:
+                try:
+                    # Get an entrypoint for the tool
+                    entrypoint = get_entrypoint_for_tool(tool, session)
+
+                    # Create a Function for the tool
+                    f = Function(
+                        name=tool.name,
+                        description=tool.description,
+                        parameters=tool.inputSchema,
+                        entrypoint=entrypoint,
+                        # Set skip_entrypoint_processing to True to avoid processing the entrypoint
+                        skip_entrypoint_processing=True,
+                    )
+
+                    # Register the Function with the toolkit
+                    self.functions[f.name] = f
+                    log_debug(f"Function: {f.name} registered with {self.name}")
+                except Exception as e:
+                    log_error(f"Failed to register tool {tool.name}: {e}")
+                    raise
+
+    async def initialize(self, session: ClientSession) -> None:
+        """Initialize the MCP toolkit by getting available tools from the MCP server"""
+
+        try:
+            # Initialize the session if not already initialized
+            await session.initialize()
+
+            self._sessions.append(session)
+            self._initialized = True
+        except Exception as e:
+            log_error(f"Failed to get MCP tools: {e}")
+            raise

agno/tools/mcp/params.py

@@ -0,0 +1,24 @@
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any, Dict, Optional
+
+
+@dataclass
+class SSEClientParams:
+    """Parameters for SSE client connection."""
+
+    url: str
+    headers: Optional[Dict[str, Any]] = None
+    timeout: Optional[float] = 5
+    sse_read_timeout: Optional[float] = 60 * 5
+
+
+@dataclass
+class StreamableHTTPClientParams:
+    """Parameters for Streamable HTTP client connection."""
+
+    url: str
+    headers: Optional[Dict[str, Any]] = None
+    timeout: Optional[timedelta] = timedelta(seconds=30)
+    sse_read_timeout: Optional[timedelta] = timedelta(seconds=60 * 5)
+    terminate_on_close: Optional[bool] = None

agno/tools/mcp_toolbox.py

@@ -0,0 +1,284 @@
+from typing import Any, Callable, Dict, List, Literal, Optional, Union
+from warnings import warn
+
+from agno.tools.function import Function
+from agno.tools.mcp import MCPTools
+from agno.utils.log import logger
+
+try:
+    from toolbox_core import ToolboxClient  # type: ignore
+except ImportError:
+    raise ImportError("`toolbox_core` not installed. Please install using `pip install -U toolbox-core`.")
+
+
+class MCPToolsMeta(type):
+    """Metaclass for MCPTools to ensure proper initialization with AgentOS"""
+
+    @property
+    def __name__(cls):
+        return "MCPTools"
+
+
+class MCPToolbox(MCPTools, metaclass=MCPToolsMeta):
+    """
+    A toolkit that combines MCPTools server connectivity with MCP Toolbox for Databases client (toolbox-core).
+
+    MCPToolbox connects to an MCP Toolbox server and registers all available tools, then uses
+    toolbox-core to filter those tools by toolset or tool name. This enables agents to
+    receive only the specific tools they need while maintaining full MCP execution capabilities.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        toolsets: Optional[List[str]] = None,
+        tool_name: Optional[str] = None,
+        headers: Optional[Dict[str, Any]] = None,
+        transport: Literal["stdio", "sse", "streamable-http"] = "streamable-http",
+        append_mcp_to_url: bool = True,
+        **kwargs,
+    ):
+        """Initialize MCPToolbox with filtering capabilities.
+
+        Args:
+            url (str): Base URL for the toolbox service.
+            toolsets (Optional[List[str]], optional): List of toolset names to filter tools by. Defaults to None.
+            tool_name (Optional[str], optional): Single tool name to load. Defaults to None.
+            headers (Optional[Dict[str, Any]], optional): Headers for toolbox-core client requests. Defaults to None.
+            transport (Literal["stdio", "sse", "streamable-http"], optional): MCP transport protocol. Defaults to "streamable-http".
+            append_mcp_to_url (bool, optional): Whether to append "/mcp" to the URL if it doesn't end with it. Defaults to True.
+
+        """
+        if append_mcp_to_url and not url.endswith("/mcp"):
+            url = url + "/mcp"
+
+        super().__init__(url=url, transport=transport, **kwargs)
+
+        self.name = "toolbox_client"
+        self.toolbox_url = url
+        self.toolsets = toolsets
+        self.tool_name = tool_name
+        self.headers = headers
+        self._core_client_initialized = False
+
+        # Validate that only one of toolsets or tool_name is provided
+        filter_params = [toolsets, tool_name]
+        non_none_params = [p for p in filter_params if p is not None]
+        if len(non_none_params) > 1:
+            raise ValueError("Only one of toolsets or tool_name can be specified")
+
+    async def connect(self):
+        """Initialize MCPToolbox instance and connect to the MCP server."""
+        # First, connect to MCP server and load all available tools
+        await super().connect()
+
+        if self._core_client_initialized:
+            return
+
+        # Then, connect to the ToolboxClient and filter tools based on toolsets or tool_name
+        await self._connect_toolbox_client()
+
+    async def _connect_toolbox_client(self):
+        try:
+            if self.toolsets is not None or self.tool_name is not None:
+                self.__core_client = ToolboxClient(
+                    url=self.toolbox_url,
+                    client_headers=self.headers,
+                )
+                self._core_client_initialized = True
+
+                if self.toolsets is not None:
+                    # Load multiple toolsets
+                    all_functions = await self.load_multiple_toolsets(toolset_names=self.toolsets)
+                    # Replace functions dict with filtered subset
+                    filtered_functions = {func.name: func for func in all_functions}
+                    self.functions = filtered_functions
+                elif self.tool_name is not None:
+                    tool = await self.load_tool(tool_name=self.tool_name)
+                    # Replace functions dict with just this single tool
+                    self.functions = {tool.name: tool}
+        except Exception as e:
+            raise RuntimeError(f"Failed to connect to ToolboxClient: {e}") from e
+
+    def _handle_auth_params(
+        self,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
+        auth_headers: Optional[dict[str, Callable[[], str]]] = None,
+    ):
+        """handle authentication parameters for toolbox-core client"""
+        if auth_tokens:
+            if auth_token_getters:
+                warn(
+                    "Both `auth_token_getters` and `auth_tokens` are provided. `auth_tokens` is deprecated, and `auth_token_getters` will be used.",
+                    DeprecationWarning,
+                )
+            else:
+                warn(
+                    "Argument `auth_tokens` is deprecated. Use `auth_token_getters` instead.",
+                    DeprecationWarning,
+                )
+                auth_token_getters = auth_tokens
+
+        if auth_headers:
+            if auth_token_getters:
+                warn(
+                    "Both `auth_token_getters` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_token_getters` will be used.",
+                    DeprecationWarning,
+                )
+            else:
+                warn(
+                    "Argument `auth_headers` is deprecated. Use `auth_token_getters` instead.",
+                    DeprecationWarning,
+                )
+                auth_token_getters = auth_headers
+        return auth_token_getters
+
+    async def load_tool(
+        self,
+        tool_name: str,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
+        auth_headers: Optional[dict[str, Callable[[], str]]] = None,
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+    ) -> Function:
+        """Loads the tool with the given tool name from the Toolbox service.
+
+        Args:
+            tool_name (str): The name of the tool to load.
+            auth_token_getters (dict[str, Callable[[], str]], optional): A mapping of authentication source names to functions that retrieve ID tokens. Defaults to {}.
+            auth_tokens (Optional[dict[str, Callable[[], str]]], optional): Deprecated. Use `auth_token_getters` instead.
+            auth_headers (Optional[dict[str, Callable[[], str]]], optional): Deprecated. Use `auth_token_getters` instead.
+            bound_params (dict[str, Union[Any, Callable[[], Any]]], optional): A mapping of parameter names to their bound values. Defaults to {}.
+
+        Raises:
+            RuntimeError: If the tool is not found in the MCP functions registry.
+
+        Returns:
+            Function: The loaded tool function.
+        """
+        auth_token_getters = self._handle_auth_params(
+            auth_token_getters=auth_token_getters,
+            auth_tokens=auth_tokens,
+            auth_headers=auth_headers,
+        )
+
+        core_sync_tool = await self.__core_client.load_tool(
+            name=tool_name,
+            auth_token_getters=auth_token_getters,
+            bound_params=bound_params,
+        )
+        # Return the Function object from our MCP functions registry
+        if core_sync_tool._name in self.functions:
+            return self.functions[core_sync_tool._name]
+        else:
+            raise RuntimeError(f"Tool '{tool_name}' was not found in MCP functions registry")
+
+    async def load_toolset(
+        self,
+        toolset_name: Optional[str] = None,
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        auth_tokens: Optional[dict[str, Callable[[], str]]] = None,
+        auth_headers: Optional[dict[str, Callable[[], str]]] = None,
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+        strict: bool = False,
+    ) -> List[Function]:
+        """Loads tools from the configured toolset.
+
+        Args:
+            toolset_name (Optional[str], optional): The name of the toolset to load. Defaults to None.
+            auth_token_getters (dict[str, Callable[[], str]], optional): A mapping of authentication source names to functions that retrieve ID tokens. Defaults to {}.
+            auth_tokens (Optional[dict[str, Callable[[], str]]], optional): Deprecated. Use `auth_token_getters` instead.
+            auth_headers (Optional[dict[str, Callable[[], str]]], optional): Deprecated. Use `auth_token_getters` instead.
+            bound_params (dict[str, Union[Any, Callable[[], Any]]], optional): A mapping of parameter names to their bound values. Defaults to {}.
+            strict (bool, optional): If True, raises an error if *any* loaded tool instance fails
+                to utilize all of the given parameters or auth tokens. (if any
+                provided). If False (default), raises an error only if a
+                user-provided parameter or auth token cannot be applied to *any*
+                loaded tool across the set.
+
+        Returns:
+            List[Function]: A list of all tools loaded from the Toolbox.
+        """
+        auth_token_getters = self._handle_auth_params(
+            auth_token_getters=auth_token_getters,
+            auth_tokens=auth_tokens,
+            auth_headers=auth_headers,
+        )
+
+        core_sync_tools = await self.__core_client.load_toolset(
+            name=toolset_name,
+            auth_token_getters=auth_token_getters,
+            bound_params=bound_params,
+            strict=strict,
+        )
+
+        tools = []
+        for core_sync_tool in core_sync_tools:
+            if core_sync_tool._name in self.functions:
+                tools.append(self.functions[core_sync_tool._name])
+            else:
+                logger.debug(f"Tool '{core_sync_tool._name}' from toolset '{toolset_name}' not available in MCP server")
+        return tools
+
+    async def load_multiple_toolsets(
+        self,
+        toolset_names: List[str],
+        auth_token_getters: dict[str, Callable[[], str]] = {},
+        bound_params: dict[str, Union[Any, Callable[[], Any]]] = {},
+        strict: bool = False,
+    ) -> List[Function]:
+        """Load tools from multiple toolsets.
+
+        Args:
+            toolset_names (List[str]): A list of toolset names to load.
+            auth_token_getters (dict[str, Callable[[], str]], optional): A mapping of authentication source names to functions that retrieve ID tokens. Defaults to {}.
+            bound_params (dict[str, Union[Any, Callable[[], Any]]], optional): A mapping of parameter names to their bound values. Defaults to {}.
+            strict (bool, optional): If True, raises an error if *any* loaded tool instance fails to utilize all of the given parameters or auth tokens. Defaults to False.
+
+        Returns:
+            List[Function]: A list of all tools loaded from the specified toolsets.
+        """
+        all_tools = []
+        for toolset_name in toolset_names:
+            tools = await self.load_toolset(
+                toolset_name=toolset_name,
+                auth_token_getters=auth_token_getters,
+                bound_params=bound_params,
+                strict=strict,
+            )
+            all_tools.extend(tools)
+        return all_tools
+
+    async def close(self):
+        """Close the underlying asynchronous client."""
+        if self._core_client_initialized and hasattr(self, "_MCPToolbox__core_client"):
+            await self.__core_client.close()
+        await super().close()
+
+    async def load_toolset_safe(self, toolset_name: str) -> List[str]:
+        """Safely load a toolset and return tool names."""
+        try:
+            tools = await self.load_toolset(toolset_name)
+            return [tool.name for tool in tools]
+        except Exception as e:
+            raise RuntimeError(f"Failed to load toolset '{toolset_name}': {e}") from e
+
+    def get_client(self) -> ToolboxClient:
+        """Get the underlying ToolboxClient."""
+        if not self._core_client_initialized:
+            raise RuntimeError("ToolboxClient not initialized. Call connect() first.")
+        return self.__core_client
+
+    async def __aenter__(self):
+        """Initialize the direct toolbox client."""
+        await super().__aenter__()
+        await self.connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        """Clean up the toolbox client."""
+        # Close ToolboxClient first, then MCP client
+        if self._core_client_initialized and hasattr(self, "_MCPToolbox__core_client"):
+            await self.__core_client.close()
+        await super().__aexit__(exc_type, exc_val, exc_tb)