mirror of
https://github.com/mcp-use/mcp-use.git
synced 2026-09-14 18:06:01 +08:00
Luigi Pederzani
f08e870e09
* chore: update imports * refactor(mcpagent): update agent creation logic and imports - Replaced `create_tool_calling_agent` with `create_agent` for improved agent instantiation. - Simplified the `_create_agent` method by removing the `AgentExecutor` and directly returning the created agent. - Adjusted type annotations for `_agent_executor` to initialize as `None` instead of using a union type. * feat(mcpagent): migrate to LangChain 1.0.0 and simplify agent methods - Updated MCPAgent to utilize LangChain 1.0.0's `create_agent()` for improved agent instantiation. - Introduced new methods `astream_simplified()` and `run_v2()` that leverage the built-in `astream()` for internal agent loop handling. - Retained legacy methods `stream()` and `run()` for backward compatibility with manual execution. - Enhanced documentation to reflect changes and provide clearer usage examples. * feat(mcpagent): implement dynamic tool updates during agent execution - Enhanced the MCPAgent to support dynamic tool updates when using server_manager mode. - Implemented logic to check for tool changes before and during execution, allowing for immediate updates and safe restarts. - Introduced a maximum restart limit to prevent infinite loops during tool updates. - Updated documentation to reflect the new behavior and usage of dynamic tool management. * refactor(mcpagent): clean up imports and streamline logging - Removed unused imports and organized existing ones for better clarity in `mcpagent.py`. - Simplified logging statements for consistency and improved readability. - Enhanced the overall structure of the code by eliminating unnecessary whitespace. - Updated unit tests to reflect changes in imports and ensure compatibility with the refactored code. * feat(mcpagent): enhance agent execution with step tracking and middleware - Introduced `ModelCallLimitMiddleware` to enforce a maximum number of model calls during agent execution. - Updated the `_consume_and_return` method to return a tuple of the final result and the number of steps taken. - Enhanced telemetry tracking to include steps taken and execution time for better performance insights. - Simplified agent creation logic by integrating middleware directly into the `create_agent` function. - Improved logging for agent execution details, including steps and tool usage. * fix(mcpagent): add null check for node output messages extraction - Updated the MCPAgent to include a null check for `node_output` before attempting to access the "messages" key, preventing potential errors during message extraction. * minor * feat(mcpagent): filter out ToolMessage from history - Updated the MCPAgent to exclude ToolMessage instances from being added to the message history, ensuring that only relevant messages are recorded. - This change helps maintain a cleaner history by avoiding intermediate tool execution details. * fix(mcpagent): set max_steps assignment * ruff * docs * e2e tests for agent * key in tests * fix manager test and split ci tests * telemetry off for tests * fix server manager test * feat(mcpagent): delegate query streaming to remote agent in remote mode - recreated functionality to the MCPAgent to yield results from a remote agent when operating in remote mode, enhancing the agent's capability to handle queries efficiently. * fix: agent unit tests * Print inputs * rollback, wrong repo * fix streaming output * stricter streaming test * pass observability callbacks * fix tests mock --------- Co-authored-by: pietrozullo <pietro.zullo@gmail.com>
119 lines
4.2 KiB
Python
119 lines
4.2 KiB
Python
"""
|
|
End-to-end integration test for server manager.
|
|
|
|
Tests the agent with custom server manager for dynamic tool management.
|
|
"""
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from langchain_core.tools import BaseTool
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
from mcp_use import MCPAgent, MCPClient
|
|
from mcp_use.agents.managers.base import BaseServerManager
|
|
from mcp_use.logging import logger
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.integration
|
|
async def test_server_manager():
|
|
"""Test agent with custom server manager for dynamic tool management."""
|
|
server_path = Path(__file__).parent.parent / "servers_for_testing" / "simple_server.py"
|
|
|
|
config = {"mcpServers": {"simple": {"command": sys.executable, "args": [str(server_path), "--transport", "stdio"]}}}
|
|
|
|
class GreetingTool(BaseTool):
|
|
"""A simple greeting tool."""
|
|
|
|
name: str = "greet"
|
|
description: str = "Returns a greeting message"
|
|
|
|
def _run(self, name: str = "World") -> str:
|
|
return f"Hello, {name}!"
|
|
|
|
async def _arun(self, name: str = "World") -> str:
|
|
return f"Hello, {name}!"
|
|
|
|
class CustomServerManager(BaseServerManager):
|
|
"""Custom server manager with dynamic tool addition."""
|
|
|
|
def __init__(self):
|
|
self._tools: list[BaseTool] = []
|
|
self._initialized = False
|
|
|
|
def add_tool(self, tool: BaseTool):
|
|
self._tools.append(tool)
|
|
|
|
async def initialize(self) -> None:
|
|
# Create the get_greeting_tool dynamically inside the manager
|
|
class GetGreetingToolTool(BaseTool):
|
|
"""A tool that adds a greeting tool to the server manager."""
|
|
|
|
name: str = "get_greeting_tool"
|
|
description: str = "Adds a greeting tool to the server manager"
|
|
manager: CustomServerManager
|
|
|
|
def _run(self) -> str:
|
|
greeting_tool = GreetingTool()
|
|
self.manager.add_tool(greeting_tool)
|
|
return f"Added greeting tool to server manager. Total tools: {len(self.manager.tools)}"
|
|
|
|
async def _arun(self) -> str:
|
|
greeting_tool = GreetingTool()
|
|
self.manager.add_tool(greeting_tool)
|
|
return f"Added greeting tool to server manager. Total tools: {len(self.manager.tools)}"
|
|
|
|
get_greeting_tool = GetGreetingToolTool(manager=self)
|
|
self.add_tool(get_greeting_tool)
|
|
self._initialized = True
|
|
|
|
@property
|
|
def tools(self) -> list[BaseTool]:
|
|
return list(self._tools)
|
|
|
|
def has_tool_changes(self, current_tool_names: set[str]) -> bool:
|
|
new_tool_names = {tool.name for tool in self.tools}
|
|
return new_tool_names != current_tool_names
|
|
|
|
client = MCPClient.from_dict(config)
|
|
llm = ChatOpenAI(model="gpt-4o")
|
|
server_manager = CustomServerManager()
|
|
|
|
agent = MCPAgent(llm=llm, client=client, use_server_manager=True, server_manager=server_manager, max_steps=10)
|
|
|
|
try:
|
|
logger.info("\n" + "=" * 80)
|
|
logger.info("TEST: test_server_manager")
|
|
logger.info("=" * 80)
|
|
|
|
await agent.initialize()
|
|
|
|
logger.info(f"Initial server manager tools: {[t.name for t in server_manager.tools]}")
|
|
|
|
assert len(server_manager.tools) == 1
|
|
assert server_manager.tools[0].name == "get_greeting_tool"
|
|
|
|
query = "Call get_greeting_tool to add the greeting tool, then use the greet tool to say hello to Alice"
|
|
logger.info(f"Query: {query}")
|
|
|
|
result = await agent.run(query)
|
|
|
|
logger.info(f"Result: {result}")
|
|
logger.info(f"Tools used: {agent.tools_used_names}")
|
|
logger.info(f"Final server manager tools: {[t.name for t in server_manager.tools]}")
|
|
logger.info("=" * 80 + "\n")
|
|
|
|
# Assert the agent used the get_greeting_tool to add the greeting tool
|
|
assert "get_greeting_tool" in agent.tools_used_names
|
|
# Assert the server manager has at least 2 tools
|
|
assert len(server_manager.tools) >= 2
|
|
# Assert the server manager has the greet tool
|
|
assert "greet" in [t.name for t in server_manager.tools]
|
|
# Assert the gree tool was used by the agent
|
|
assert "greet" in agent.tools_used_names
|
|
|
|
finally:
|
|
await agent.close()
|