Add MCPAgent

Author: RussellLuo

coagent/agents/__init__.py

@@ -1,6 +1,7 @@
 # ruff: noqa: F401
 from .chat_agent import ChatAgent, confirm, submit, RunContext, StreamChatAgent, tool
 from .dynamic_triage import DynamicTriage
+from .mcp_agent import MCPAgent
 from .messages import ChatHistory, ChatMessage
 from .model_client import ModelClient
 from .parallel import Aggregator, AggregationResult, Parallel

coagent/agents/aswarm/util.py

@@ -131,6 +131,14 @@ def greet(
     Then you will get a JSON schema with per-parameter descriptions.
     """
 
+    if hasattr(func, "__mcp_tool_schema__"):
+        # If the function already has a schema, return it.
+        # This is the case for tools used in MCPAgent.
+        return dict(
+            type="function",
+            function=func.__mcp_tool_schema__
+        )
+
     # Construct the pydantic mdoel for the _under_fn's function signature parameters.
     # 1. Get the function signature.
 

coagent/agents/chat_agent.py

@@ -234,6 +234,9 @@ def system(self) -> str:
     def client(self) -> ModelClient:
         return self._client
 
+    async def get_swarm_agent(self) -> SwarmAgent:
+        return self._swarm_agent
+
     async def agent(self, agent_type: str) -> AsyncIterator[ChatMessage]:
         """The candidate agent to delegate the conversation to."""
         async for chunk in StreamDelegate(self, agent_type).handle(self._history):
@@ -265,8 +268,10 @@ async def _handle_history(
         await self.update_user_confirmed(msg)
         await self.update_user_submitted(msg)
 
+        swarm_agent = await self.get_swarm_agent()
+
         response = self._swarm_client.run_and_stream(
-            agent=self._swarm_agent,
+            agent=swarm_agent,
             messages=[m.model_dump() for m in msg.messages],
             context_variables=msg.extensions,
         )

coagent/agents/mcp_agent.py

@@ -0,0 +1,114 @@
+from typing import Any, AsyncContextManager, Callable
+from urllib.parse import urljoin
+
+from mcp import ClientSession, Tool
+from mcp.types import ImageContent, TextContent
+from mcp.client.sse import sse_client
+import jsonschema
+
+from .aswarm import Agent as SwarmAgent
+from .chat_agent import ChatAgent, wrap_error
+from .model_client import default_model_client, ModelClient
+
+
+class MCPAgent(ChatAgent):
+    """An agent that can use tools provided by MCP (Model Context Protocol) servers."""
+
+    def __init__(
+        self,
+        system: str = "",
+        mcp_server_base_url: str = "",
+        client: ModelClient = default_model_client,
+    ) -> None:
+        super().__init__(system=system, client=client)
+
+        self._mcp_server_base_url: str = mcp_server_base_url
+        self._mcp_sse_client: AsyncContextManager[tuple] | None = None
+        self._mcp_client_session: ClientSession | None = None
+
+        self._mcp_swarm_agent: SwarmAgent | None = None
+
+    @property
+    def mcp_server_base_url(self) -> str:
+        if not self._mcp_server_base_url:
+            raise ValueError("MCP server base URL is empty")
+        return self._mcp_server_base_url
+
+    def make_tool(self, t: Tool) -> Callable:
+        async def tool(**kwargs) -> Any:
+            # Validate the input against the schema
+            jsonschema.validate(instance=kwargs, schema=t.inputSchema)
+            # Actually call the tool.
+            result = await self._mcp_client_session.call_tool(t.name, arguments=kwargs)
+            if not result.content:
+                return ""
+            content = result.content[0]
+
+            if result.isError:
+                raise ValueError(content.text)
+
+            match content:
+                case TextContent():
+                    return content.text
+                case ImageContent():
+                    return content.data
+                case _:  # EmbeddedResource() or other types
+                    return ""
+
+        tool.__name__ = t.name
+        tool.__doc__ = t.description
+
+        # Attach the schema and arguments to the tool.
+        tool.__mcp_tool_schema__ = dict(
+            name=t.name,
+            description=t.description,
+            parameters=t.inputSchema,
+        )
+        tool.__mcp_tool_args__ = t.inputSchema["properties"].keys()
+        return tool
+
+    async def get_tools(self) -> list[Callable]:
+        result = await self._mcp_client_session.list_tools()
+        tools = [self.make_tool(t) for t in result.tools]
+        return tools
+
+    async def get_swarm_agent(self) -> SwarmAgent:
+        if not self._mcp_swarm_agent:
+            tools = await self.get_tools()
+            self._mcp_swarm_agent = SwarmAgent(
+                name=self.name,
+                model=self.client.model,
+                instructions=self.system,
+                functions=[wrap_error(t) for t in tools],
+            )
+        return self._mcp_swarm_agent
+
+    async def started(self) -> None:
+        """
+        Combining `started` and `stopped` to achieve the following behavior:
+
+            async with sse_client(url=url) as (read, write):
+                async with ClientSession(read, write) as session:
+                    pass
+        """
+        url = urljoin(self.mcp_server_base_url, "sse")
+        self._mcp_sse_client = sse_client(url=url)
+        read, write = await self._mcp_sse_client.__aenter__()
+
+        self._mcp_client_session = ClientSession(read, write)
+        await self._mcp_client_session.__aenter__()
+
+        # Initialize the connection
+        await self._mcp_client_session.initialize()
+
+    async def stopped(self) -> None:
+        await self._mcp_client_session.__aexit__(None, None, None)
+        await self._mcp_sse_client.__aexit__(None, None, None)
+
+    async def _handle_data(self) -> None:
+        """Override the method to handle exceptions properly."""
+        try:
+            await super()._handle_data()
+        finally:
+            # Ensure the resources created in `started` are properly cleaned up.
+            await self.stopped()

coagent/core/util.py

@@ -43,6 +43,9 @@ def direct_values(self, prefix: str) -> list[Any]:
 
 
 def get_func_args(func) -> set[str]:
+    if hasattr(func, "__mcp_tool_args__"):
+        return set(func.__mcp_tool_args__)
+
     hints = get_type_hints(func)
     hints.pop("return", None)  # Ignore the return type.
     return set(hints.keys())

poetry.lock

@@ -30,6 +30,7 @@ blinker = "1.9.0"
 loguru = "0.7.3"
 jq = "1.8.0"
 litellm = "1.55.12"
+mcp = "1.2.0"
 
 [tool.pyright]
 # https://github.com/microsoft/pyright/blob/main/docs/configuration.md