diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py index ed3a5d6c0..968f84334 100644 --- a/src/fastmcp/tools/tool_transform.py +++ b/src/fastmcp/tools/tool_transform.py @@ -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. diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py index 1067d36bc..6217b441e 100644 --- a/tests/tools/test_tool_transform.py +++ b/tests/tools/test_tool_transform.py @@ -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: