mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-20 04:24:17 +02:00
Support enabled/disabled tools
This commit is contained in:
parent
17a71278aa
commit
6e92335510
4 changed files with 100 additions and 0 deletions
|
|
@ -625,6 +625,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | dict[str, Any] | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> FunctionTool: ...
|
||||
|
||||
@overload
|
||||
|
|
@ -637,6 +638,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | dict[str, Any] | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> Callable[[AnyFunction], FunctionTool]: ...
|
||||
|
||||
def tool(
|
||||
|
|
@ -648,6 +650,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | dict[str, Any] | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> Callable[[AnyFunction], FunctionTool] | FunctionTool:
|
||||
"""Decorator to register a tool.
|
||||
|
||||
|
|
@ -721,6 +724,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
annotations=annotations,
|
||||
exclude_args=exclude_args,
|
||||
serializer=self._tool_serializer,
|
||||
enabled=enabled,
|
||||
)
|
||||
self.add_tool(tool)
|
||||
return tool
|
||||
|
|
@ -749,6 +753,7 @@ class FastMCP(Generic[LifespanResultT]):
|
|||
tags=tags,
|
||||
annotations=annotations,
|
||||
exclude_args=exclude_args,
|
||||
enabled=enabled,
|
||||
)
|
||||
|
||||
def add_resource(self, resource: Resource, key: str | None = None) -> None:
|
||||
|
|
|
|||
|
|
@ -62,6 +62,7 @@ class Tool(FastMCPComponent, ABC):
|
|||
annotations: ToolAnnotations | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
serializer: Callable[[Any], str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> FunctionTool:
|
||||
"""Create a Tool from a function."""
|
||||
return FunctionTool.from_function(
|
||||
|
|
@ -72,6 +73,7 @@ class Tool(FastMCPComponent, ABC):
|
|||
annotations=annotations,
|
||||
exclude_args=exclude_args,
|
||||
serializer=serializer,
|
||||
enabled=enabled,
|
||||
)
|
||||
|
||||
@abstractmethod
|
||||
|
|
@ -92,6 +94,7 @@ class Tool(FastMCPComponent, ABC):
|
|||
tags: set[str] | None = None,
|
||||
annotations: ToolAnnotations | None = None,
|
||||
serializer: Callable[[Any], str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> TransformedTool:
|
||||
from fastmcp.tools.tool_transform import TransformedTool
|
||||
|
||||
|
|
@ -104,6 +107,7 @@ class Tool(FastMCPComponent, ABC):
|
|||
tags=tags,
|
||||
annotations=annotations,
|
||||
serializer=serializer,
|
||||
enabled=enabled,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -120,6 +124,7 @@ class FunctionTool(Tool):
|
|||
annotations: ToolAnnotations | None = None,
|
||||
exclude_args: list[str] | None = None,
|
||||
serializer: Callable[[Any], str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> FunctionTool:
|
||||
"""Create a Tool from a function."""
|
||||
|
||||
|
|
@ -136,6 +141,7 @@ class FunctionTool(Tool):
|
|||
tags=tags or set(),
|
||||
annotations=annotations,
|
||||
serializer=serializer,
|
||||
enabled=enabled if enabled is not None else True,
|
||||
)
|
||||
|
||||
async def run(
|
||||
|
|
|
|||
|
|
@ -267,6 +267,7 @@ class TransformedTool(Tool):
|
|||
transform_args: dict[str, ArgTransform] | None = None,
|
||||
annotations: ToolAnnotations | None = None,
|
||||
serializer: Callable[[Any], str] | None = None,
|
||||
enabled: bool | None = None,
|
||||
) -> TransformedTool:
|
||||
"""Create a transformed tool from a parent tool.
|
||||
|
||||
|
|
@ -399,6 +400,7 @@ class TransformedTool(Tool):
|
|||
annotations=annotations or tool.annotations,
|
||||
serializer=serializer or tool.serializer,
|
||||
transform_args=transform_args,
|
||||
enabled=enabled if enabled is not None else True,
|
||||
)
|
||||
|
||||
return transformed_tool
|
||||
|
|
|
|||
|
|
@ -699,6 +699,93 @@ class TestToolContextInjection:
|
|||
assert result[0].text == "3" # type: ignore[attr-defined]
|
||||
|
||||
|
||||
class TestToolEnabled:
|
||||
async def test_toggle_enabled(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def sample_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
assert sample_tool.enabled
|
||||
|
||||
tool = await mcp.get_tool("sample_tool")
|
||||
assert tool.enabled
|
||||
|
||||
tool.disable()
|
||||
|
||||
assert not tool.enabled
|
||||
assert not sample_tool.enabled
|
||||
|
||||
tool.enable()
|
||||
assert tool.enabled
|
||||
assert sample_tool.enabled
|
||||
|
||||
async def test_tool_disabled_in_decorator(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool(enabled=False)
|
||||
def sample_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools) == 0
|
||||
|
||||
async def test_tool_toggle_enabled(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool(enabled=False)
|
||||
def sample_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
sample_tool.enable()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools) == 1
|
||||
|
||||
async def test_tool_toggle_disabled(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def sample_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
sample_tool.disable()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tools = await client.list_tools()
|
||||
assert len(tools) == 0
|
||||
|
||||
async def test_get_tool_and_disable(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def sample_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
tool = await mcp.get_tool("sample_tool")
|
||||
assert tool.enabled
|
||||
|
||||
sample_tool.disable()
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.list_tools()
|
||||
assert len(result) == 0
|
||||
|
||||
async def test_cant_call_disabled_tool(self):
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool(enabled=False)
|
||||
def sample_tool(x: int) -> int:
|
||||
return x * 2
|
||||
|
||||
with pytest.raises(Exception, match="Unknown tool"):
|
||||
async with Client(mcp) as client:
|
||||
await client.call_tool("sample_tool", {"x": 5})
|
||||
|
||||
|
||||
class TestResource:
|
||||
async def test_text_resource(self):
|
||||
mcp = FastMCP()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue