ceos/fastmcp
1
0
Fork 0
mirror of https://github.com/PrefectHQ/fastmcp.git synced 2026-08-22 13:34:17 +02:00
Code Issues Projects Releases Packages Wiki Activity

Allow setting type with ArgTransform

Browse source
This commit is contained in:
Jeremiah Lowin 2025-06-07 17:32:22 -04:00
commit 642f9149d4
2 changed files with 308 additions and 45 deletions
Download patch file Download diff file Expand all files Collapse all files

135
src/fastmcp/tools/tool_transform.py
View file

@ -112,6 +112,7 @@ from mcp.types import EmbeddedResource, ImageContent, TextContent, ToolAnnotatio
from fastmcp.tools.tool import ParsedFunction, Tool
from fastmcp.utilities.logging import get_logger
from fastmcp.utilities.types import get_cached_typeadapter
if TYPE_CHECKING:
pass
@ -193,6 +194,7 @@ class ArgTransform:
name: New name for the argument. Use None to keep original name, or ... for no change.
description: New description for the argument. Use None to remove description, or ... for no change.
default: New default value for the argument. Use ... for no change.
type: New type for the argument. Use ... for no change.
drop: If True, remove this argument from the transformed tool's schema.
Examples:
@ -205,16 +207,20 @@ class ArgTransform:
# Add a default value (makes argument optional)
ArgTransform(default=42)
# Change the type
ArgTransform(type=str)
# Drop the argument entirely
ArgTransform(drop=True)
# Combine multiple transformations
ArgTransform(name="new_name", description="New desc", default=None)
ArgTransform(name="new_name", description="New desc", default=None, type=int)
"""
name: str | None | EllipsisType = ...
description: str | None | EllipsisType = ...
default: Any | EllipsisType = ...
type: Any | EllipsisType = ...
drop: bool = False
@ -259,9 +265,18 @@ class TransformedTool(Tool):
"""
from fastmcp.tools.tool import _convert_to_content
# Fill in missing arguments with schema defaults to ensure
# ArgTransform defaults take precedence over function defaults
filled_arguments = arguments.copy()
properties = self.parameters.get("properties", {})
for param_name, param_schema in properties.items():
if param_name not in filled_arguments and "default" in param_schema:
filled_arguments[param_name] = param_schema["default"]
token = _current_tool.set(self)
try:
result = await self.fn(**arguments)
result = await self.fn(**filled_arguments)
return _convert_to_content(result, serializer=self.serializer)
finally:
_current_tool.reset(token)
@ -357,51 +372,20 @@ class TransformedTool(Tool):
f"Function declares: {', '.join(sorted(fn_params))}"
)
# The function defines the final schema
final_schema = parsed_fn.parameters.copy()
# Inherit descriptions from transformed parent where possible
fn_props = final_schema.get("properties", {})
transformed_props = schema.get("properties", {})
for param_name in fn_props:
if param_name in transformed_props:
parent_desc = transformed_props[param_name].get("description")
if parent_desc and "description" not in fn_props[param_name]:
fn_props[param_name]["description"] = parent_desc
# ArgTransform takes precedence over function signature
# Start with function schema as base, then override with transformed schema
final_schema = cls._merge_schema_with_precedence(
parsed_fn.parameters, schema
)
else:
# With **kwargs, function can access all transformed params
# Function params override transformed params if they overlap
# ArgTransform takes precedence over function signature
# No validation needed - kwargs makes everything accessible
# Function accepts **kwargs, so use transformed schema as base
# and let function override specific parameters
fn_props = parsed_fn.parameters.get("properties", {})
fn_required = set(parsed_fn.parameters.get("required", []))
final_props = schema.get("properties", {}).copy()
final_required = set(schema.get("required", []))
# Override with function's parameters
for param_name, param_schema in fn_props.items():
# Inherit description from transformed parent if function doesn't provide one
if param_name in final_props and "description" not in param_schema:
param_schema = param_schema.copy()
param_schema["description"] = final_props[param_name].get(
"description"
)
final_props[param_name] = param_schema
if param_name in fn_required:
final_required.add(param_name)
else:
final_required.discard(param_name)
final_schema = {
"type": "object",
"properties": final_props,
"required": list(final_required),
}
# Start with function schema as base, then override with transformed schema
final_schema = cls._merge_schema_with_precedence(
parsed_fn.parameters, schema
)
# Additional validation: check for naming conflicts after transformation
if transform_args:
@ -578,11 +562,76 @@ class TransformedTool(Tool):
if transform.default is not ...:
new_schema["default"] = transform.default
is_required = False
if transform.type is not ...:
# Use TypeAdapter to get proper JSON schema for the type
type_schema = get_cached_typeadapter(transform.type).json_schema()
# Update the schema with the type information from TypeAdapter
new_schema.update(type_schema)
return new_name, new_schema, is_required # type: ignore[return-value]
raise ValueError(f"Invalid transform: {transform}")
@staticmethod
def _merge_schema_with_precedence(
base_schema: dict[str, Any], override_schema: dict[str, Any]
) -> dict[str, Any]:
"""Merge two schemas, with the override schema taking precedence.
Args:
base_schema: Base schema to start with
override_schema: Schema that takes precedence for overlapping properties
Returns:
Merged schema with override taking precedence
"""
merged_props = base_schema.get("properties", {}).copy()
merged_required = set(base_schema.get("required", []))
override_props = override_schema.get("properties", {})
override_required = set(override_schema.get("required", []))
# Override properties
for param_name, param_schema in override_props.items():
if param_name in merged_props:
# Merge the schemas, with override taking precedence
base_param = merged_props[param_name].copy()
base_param.update(param_schema)
merged_props[param_name] = base_param
else:
merged_props[param_name] = param_schema.copy()
# Handle required parameters - override takes complete precedence
# Start with override's required set
final_required = override_required.copy()
# For parameters not in override, inherit base requirement status
# but only if they don't have a default in the final merged properties
for param_name in merged_required:
if param_name not in override_props:
# Parameter not mentioned in override, keep base requirement status
final_required.add(param_name)
elif (
param_name in override_props
and "default" not in merged_props[param_name]
):
# Parameter in override but no default, keep required if it was required in base
if param_name not in override_required:
# Override doesn't specify it as required, and it has no default,
# so inherit from base
final_required.add(param_name)
# Remove any parameters that have defaults (they become optional)
for param_name, param_schema in merged_props.items():
if "default" in param_schema:
final_required.discard(param_name)
return {
"type": "object",
"properties": merged_props,
"required": list(final_required),
}
@staticmethod
def _function_has_kwargs(fn: Callable[..., Any]) -> bool:
"""Check if function accepts **kwargs.

218
tests/tools/test_tool_transform.py
View file

@ -1,9 +1,10 @@
import re
from typing import Annotated, Any
from dataclasses import dataclass
from typing import Annotated, Any, TypedDict
import pytest
from dirty_equals import IsList
from pydantic import Field
from pydantic import BaseModel, Field
from rich import print # type: ignore
from fastmcp import FastMCP
@ -388,6 +389,219 @@ async def test_chaining_transformations(add_tool):
assert "Chained:" in result[0].text # type: ignore
class MyModel(BaseModel):
x: int
y: str
@dataclass
class MyDataclass:
x: int
y: str
class MyTypedDict(TypedDict):
x: int
y: str
@pytest.mark.parametrize(
"py_type, json_type",
[
(int, "integer"),
(float, "number"),
(str, "string"),
(bool, "boolean"),
(list, "array"),
(list[int], "array"),
(dict, "object"),
(dict[str, int], "object"),
(MyModel, "object"),
(MyDataclass, "object"),
(MyTypedDict, "object"),
],
)
def test_arg_transform_type_handling(add_tool, py_type, json_type):
"""Test that ArgTransform type attribute gets applied to schema."""
new_tool = Tool.from_tool(
add_tool, transform_args={"old_x": ArgTransform(type=py_type)}
)
# Check that the type was changed in the schema
x_prop = get_property(new_tool, "old_x")
assert x_prop["type"] == json_type
def test_arg_transform_annotated_types(add_tool):
"""Test that ArgTransform works with annotated types and complex types."""
from typing import Annotated
from pydantic import Field
# Test with Annotated types
tool = Tool.from_tool(
add_tool,
transform_args={
"old_x": ArgTransform(
type=Annotated[int, Field(description="An annotated integer")]
)
},
)
x_prop = get_property(tool, "old_x")
assert x_prop["type"] == "integer"
# The ArgTransform description should override the annotation description
# (since we didn't set a description in ArgTransform, it should use the original)
# Test with Annotated string that has constraints
tool2 = Tool.from_tool(
add_tool,
transform_args={
"old_x": ArgTransform(
type=Annotated[str, Field(min_length=1, max_length=10)]
)
},
)
x_prop2 = get_property(tool2, "old_x")
assert x_prop2["type"] == "string"
assert x_prop2["minLength"] == 1
assert x_prop2["maxLength"] == 10
def test_arg_transform_precedence_over_function_without_kwargs():
"""Test that ArgTransform attributes take precedence over function signature (no **kwargs)."""
@Tool.from_function
def base(x: int, y: str = "default") -> str:
return f"{x}: {y}"
# Function signature says x: int with no default, y: str = "function_default"
# ArgTransform should override these
def custom_fn(x: str = "transform_default", y: int = 99) -> str:
return f"custom: {x}, {y}"
tool = Tool.from_tool(
base,
transform_fn=custom_fn,
transform_args={
"x": ArgTransform(type=str, default="transform_default"),
"y": ArgTransform(type=int, default=99),
},
)
# ArgTransform should take precedence
x_prop = get_property(tool, "x")
y_prop = get_property(tool, "y")
assert x_prop["type"] == "string" # ArgTransform type wins
assert x_prop["default"] == "transform_default" # ArgTransform default wins
assert y_prop["type"] == "integer" # ArgTransform type wins
assert y_prop["default"] == 99 # ArgTransform default wins
# Neither parameter should be required due to ArgTransform defaults
assert "x" not in tool.parameters["required"]
assert "y" not in tool.parameters["required"]
async def test_arg_transform_precedence_over_function_with_kwargs():
"""Test that ArgTransform attributes take precedence over function signature (with **kwargs)."""
@Tool.from_function
def base(x: int, y: str = "base_default") -> str:
return f"{x}: {y}"
# Function signature has different types/defaults than ArgTransform
async def custom_fn(x: str = "function_default", **kwargs) -> str:
result = await forward(x=x, **kwargs)
return f"custom: {result}"
tool = Tool.from_tool(
base,
transform_fn=custom_fn,
transform_args={
"x": ArgTransform(type=int, default=42), # Different type and default
"y": ArgTransform(description="ArgTransform description"),
},
)
# ArgTransform should take precedence
x_prop = get_property(tool, "x")
y_prop = get_property(tool, "y")
assert x_prop["type"] == "integer" # ArgTransform type wins over function's str
assert x_prop["default"] == 42 # ArgTransform default wins over function's default
assert (
y_prop["description"] == "ArgTransform description"
) # ArgTransform description
# x should not be required due to ArgTransform default
assert "x" not in tool.parameters["required"]
# Test it works at runtime
result = await tool.run(arguments={"y": "test"})
# Should use ArgTransform default of 42
assert "42: test" in result[0].text # type: ignore
def test_arg_transform_combined_attributes():
"""Test that multiple ArgTransform attributes work together."""
@Tool.from_function
def base(param: int) -> str:
return str(param)
tool = Tool.from_tool(
base,
transform_args={
"param": ArgTransform(
name="renamed_param",
type=str,
description="New description",
default="default_value",
)
},
)
# Check all attributes were applied
assert "renamed_param" in tool.parameters["properties"]
assert "param" not in tool.parameters["properties"]
prop = get_property(tool, "renamed_param")
assert prop["type"] == "string"
assert prop["description"] == "New description"
assert prop["default"] == "default_value"
assert "renamed_param" not in tool.parameters["required"] # Has default
async def test_arg_transform_type_precedence_runtime():
"""Test that ArgTransform type changes work correctly at runtime."""
@Tool.from_function
def base(x: int, y: int = 10) -> int:
return x + y
# Transform x to string type but keep same logic
async def custom_fn(x: str, y: int = 10) -> str:
# Convert string back to int for the original function
result = await forward_raw(x=int(x), y=y)
# Extract the text from the result
result_text = result[0].text
return f"String input '{x}' converted to result: {result_text}"
tool = Tool.from_tool(
base, transform_fn=custom_fn, transform_args={"x": ArgTransform(type=str)}
)
# Verify schema shows string type
assert get_property(tool, "x")["type"] == "string"
# Test it works with string input
result = await tool.run(arguments={"x": "5", "y": 3})
assert "String input '5'" in result[0].text # type: ignore
assert "result: 8" in result[0].text # type: ignore
class TestProxy:
@pytest.fixture
def mcp_server(self) -> FastMCP: