Skip to content

Commit

Permalink
Add MCPAgent
Browse files Browse the repository at this point in the history
  • Loading branch information
RussellLuo committed Jan 18, 2025
1 parent 9c11a73 commit 2f14a1d
Show file tree
Hide file tree
Showing 7 changed files with 193 additions and 2 deletions.
1 change: 1 addition & 0 deletions coagent/agents/__init__.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
8 changes: 8 additions & 0 deletions coagent/agents/aswarm/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
7 changes: 6 additions & 1 deletion coagent/agents/chat_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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,
)
Expand Down
3 changes: 3 additions & 0 deletions coagent/core/util.py
Original file line number Diff line number Diff line change
Expand Up @@ -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())
Expand Down
102 changes: 101 additions & 1 deletion poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
73 changes: 73 additions & 0 deletions uv.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

0 comments on commit 2f14a1d

Please sign in to comment.