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/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