Add MCPTools for loading a subset of tools

Author: RussellLuo

coagent/agents/react_agent/__init__.py

@@ -1,6 +1,7 @@
 # ruff: noqa: F401
 from .agent import ReActAgent
 from .context import RunContext
+from .mcp import MCPTools
 from .messages import InputMessage, InputHistory, OutputMessage
 from .types import (
     MessageOutputItem,

coagent/agents/react_agent/agent.py

@@ -38,6 +38,7 @@
     MessageOutputItem,
     StreamEvent,
 )
+from .mcp import MCPTools
 from .messages import InputHistory, OutputMessage
 from ..model import default_model, Model
 from .converter import Converter
@@ -101,19 +102,19 @@ def model_settings(self) -> ModelSettings:
     async def started(self) -> None:
         final_tools: list[Callable] = []
 
-        mcp_clients: list[mcputil.Client] = []
+        mcp_loaders: list[MCPTools] = []
         subagents: list[Subagent] = []
 
         for tool in self._tools:
-            if isinstance(tool, mcputil.Client):
-                mcp_clients.append(tool)
+            if isinstance(tool, MCPTools):
+                mcp_loaders.append(tool)
             elif isinstance(tool, Subagent):
                 subagents.append(tool)
             else:
                 final_tools.append(tool)
 
-        # Load tools from MCP clients
-        mcp_tools = await self._load_mcp_tools(mcp_clients)
+        # Load tools from MCP loaders
+        mcp_tools = await self._load_mcp_tools(mcp_loaders)
         final_tools.extend(mcp_tools)
 
         # Create tools from subagents
@@ -124,22 +125,20 @@ async def started(self) -> None:
 
         self._tools = final_tools
 
-    async def _load_mcp_tools(
-        self, mcp_clients: list[mcputil.Client]
-    ) -> list[Callable]:
+    async def _load_mcp_tools(self, mcp_loaders: list[MCPTools]) -> list[mcputil.Tool]:
         """Load tools from all MCP clients concurrently."""
 
-        async def get_client_tools(client):
+        async def load_tools(loader: MCPTools) -> list[mcputil.Tool]:
             try:
-                return await client.get_tools()
+                return await loader.load()
             except Exception as exc:
                 # Log error but continue with empty tools list
                 logger.error(f"Error getting tools from MCP client: {exc}")
                 return []
 
         # Fetch all tools concurrently and flatten the results
         all_mcp_tools = await asyncio.gather(
-            *[get_client_tools(client) for client in mcp_clients]
+            *[load_tools(loader) for loader in mcp_loaders]
         )
 
         return [tool for sublist in all_mcp_tools for tool in sublist]

coagent/agents/react_agent/mcp.py

@@ -0,0 +1,26 @@
+from dataclasses import dataclass
+
+import mcputil
+
+
+@dataclass
+class MCPTools:
+    """A loader for MCP tools. This class can be used to load a subset of tools from an MCP server."""
+
+    client: mcputil.Client
+    """The MCP client to use for loading the tools."""
+
+    include: list[str] | None = None
+    """The list of tool names to include. If None, all tools are included."""
+
+    exclude: list[str] | None = None
+    """The list of tool names to exclude. If None, no tools are excluded."""
+
+    async def load(self) -> list[mcputil.Tool]:
+        return await self.client.get_tools(include=self.include, exclude=self.exclude)
+
+    async def connect(self) -> None:
+        await self.client.connect()
+
+    async def close(self) -> None:
+        await self.client.close()

examples/react-mcp/agent.py

@@ -4,6 +4,7 @@
     InputMessage,
     InputHistory,
     OutputMessage,
+    MCPTools,
     MessageOutputItem,
     ToolCallItem,
     ToolCallOutputItem,
@@ -32,7 +33,7 @@ async def main():
                     ReActAgent,
                     name="reporter",
                     system="You are a helpful weather reporter",
-                    tools=[mcp_client],
+                    tools=[MCPTools(mcp_client)],
                 ),
             )
             await runtime.register(reporter)

poetry.lock

@@ -38,7 +38,7 @@ litellm = "1.60.4"
 mcp = ">=1.19.0"
 jinja2 = "3.1.5"
 aiorwlock = ">=1.5.0"
-mcputil = "0.2.0"
+mcputil = "0.3.0"
 
 # A list of optional dependencies, which are included in the
 # below `extras`.