mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-28 09:38:17 +00:00
Support to filter available MCP Tools (#345)
* feat: support to complex filter on ToolCollection * refactor: use proper tool collection methot to filter tool in CrewAiEnterpriseTools * feat: allow to filter available MCP tools
This commit is contained in:
@@ -19,6 +19,11 @@ def echo_server_script():
|
||||
"""Echo the input text"""
|
||||
return f"Echo: {text}"
|
||||
|
||||
@mcp.tool()
|
||||
def calc_tool(a: int, b: int) -> int:
|
||||
"""Calculate a + b"""
|
||||
return a + b
|
||||
|
||||
mcp.run()
|
||||
'''
|
||||
)
|
||||
@@ -37,6 +42,11 @@ def echo_server_sse_script():
|
||||
"""Echo the input text"""
|
||||
return f"Echo: {text}"
|
||||
|
||||
@mcp.tool()
|
||||
def calc_tool(a: int, b: int) -> int:
|
||||
"""Calculate a + b"""
|
||||
return a + b
|
||||
|
||||
mcp.run("sse")
|
||||
'''
|
||||
)
|
||||
@@ -69,16 +79,20 @@ def test_context_manager_syntax(echo_server_script):
|
||||
)
|
||||
with MCPServerAdapter(serverparams) as tools:
|
||||
assert isinstance(tools, ToolCollection)
|
||||
assert len(tools) == 1
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "echo_tool"
|
||||
assert tools[1].name == "calc_tool"
|
||||
assert tools[0].run(text="hello") == "Echo: hello"
|
||||
assert tools[1].run(a=5, b=3) == '8'
|
||||
|
||||
def test_context_manager_syntax_sse(echo_sse_server):
|
||||
sse_serverparams = echo_sse_server
|
||||
with MCPServerAdapter(sse_serverparams) as tools:
|
||||
assert len(tools) == 1
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "echo_tool"
|
||||
assert tools[1].name == "calc_tool"
|
||||
assert tools[0].run(text="hello") == "Echo: hello"
|
||||
assert tools[1].run(a=5, b=3) == '8'
|
||||
|
||||
def test_try_finally_syntax(echo_server_script):
|
||||
serverparams = StdioServerParameters(
|
||||
@@ -87,9 +101,11 @@ def test_try_finally_syntax(echo_server_script):
|
||||
try:
|
||||
mcp_server_adapter = MCPServerAdapter(serverparams)
|
||||
tools = mcp_server_adapter.tools
|
||||
assert len(tools) == 1
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "echo_tool"
|
||||
assert tools[1].name == "calc_tool"
|
||||
assert tools[0].run(text="hello") == "Echo: hello"
|
||||
assert tools[1].run(a=5, b=3) == '8'
|
||||
finally:
|
||||
mcp_server_adapter.stop()
|
||||
|
||||
@@ -98,8 +114,76 @@ def test_try_finally_syntax_sse(echo_sse_server):
|
||||
mcp_server_adapter = MCPServerAdapter(sse_serverparams)
|
||||
try:
|
||||
tools = mcp_server_adapter.tools
|
||||
assert len(tools) == 2
|
||||
assert tools[0].name == "echo_tool"
|
||||
assert tools[1].name == "calc_tool"
|
||||
assert tools[0].run(text="hello") == "Echo: hello"
|
||||
assert tools[1].run(a=5, b=3) == '8'
|
||||
finally:
|
||||
mcp_server_adapter.stop()
|
||||
|
||||
def test_context_manager_with_filtered_tools(echo_server_script):
|
||||
serverparams = StdioServerParameters(
|
||||
command="uv", args=["run", "python", "-c", echo_server_script]
|
||||
)
|
||||
# Only select the echo_tool
|
||||
with MCPServerAdapter(serverparams, "echo_tool") as tools:
|
||||
assert isinstance(tools, ToolCollection)
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "echo_tool"
|
||||
assert tools[0].run(text="hello") == "Echo: hello"
|
||||
# Check that calc_tool is not present
|
||||
with pytest.raises(IndexError):
|
||||
_ = tools[1]
|
||||
with pytest.raises(KeyError):
|
||||
_ = tools["calc_tool"]
|
||||
|
||||
def test_context_manager_sse_with_filtered_tools(echo_sse_server):
|
||||
sse_serverparams = echo_sse_server
|
||||
# Only select the calc_tool
|
||||
with MCPServerAdapter(sse_serverparams, "calc_tool") as tools:
|
||||
assert isinstance(tools, ToolCollection)
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "calc_tool"
|
||||
assert tools[0].run(a=10, b=5) == '15'
|
||||
# Check that echo_tool is not present
|
||||
with pytest.raises(IndexError):
|
||||
_ = tools[1]
|
||||
with pytest.raises(KeyError):
|
||||
_ = tools["echo_tool"]
|
||||
|
||||
def test_try_finally_with_filtered_tools(echo_server_script):
|
||||
serverparams = StdioServerParameters(
|
||||
command="uv", args=["run", "python", "-c", echo_server_script]
|
||||
)
|
||||
try:
|
||||
# Select both tools but in reverse order
|
||||
mcp_server_adapter = MCPServerAdapter(serverparams, "calc_tool", "echo_tool")
|
||||
tools = mcp_server_adapter.tools
|
||||
assert len(tools) == 2
|
||||
# The order of tools is based on filter_by_names which preserves
|
||||
# the original order from the collection
|
||||
assert tools[0].name == "calc_tool"
|
||||
assert tools[1].name == "echo_tool"
|
||||
finally:
|
||||
mcp_server_adapter.stop()
|
||||
|
||||
def test_filter_with_nonexistent_tool(echo_server_script):
|
||||
serverparams = StdioServerParameters(
|
||||
command="uv", args=["run", "python", "-c", echo_server_script]
|
||||
)
|
||||
# Include a tool that doesn't exist
|
||||
with MCPServerAdapter(serverparams, "echo_tool", "nonexistent_tool") as tools:
|
||||
# Only echo_tool should be in the result
|
||||
assert len(tools) == 1
|
||||
assert tools[0].name == "echo_tool"
|
||||
|
||||
def test_filter_with_only_nonexistent_tools(echo_server_script):
|
||||
serverparams = StdioServerParameters(
|
||||
command="uv", args=["run", "python", "-c", echo_server_script]
|
||||
)
|
||||
# All requested tools don't exist
|
||||
with MCPServerAdapter(serverparams, "nonexistent1", "nonexistent2") as tools:
|
||||
# Should return an empty tool collection
|
||||
assert isinstance(tools, ToolCollection)
|
||||
assert len(tools) == 0
|
||||
|
||||
Reference in New Issue
Block a user