Add MCP adapters to autogen-ext

python/packages/autogen-ext/pyproject.toml

@@ -112,6 +112,11 @@ semantic-kernel-all = [
 
 rich = ["rich>=13.9.4"]
 
+mcp = [
+    "mcp>=1.1.3",
+    "json-schema-to-pydantic>=0.2.2"
+]
+
 [tool.hatch.build.targets.wheel]
 packages = ["src/autogen_ext"]
 

python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py

@@ -0,0 +1,15 @@
+from ._config import McpServerParams, SseServerParams, StdioServerParams
+from ._factory import mcp_server_tools
+from ._sse import SseMcpToolAdapter, SseMcpToolAdapterConfig
+from ._stdio import StdioMcpToolAdapter, StdioMcpToolAdapterConfig
+
+__all__ = [
+    "StdioMcpToolAdapter",
+    "StdioMcpToolAdapterConfig",
+    "StdioServerParams",
+    "SseMcpToolAdapter",
+    "SseMcpToolAdapterConfig",
+    "SseServerParams",
+    "McpServerParams",
+    "mcp_server_tools",
+]

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py

@@ -0,0 +1,101 @@
+from abc import ABC
+from typing import Any, Generic, Type, TypeVar
+
+from autogen_core import CancellationToken
+from autogen_core.tools import BaseTool
+from json_schema_to_pydantic import create_model
+from mcp import Tool
+from pydantic import BaseModel
+
+from ._config import McpServerParams
+from ._session import create_mcp_server_session
+
+TServerParams = TypeVar("TServerParams", bound=McpServerParams)
+
+
+class McpToolAdapter(BaseTool[BaseModel, Any], ABC, Generic[TServerParams]):
+    """
+    Base adapter class for MCP tools to make them compatible with AutoGen.
+
+    Args:
+        server_params (TServerParams): Parameters for the MCP server connection.
+        tool (Tool): The MCP tool to wrap.
+    """
+
+    component_type = "tool"
+
+    def __init__(self, server_params: TServerParams, tool: Tool) -> None:
+        self._tool = tool
+        self._server_params = server_params
+
+        # Extract name and description
+        name = tool.name
+        description = tool.description or ""
+
+        # Create the input model from the tool's schema
+        input_model = create_model(tool.inputSchema)
+
+        # Use Any as return type since MCP tool returns can vary
+        return_type: Type[Any] = object
+
+        super().__init__(input_model, return_type, name, description)
+
+    async def run(self, args: BaseModel, cancellation_token: CancellationToken) -> Any:
+        """
+        Run the MCP tool with the provided arguments.
+
+        Args:
+            args (BaseModel): The arguments to pass to the tool.
+            cancellation_token (CancellationToken): Token to signal cancellation.
+
+        Returns:
+            Any: The result of the tool execution.
+
+        Raises:
+            Exception: If the operation is cancelled or the tool execution fails.
+        """
+        kwargs = args.model_dump()
+
+        try:
+            async with create_mcp_server_session(self._server_params) as session:
+                await session.initialize()
+
+                if cancellation_token.is_cancelled():
+                    raise Exception("Operation cancelled")
+
+                result = await session.call_tool(self._tool.name, kwargs)  # type: ignore
+
+                if result.isError:
+                    raise Exception(f"MCP tool execution failed: {result.content}")
+
+                return result.content
+        except Exception as e:
+            raise Exception(str(e)) from e
+
+    @classmethod
+    async def from_server_params(cls, server_params: TServerParams, tool_name: str) -> "McpToolAdapter[TServerParams]":
+        """
+        Create an instance of McpToolAdapter from server parameters and tool name.
+
+        Args:
+            server_params (TServerParams): Parameters for the MCP server connection.
+            tool_name (str): The name of the tool to wrap.
+
+        Returns:
+            McpToolAdapter[TServerParams]: An instance of McpToolAdapter.
+
+        Raises:
+            ValueError: If the tool with the specified name is not found.
+        """
+        async with create_mcp_server_session(server_params) as session:
+            await session.initialize()
+
+            tools_response = await session.list_tools()
+            matching_tool = next((t for t in tools_response.tools if t.name == tool_name), None)
+
+            if matching_tool is None:
+                raise ValueError(
+                    f"Tool '{tool_name}' not found, available tools: {', '.join([t.name for t in tools_response.tools])}"
+                )
+
+        return cls(server_params=server_params, tool=matching_tool)

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py

@@ -0,0 +1,22 @@
+from typing import Any, TypeAlias
+
+from mcp import StdioServerParameters
+from pydantic import BaseModel
+
+
+class StdioServerParams(StdioServerParameters):
+    """Parameters for connecting to an MCP server over STDIO."""
+
+    pass
+
+
+class SseServerParams(BaseModel):
+    """Parameters for connecting to an MCP server over SSE."""
+
+    url: str
+    headers: dict[str, Any] | None = None
+    timeout: float = 5
+    sse_read_timeout: float = 60 * 5
+
+
+McpServerParams: TypeAlias = StdioServerParams | SseServerParams

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py

@@ -0,0 +1,94 @@
+from ._config import McpServerParams, SseServerParams, StdioServerParams
+from ._session import create_mcp_server_session
+from ._sse import SseMcpToolAdapter
+from ._stdio import StdioMcpToolAdapter
+
+
+async def mcp_server_tools(
+    server_params: McpServerParams,
+) -> list[StdioMcpToolAdapter | SseMcpToolAdapter]:
+    """Creates a list of MCP tool adapters that can be used with AutoGen agents.
+
+    This factory function connects to an MCP server and returns adapters for all available tools.
+    The adapters can be directly assigned to an AutoGen agent's tools list.
+
+    Args:
+        server_params (McpServerParams): Connection parameters for the MCP server.
+            Can be either StdioServerParams for command-line tools or
+            SseServerParams for HTTP/SSE services.
+
+    Returns:
+        list[StdioMcpToolAdapter | SseMcpToolAdapter]: A list of tool adapters ready to use
+            with AutoGen agents.
+
+    Examples:
+        Create an agent that can use all tools from a local filesystem MCP server:
+
+        .. code-block:: python
+
+            import asyncio
+            from pathlib import Path
+            from autogen_ext.models.openai import OpenAIChatCompletionClient
+            from autogen_ext.tools.mcp import StdioServerParams, mcp_server_tools
+            from autogen_agentchat.agents import AssistantAgent
+            from autogen_core import CancellationToken
+
+
+            async def main() -> None:
+                # Setup server params for local filesystem access
+                desktop = str(Path.home() / "Desktop")
+                server_params = StdioServerParams(
+                    command="npx.cmd", args=["-y", "@modelcontextprotocol/server-filesystem", desktop]
+                )
+
+                # Get all available tools from the server
+                tools = await mcp_server_tools(server_params)
+
+                # Create an agent that can use all the tools
+                agent = AssistantAgent(
+                    name="file_manager",
+                    model_client=OpenAIChatCompletionClient(model="gpt-4"),
+                    tools=tools,  # Assign all tools to the agent
+                )
+
+                # The agent can now use any of the filesystem tools
+                await agent.run(
+                    task="Create a file called test.txt with some content", cancellation_token=CancellationToken()
+                )
+
+
+            if __name__ == "__main__":
+                asyncio.run(main())
+
+        Or connect to a remote MCP service over SSE:
+
+        .. code-block:: python
+
+            from autogen_ext_mcp.tools import SseServerParams, mcp_server_tools
+
+
+            async def main() -> None:
+                # Setup server params for remote service
+                server_params = SseServerParams(
+                    url="https://api.example.com/mcp", headers={"Authorization": "Bearer token"}
+                )
+
+                # Get all available tools
+                tools = await mcp_server_tools(server_params)
+
+                # Create an agent with all tools
+                agent = AssistantAgent(
+                    name="tool_user", model_client=OpenAIChatCompletionClient(model="gpt-4"), tools=tools
+                )
+
+    For more examples and detailed usage, see the samples directory in the package repository.
+    """
+    async with create_mcp_server_session(server_params) as session:
+        await session.initialize()
+
+        tools = await session.list_tools()
+
+    if isinstance(server_params, StdioServerParams):
+        return [StdioMcpToolAdapter(server_params=server_params, tool=tool) for tool in tools.tools]
+    elif isinstance(server_params, SseServerParams):
+        return [SseMcpToolAdapter(server_params=server_params, tool=tool) for tool in tools.tools]

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py

@@ -0,0 +1,23 @@
+from contextlib import asynccontextmanager
+from typing import AsyncGenerator
+
+from mcp import ClientSession
+from mcp.client.sse import sse_client
+from mcp.client.stdio import stdio_client
+
+from ._config import McpServerParams, SseServerParams, StdioServerParams
+
+
+@asynccontextmanager
+async def create_mcp_server_session(
+    server_params: McpServerParams,
+) -> AsyncGenerator[ClientSession, None]:
+    """Create an MCP client session for the given server parameters."""
+    if isinstance(server_params, StdioServerParams):
+        async with stdio_client(server_params) as (read, write):
+            async with ClientSession(read_stream=read, write_stream=write) as session:
+                yield session
+    elif isinstance(server_params, SseServerParams):
+        async with sse_client(**server_params.model_dump()) as (read, write):
+            async with ClientSession(read_stream=read, write_stream=write) as session:
+                yield session

python/packages/autogen-ext/src/autogen_ext/tools/mcp/_sse.py

@@ -0,0 +1,107 @@
+from typing import Self
+
+from autogen_core import Component
+from mcp import Tool
+from pydantic import BaseModel
+
+from ._base import McpToolAdapter
+from ._config import SseServerParams
+
+
+class SseMcpToolAdapterConfig(BaseModel):
+    """Configuration for the MCP tool adapter."""
+
+    server_params: SseServerParams
+    tool: Tool
+
+
+class SseMcpToolAdapter(
+    McpToolAdapter[SseServerParams],
+    Component[SseMcpToolAdapterConfig],
+):
+    """
+    Allows you to wrap an MCP tool running over Server-Sent Events (SSE) and make it available to AutoGen.
+
+    This adapter enables using MCP-compatible tools that communicate over HTTP with SSE
+    with AutoGen agents. Common use cases include integrating with remote MCP services,
+    cloud-based tools, and web APIs that implement the Model Context Protocol (MCP).
+
+    Args:
+        server_params (SseServerParameters): Parameters for the MCP server connection,
+            including URL, headers, and timeouts
+        tool (Tool): The MCP tool to wrap
+
+    Examples:
+        Use a remote translation service that implements MCP over SSE to create tools
+        that allow AutoGen agents to perform translations:
+
+        .. code-block:: python
+
+            import asyncio
+            from autogen_ext.models.openai import OpenAIChatCompletionClient
+            from autogen_ext.tools.mcp import SseMcpToolAdapter, SseServerParams
+            from autogen_agentchat.agents import AssistantAgent
+            from autogen_agentchat.ui import Console
+            from autogen_core import CancellationToken
+
+
+            async def main() -> None:
+                # Create server params for the remote MCP service
+                server_params = SseServerParams(
+                    url="https://api.example.com/mcp",
+                    headers={"Authorization": "Bearer your-api-key", "Content-Type": "application/json"},
+                    timeout=30,  # Connection timeout in seconds
+                )
+
+                # Get the translation tool from the server
+                adapter = await SseMcpToolAdapter.from_server_params(server_params, "translate")
+
+                # Create an agent that can use the translation tool
+                model_client = OpenAIChatCompletionClient(model="gpt-4")
+                agent = AssistantAgent(
+                    name="translator",
+                    model_client=model_client,
+                    tools=[adapter],
+                    system_message="You are a helpful translation assistant.",
+                )
+
+                # Let the agent translate some text
+                await Console(
+                    agent.run_stream(
+                        task="Translate 'Hello, how are you?' to Spanish", cancellation_token=CancellationToken()
+                    )
+                )
+
+
+            if __name__ == "__main__":
+                asyncio.run(main())
+
+    """
+
+    component_config_schema = SseMcpToolAdapterConfig
+    component_provider_override = "autogen-ext.tools.mcp.SseMcpToolAdapter"
+
+    def __init__(self, server_params: SseServerParams, tool: Tool) -> None:
+        super().__init__(server_params=server_params, tool=tool)
+
+    def _to_config(self) -> SseMcpToolAdapterConfig:
+        """
+        Convert the adapter to its configuration representation.
+
+        Returns:
+            SseMcpToolAdapterConfig: The configuration of the adapter.
+        """
+        return SseMcpToolAdapterConfig(server_params=self._server_params, tool=self._tool)
+
+    @classmethod
+    def _from_config(cls, config: SseMcpToolAdapterConfig) -> Self:
+        """
+        Create an instance of SseMcpToolAdapter from its configuration.
+
+        Args:
+            config (SseMcpToolAdapterConfig): The configuration of the adapter.
+
+        Returns:
+            SseMcpToolAdapter: An instance of SseMcpToolAdapter.
+        """
+        return cls(server_params=config.server_params, tool=config.tool)