Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add MCP adapters to autogen-ext #5251

Merged
merged 17 commits into from
Feb 9, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions python/packages/autogen-core/docs/src/reference/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ python/autogen_ext.models.replay
python/autogen_ext.models.azure
python/autogen_ext.models.semantic_kernel
python/autogen_ext.tools.langchain
python/autogen_ext.tools.mcp
python/autogen_ext.tools.graphrag
python/autogen_ext.tools.code_execution
python/autogen_ext.tools.semantic_kernel
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
autogen\_ext.tools.mcp
======================


.. automodule:: autogen_ext.tools.mcp
:members:
:undoc-members:
:show-inheritance:
5 changes: 5 additions & 0 deletions python/packages/autogen-ext/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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"]

Expand Down
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import json
from typing import Any, Literal, Mapping, Optional, Sequence
import warnings
from typing import Any, Literal, Mapping, Optional, Sequence

from autogen_core import FunctionCall
from autogen_core._cancellation_token import CancellationToken
Expand Down
13 changes: 13 additions & 0 deletions python/packages/autogen-ext/src/autogen_ext/tools/mcp/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,13 @@
from ._config import McpServerParams, SseServerParams, StdioServerParams
from ._factory import mcp_server_tools
from ._sse import SseMcpToolAdapter
from ._stdio import StdioMcpToolAdapter

__all__ = [
"StdioMcpToolAdapter",
"StdioServerParams",
"SseMcpToolAdapter",
"SseServerParams",
"McpServerParams",
"mcp_server_tools",
]
101 changes: 101 additions & 0 deletions python/packages/autogen-ext/src/autogen_ext/tools/mcp/_base.py
Original file line number Diff line number Diff line change
@@ -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)
22 changes: 22 additions & 0 deletions python/packages/autogen-ext/src/autogen_ext/tools/mcp/_config.py
Original file line number Diff line number Diff line change
@@ -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
134 changes: 134 additions & 0 deletions python/packages/autogen-ext/src/autogen_ext/tools/mcp/_factory.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,134 @@
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:

**Local file system MCP service over standard I/O example:**

Install the filesystem server package from npm (requires Node.js 16+ and npm).

.. code-block:: bash

npm install -g @modelcontextprotocol/server-filesystem

Create an agent that can use all tools from the 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, # type: ignore
)

# 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())

**Local fetch MCP service over standard I/O example:**

Install the `mcp-server-fetch` package.

.. code-block:: bash

pip install mcp-server-fetch

Create an agent that can use the `fetch` tool from the local MCP server.

.. code-block:: python

import asyncio

from autogen_agentchat.agents import AssistantAgent
from autogen_ext.models.openai import OpenAIChatCompletionClient
from autogen_ext.tools.mcp import StdioServerParams, mcp_server_tools


async def main() -> None:
# Get the fetch tool from mcp-server-fetch.
fetch_mcp_server = StdioServerParams(command="uvx", args=["mcp-server-fetch"])
tools = await mcp_server_tools(fetch_mcp_server)

# Create an agent that can use the fetch tool.
model_client = OpenAIChatCompletionClient(model="gpt-4o")
agent = AssistantAgent(name="fetcher", model_client=model_client, tools=tools, reflect_on_tool_use=True) # type: ignore

# Let the agent fetch the content of a URL and summarize it.
result = await agent.run(task="Summarize the content of https://en.wikipedia.org/wiki/Seattle")
print(result.messages[-1].content)


asyncio.run(main())


**Remote MCP service over SSE example:**

.. code-block:: python

from autogen_ext.tools.mcp 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) # type: ignore

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]
raise ValueError(f"Unsupported server params type: {type(server_params)}")
23 changes: 23 additions & 0 deletions python/packages/autogen-ext/src/autogen_ext/tools/mcp/_session.py
Original file line number Diff line number Diff line change
@@ -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(
ekzhu marked this conversation as resolved.
Show resolved Hide resolved
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
Loading
Loading