mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-27 07:50:43 +02:00
Allow tool manager to store an alternative name for tools
This commit is contained in:
parent
2d86b326ba
commit
a691b40cf7
6 changed files with 279 additions and 43 deletions
|
|
@ -195,17 +195,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
|
||||
See `list_tools` for a more ergonomic way to list tools.
|
||||
"""
|
||||
|
||||
tools = self.list_tools()
|
||||
|
||||
return [
|
||||
MCPTool(
|
||||
name=info.name,
|
||||
description=info.description,
|
||||
inputSchema=info.parameters,
|
||||
)
|
||||
for info in tools
|
||||
]
|
||||
return self._tool_manager.list_mcp_tools()
|
||||
|
||||
def get_context(self) -> "Context[ServerSession, LifespanResultT]":
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import inspect
|
|||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Annotated, Any
|
||||
|
||||
from mcp.types import Tool as MCPTool
|
||||
from pydantic import BaseModel, BeforeValidator, Field
|
||||
|
||||
from fastmcp.exceptions import ToolError
|
||||
|
|
@ -101,6 +102,14 @@ class Tool(BaseModel):
|
|||
except Exception as e:
|
||||
raise ToolError(f"Error executing tool {self.name}: {e}") from e
|
||||
|
||||
def to_mcp_tool(self, **overrides: Any) -> MCPTool:
|
||||
kwargs = {
|
||||
"name": self.name,
|
||||
"description": self.description,
|
||||
"inputSchema": self.parameters,
|
||||
}
|
||||
return MCPTool(**kwargs | overrides)
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if not isinstance(other, Tool):
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
from __future__ import annotations as _annotations
|
||||
|
||||
import copy
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
|
|
@ -8,7 +7,7 @@ from mcp.shared.context import LifespanContextT
|
|||
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.settings import DuplicateBehavior
|
||||
from fastmcp.tools.tool import Tool
|
||||
from fastmcp.tools.tool import MCPTool, Tool
|
||||
from fastmcp.utilities.logging import get_logger
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -34,6 +33,10 @@ class ToolManager:
|
|||
"""List all registered tools."""
|
||||
return list(self._tools.values())
|
||||
|
||||
def list_mcp_tools(self) -> list[MCPTool]:
|
||||
"""List all registered tools in the format expected by the low-level MCP server."""
|
||||
return [tool.to_mcp_tool(name=name) for name, tool in self._tools.items()]
|
||||
|
||||
def add_tool_from_fn(
|
||||
self,
|
||||
fn: Callable[..., Any],
|
||||
|
|
@ -45,20 +48,21 @@ class ToolManager:
|
|||
tool = Tool.from_function(fn, name=name, description=description, tags=tags)
|
||||
return self.add_tool(tool)
|
||||
|
||||
def add_tool(self, tool: Tool) -> Tool:
|
||||
def add_tool(self, tool: Tool, name: str | None = None) -> Tool:
|
||||
"""Register a tool with the server."""
|
||||
existing = self._tools.get(tool.name)
|
||||
name = name or tool.name
|
||||
existing = self._tools.get(name)
|
||||
if existing:
|
||||
if self.duplicate_behavior == DuplicateBehavior.WARN:
|
||||
logger.warning(f"Tool already exists: {tool.name}")
|
||||
self._tools[tool.name] = tool
|
||||
logger.warning(f"Tool already exists: {name}")
|
||||
self._tools[name] = tool
|
||||
elif self.duplicate_behavior == DuplicateBehavior.REPLACE:
|
||||
self._tools[tool.name] = tool
|
||||
self._tools[name] = tool
|
||||
elif self.duplicate_behavior == DuplicateBehavior.ERROR:
|
||||
raise ValueError(f"Tool already exists: {tool.name}")
|
||||
raise ValueError(f"Tool already exists: {name}")
|
||||
elif self.duplicate_behavior == DuplicateBehavior.IGNORE:
|
||||
pass
|
||||
self._tools[tool.name] = tool
|
||||
self._tools[name] = tool
|
||||
return tool
|
||||
|
||||
async def call_tool(
|
||||
|
|
@ -90,10 +94,5 @@ class ToolManager:
|
|||
"""
|
||||
for name, tool in tool_manager._tools.items():
|
||||
prefixed_name = f"{prefix}{name}" if prefix else name
|
||||
|
||||
new_tool = copy.copy(tool)
|
||||
new_tool.name = prefixed_name
|
||||
|
||||
# Store the copied tool
|
||||
self.add_tool(new_tool)
|
||||
logger.debug(f'Imported tool "{name}" as "{prefixed_name}"')
|
||||
self.add_tool(tool, name=prefixed_name)
|
||||
logger.debug(f'Imported tool "{tool.name}" as "{prefixed_name}"')
|
||||
|
|
|
|||
|
|
@ -24,8 +24,9 @@ async def test_mount_basic_functionality():
|
|||
assert "sub_tool" in sub_app._tool_manager._tools
|
||||
|
||||
# Verify the original tool still exists in the sub-app
|
||||
tool = main_app._tool_manager._tools["sub_sub_tool"]
|
||||
assert tool.name == "sub_sub_tool"
|
||||
tool = main_app._tool_manager.get_tool("sub_sub_tool")
|
||||
assert tool is not None
|
||||
assert tool.name == "sub_tool"
|
||||
assert callable(tool.fn)
|
||||
|
||||
|
||||
|
|
@ -230,3 +231,78 @@ async def test_mount_lifespan():
|
|||
"exit SubApp",
|
||||
"exit MainApp",
|
||||
]
|
||||
|
||||
|
||||
async def test_mount_with_proxy_tools():
|
||||
"""Test mounting with tools that have custom names (proxy tools)."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
api_app = FastMCP("APIApp")
|
||||
|
||||
# Create a tool function
|
||||
def fetch_data(query: str) -> str:
|
||||
return f"Data for query: {query}"
|
||||
|
||||
# Add the tool to the API app with a custom name
|
||||
api_app.add_tool(fetch_data, name="get_data")
|
||||
|
||||
# Verify the tool is registered with the custom name in the source app
|
||||
assert api_app._tool_manager.get_tool("get_data") is not None
|
||||
|
||||
# Mount the API app to the main app
|
||||
main_app.mount("api", api_app)
|
||||
|
||||
# Verify the tool was imported with the prefixed custom name
|
||||
tool = main_app._tool_manager.get_tool("api_get_data")
|
||||
assert tool is not None
|
||||
|
||||
# The internal function name should be preserved
|
||||
assert tool.fn.__name__ == "fetch_data"
|
||||
|
||||
# The tool should be callable through the mounted name
|
||||
context = main_app.get_context()
|
||||
result = await main_app._tool_manager.call_tool(
|
||||
"api_get_data", {"query": "test"}, context=context
|
||||
)
|
||||
assert result == "Data for query: test"
|
||||
|
||||
|
||||
async def test_mount_nested_prefixed_tools():
|
||||
"""Test mounting tools with multiple layers of prefixes."""
|
||||
# Create apps
|
||||
main_app = FastMCP("MainApp")
|
||||
service_app = FastMCP("ServiceApp")
|
||||
provider_app = FastMCP("ProviderApp")
|
||||
|
||||
# Create a tool function
|
||||
def calculate_value(input: int) -> int:
|
||||
return input * 2
|
||||
|
||||
# Add the tool to the provider app with a custom name
|
||||
provider_app.add_tool(calculate_value, name="compute")
|
||||
|
||||
# The provider has a tool registered with a custom name
|
||||
assert provider_app._tool_manager.get_tool("compute") is not None
|
||||
|
||||
# First mount: Mount the provider app to the service app
|
||||
service_app.mount("provider", provider_app)
|
||||
|
||||
# Verify the tool is accessible in the service app with the first prefix
|
||||
assert service_app._tool_manager.get_tool("provider_compute") is not None
|
||||
|
||||
# Second mount: Mount the service app to the main app
|
||||
main_app.mount("service", service_app)
|
||||
|
||||
# Verify the tool is accessible in the main app with both prefixes
|
||||
nested_tool = main_app._tool_manager.get_tool("service_provider_compute")
|
||||
assert nested_tool is not None
|
||||
|
||||
# The internal function name should still be preserved after multiple mounts
|
||||
assert nested_tool.fn.__name__ == "calculate_value"
|
||||
|
||||
# The tool should be callable through the fully-qualified name
|
||||
context = main_app.get_context()
|
||||
result = await main_app._tool_manager.call_tool(
|
||||
"service_provider_compute", {"input": 21}, context=context
|
||||
)
|
||||
assert result == 42
|
||||
|
|
|
|||
|
|
@ -204,6 +204,30 @@ class TestToolDecorator:
|
|||
assert len(tools) == 1
|
||||
assert tools[0].tags == {"example", "test-tag"}
|
||||
|
||||
async def test_add_tool_with_custom_name(self):
|
||||
"""Test adding a tool with a custom name using server.add_tool()."""
|
||||
mcp = FastMCP()
|
||||
|
||||
def multiply(a: int, b: int) -> int:
|
||||
"""Multiply two numbers."""
|
||||
return a * b
|
||||
|
||||
# Add the tool with a custom name
|
||||
mcp.add_tool(multiply, name="custom_multiply")
|
||||
|
||||
# Check that the tool is registered with the custom name
|
||||
tools = mcp.list_tools()
|
||||
tool_names = [t.name for t in tools]
|
||||
assert "custom_multiply" in tool_names
|
||||
|
||||
# Call the tool by its custom name
|
||||
result = await mcp.call_tool("custom_multiply", {"a": 5, "b": 3})
|
||||
assert isinstance(result[0], TextContent)
|
||||
assert result[0].text == "15"
|
||||
|
||||
# Original name should not be registered
|
||||
assert "multiply" not in tool_names
|
||||
|
||||
|
||||
class TestResourceDecorator:
|
||||
async def test_no_resources_before_decorator(self):
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from pydantic import BaseModel
|
|||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.settings import DuplicateBehavior
|
||||
from fastmcp.tools import ToolManager
|
||||
from fastmcp.tools.tool import Tool
|
||||
|
||||
|
||||
class TestAddTools:
|
||||
|
|
@ -28,7 +29,6 @@ class TestAddTools:
|
|||
assert tool.parameters["properties"]["a"]["type"] == "integer"
|
||||
assert tool.parameters["properties"]["b"]["type"] == "integer"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_async_function(self):
|
||||
"""Test registering and running an async function."""
|
||||
|
||||
|
|
@ -137,8 +137,15 @@ class TestAddTools:
|
|||
|
||||
# Should have replaced the first tool with the second
|
||||
stored_tool = manager.get_tool("test_tool")
|
||||
assert stored_tool is not None
|
||||
assert stored_tool == replacement_tool
|
||||
|
||||
# The name should still be the same
|
||||
assert stored_tool.name == "test_tool"
|
||||
|
||||
# But the function is different
|
||||
assert stored_tool.fn.__name__ == "replacement_fn"
|
||||
|
||||
|
||||
class TestToolTags:
|
||||
"""Test functionality related to tool tags."""
|
||||
|
|
@ -232,7 +239,6 @@ class TestToolTags:
|
|||
|
||||
|
||||
class TestCallTools:
|
||||
@pytest.mark.anyio
|
||||
async def test_call_tool(self):
|
||||
def add(a: int, b: int) -> int:
|
||||
"""Add two numbers."""
|
||||
|
|
@ -243,7 +249,6 @@ class TestCallTools:
|
|||
result = await manager.call_tool("add", {"a": 1, "b": 2})
|
||||
assert result == 3
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_async_tool(self):
|
||||
async def double(n: int) -> int:
|
||||
"""Double a number."""
|
||||
|
|
@ -254,7 +259,6 @@ class TestCallTools:
|
|||
result = await manager.call_tool("double", {"n": 5})
|
||||
assert result == 10
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_tool_with_default_args(self):
|
||||
def add(a: int, b: int = 1) -> int:
|
||||
"""Add two numbers."""
|
||||
|
|
@ -265,7 +269,6 @@ class TestCallTools:
|
|||
result = await manager.call_tool("add", {"a": 1})
|
||||
assert result == 2
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_tool_with_missing_args(self):
|
||||
def add(a: int, b: int) -> int:
|
||||
"""Add two numbers."""
|
||||
|
|
@ -276,13 +279,11 @@ class TestCallTools:
|
|||
with pytest.raises(ToolError):
|
||||
await manager.call_tool("add", {"a": 1})
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_unknown_tool(self):
|
||||
manager = ToolManager()
|
||||
with pytest.raises(ToolError):
|
||||
await manager.call_tool("unknown", {"a": 1})
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_tool_with_list_int_input(self):
|
||||
def sum_vals(vals: list[int]) -> int:
|
||||
return sum(vals)
|
||||
|
|
@ -295,7 +296,6 @@ class TestCallTools:
|
|||
result = await manager.call_tool("sum_vals", {"vals": [1, 2, 3]})
|
||||
assert result == 6
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_tool_with_list_str_or_str_input(self):
|
||||
def concat_strs(vals: list[str] | str) -> str:
|
||||
return vals if isinstance(vals, str) else "".join(vals)
|
||||
|
|
@ -312,7 +312,6 @@ class TestCallTools:
|
|||
result = await manager.call_tool("concat_strs", {"vals": '"a"'})
|
||||
assert result == '"a"'
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_call_tool_with_complex_model(self):
|
||||
from fastmcp import Context
|
||||
|
||||
|
|
@ -341,7 +340,6 @@ class TestCallTools:
|
|||
|
||||
|
||||
class TestToolSchema:
|
||||
@pytest.mark.anyio
|
||||
async def test_context_arg_excluded_from_schema(self):
|
||||
from fastmcp import Context
|
||||
|
||||
|
|
@ -376,7 +374,6 @@ class TestContextHandling:
|
|||
tool = manager.add_tool_from_fn(tool_without_context)
|
||||
assert tool.context_kwarg is None
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_context_injection(self):
|
||||
"""Test that context is properly injected during tool execution."""
|
||||
from fastmcp import Context, FastMCP
|
||||
|
|
@ -393,7 +390,6 @@ class TestContextHandling:
|
|||
result = await manager.call_tool("tool_with_context", {"x": 42}, context=ctx)
|
||||
assert result == "42"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_context_injection_async(self):
|
||||
"""Test that context is properly injected in async tools."""
|
||||
from fastmcp import Context, FastMCP
|
||||
|
|
@ -410,7 +406,6 @@ class TestContextHandling:
|
|||
result = await manager.call_tool("async_tool", {"x": 42}, context=ctx)
|
||||
assert result == "42"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_context_optional(self):
|
||||
"""Test that context is optional when calling tools."""
|
||||
from fastmcp import Context
|
||||
|
|
@ -424,7 +419,6 @@ class TestContextHandling:
|
|||
result = await manager.call_tool("tool_with_context", {"x": 42})
|
||||
assert result == "42"
|
||||
|
||||
@pytest.mark.anyio
|
||||
async def test_context_error_handling(self):
|
||||
"""Test error handling when context injection fails."""
|
||||
from fastmcp import Context, FastMCP
|
||||
|
|
@ -545,3 +539,147 @@ class TestImportTools:
|
|||
assert (
|
||||
main_manager._tools["news/headlines"].fn.__name__ == headlines_fn.__name__
|
||||
)
|
||||
|
||||
|
||||
class TestCustomToolNames:
|
||||
"""Test adding tools with custom names that differ from their function names."""
|
||||
|
||||
def test_add_tool_with_custom_name(self):
|
||||
"""Test adding a tool with a custom name parameter using add_tool_from_fn."""
|
||||
|
||||
def original_fn(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
manager = ToolManager()
|
||||
tool = manager.add_tool_from_fn(original_fn, name="custom_name")
|
||||
|
||||
# The tool is stored under the custom name and its .name is also set to custom_name
|
||||
assert manager.get_tool("custom_name") is not None
|
||||
assert tool.name == "custom_name"
|
||||
assert tool.fn.__name__ == "original_fn"
|
||||
# The tool should not be accessible via its original function name
|
||||
assert manager.get_tool("original_fn") is None
|
||||
|
||||
def test_add_tool_object_with_custom_storage_name(self):
|
||||
"""Test adding a Tool object with a custom storage name using add_tool()."""
|
||||
|
||||
def fn(x: int) -> int:
|
||||
return x + 1
|
||||
|
||||
# Create a tool with a specific name
|
||||
tool = Tool.from_function(fn, name="my_tool")
|
||||
manager = ToolManager()
|
||||
# Store it under a different name
|
||||
manager.add_tool(tool, name="proxy_tool")
|
||||
# The tool is accessible under the storage name
|
||||
stored = manager.get_tool("proxy_tool")
|
||||
assert stored is not None
|
||||
# But the tool's .name is unchanged
|
||||
assert stored.name == "my_tool"
|
||||
# The tool is not accessible under its original name
|
||||
assert manager.get_tool("my_tool") is None
|
||||
|
||||
async def test_call_tool_with_custom_name(self):
|
||||
"""Test calling a tool added with a custom name."""
|
||||
|
||||
def multiply(a: int, b: int) -> int:
|
||||
"""Multiply two numbers."""
|
||||
return a * b
|
||||
|
||||
manager = ToolManager()
|
||||
manager.add_tool_from_fn(multiply, name="custom_multiply")
|
||||
|
||||
# Tool should be callable by its custom name
|
||||
result = await manager.call_tool("custom_multiply", {"a": 5, "b": 3})
|
||||
assert result == 15
|
||||
|
||||
# Original name should not be registered
|
||||
with pytest.raises(ToolError):
|
||||
await manager.call_tool("multiply", {"a": 5, "b": 3})
|
||||
|
||||
def test_tool_to_mcp_tool_with_custom_name(self):
|
||||
"""Test that to_mcp_tool uses the storage name, not the internal name."""
|
||||
|
||||
def some_function(x: int) -> int:
|
||||
return x
|
||||
|
||||
manager = ToolManager()
|
||||
manager.add_tool_from_fn(some_function, name="api_function")
|
||||
|
||||
# When listing tools for MCP, the custom name should be used
|
||||
mcp_tools = manager.list_mcp_tools()
|
||||
assert len(mcp_tools) == 1
|
||||
assert mcp_tools[0].name == "api_function"
|
||||
|
||||
def test_import_tools_with_custom_names(self):
|
||||
"""Test importing tools with custom names."""
|
||||
|
||||
def source_fn(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
# Create a source manager with a tool using custom name
|
||||
source_manager = ToolManager()
|
||||
source_manager.add_tool_from_fn(source_fn, name="custom_source")
|
||||
|
||||
# Import the tools to a target manager with a prefix
|
||||
target_manager = ToolManager()
|
||||
target_manager.import_tools(source_manager, "prefix/")
|
||||
|
||||
# The tool should be imported with the prefixed custom name
|
||||
assert target_manager.get_tool("prefix/custom_source") is not None
|
||||
assert target_manager.get_tool("prefix/source_fn") is None
|
||||
|
||||
def test_replace_tool_keeps_original_name(self):
|
||||
"""Test that replacing a tool with DuplicateBehavior.REPLACE keeps the original name."""
|
||||
|
||||
def original_fn(x: int) -> int:
|
||||
return x
|
||||
|
||||
def replacement_fn(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
# Create a manager with REPLACE behavior
|
||||
manager = ToolManager(duplicate_behavior=DuplicateBehavior.REPLACE)
|
||||
|
||||
# Add the original tool
|
||||
original_tool = manager.add_tool_from_fn(original_fn, name="test_tool")
|
||||
assert original_tool.name == "test_tool"
|
||||
|
||||
# Replace with a new function but keep the same registered name
|
||||
replacement_tool = manager.add_tool_from_fn(replacement_fn, name="test_tool")
|
||||
|
||||
# The tool object should have been replaced
|
||||
stored_tool = manager.get_tool("test_tool")
|
||||
assert stored_tool is not None
|
||||
assert stored_tool == replacement_tool
|
||||
|
||||
# The name should still be the same
|
||||
assert stored_tool.name == "test_tool"
|
||||
|
||||
# But the function is different
|
||||
assert stored_tool.fn.__name__ == "replacement_fn"
|
||||
|
||||
def test_mcp_tool_name_for_add_tool(self):
|
||||
"""Test MCPTool name for add_tool (storage name != tool.name)."""
|
||||
|
||||
def fn(x: int) -> int:
|
||||
return x + 1
|
||||
|
||||
tool = Tool.from_function(fn, name="my_tool")
|
||||
manager = ToolManager()
|
||||
manager.add_tool(tool, name="proxy_tool")
|
||||
mcp_tools = manager.list_mcp_tools()
|
||||
assert len(mcp_tools) == 1
|
||||
assert mcp_tools[0].name == "proxy_tool"
|
||||
|
||||
def test_mcp_tool_name_for_add_tool_from_fn(self):
|
||||
"""Test MCPTool name for add_tool_from_fn (storage name == tool.name)."""
|
||||
|
||||
def fn(x: int) -> int:
|
||||
return x + 1
|
||||
|
||||
manager = ToolManager()
|
||||
manager.add_tool_from_fn(fn, name="custom_fn")
|
||||
mcp_tools = manager.list_mcp_tools()
|
||||
assert len(mcp_tools) == 1
|
||||
assert mcp_tools[0].name == "custom_fn"
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue