From a7eb2da273d7bc607a6a5b268d206bea71c919b2 Mon Sep 17 00:00:00 2001 From: Jeremiah Lowin <153965+jlowin@users.noreply.github.com> Date: Fri, 27 Jun 2025 20:17:10 -0400 Subject: [PATCH] Update output schema control --- src/fastmcp/tools/tool.py | 48 +++++++++++++++++++------- src/fastmcp/tools/tool_transform.py | 53 ++++++++++++++++++++++++----- tests/tools/test_tool.py | 40 +++++++++++----------- tests/tools/test_tool_transform.py | 38 ++++++++++----------- 4 files changed, 119 insertions(+), 60 deletions(-) diff --git a/src/fastmcp/tools/tool.py b/src/fastmcp/tools/tool.py index 3fd4f2004..9a53acb69 100644 --- a/src/fastmcp/tools/tool.py +++ b/src/fastmcp/tools/tool.py @@ -3,7 +3,7 @@ from __future__ import annotations import inspect from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, Annotated, Any +from typing import TYPE_CHECKING, Annotated, Any, Literal import mcp.types import pydantic_core @@ -40,6 +40,27 @@ def default_serializer(data: Any) -> str: return pydantic_core.to_json(data, fallback=str, indent=2).decode() +def _wrap_schema_if_needed(schema: dict[str, Any] | None) -> dict[str, Any] | None: + """Wrap non-object schemas with result property for structured output. + + This wrapping allows primitive types (int, str, etc.) to be returned as + structured content by placing them under a "result" key. + + Args: + schema: The JSON schema to potentially wrap + + Returns: + Wrapped schema if needed, or original schema if already an object type + """ + if schema and schema.get("type") != "object": + return { + "type": "object", + "properties": {"result": schema}, + "x-fastmcp-wrap-result": True, + } + return schema + + class ToolResult: def __init__( self, @@ -64,7 +85,11 @@ class ToolResult: ) raise if not isinstance(structured_content, dict): - structured_content = {"result": structured_content} + raise ValueError( + "structured_content must be a dict or None. " + f"Got {type(structured_content).__name__}: {structured_content!r}. " + "Tools should wrap non-dict values based on their output_schema." + ) self.structured_content: dict[str, Any] | None = structured_content def to_mcp_result( @@ -127,7 +152,7 @@ class Tool(FastMCPComponent): tags: set[str] | None = None, annotations: ToolAnnotations | None = None, exclude_args: list[str] | None = None, - output_schema: dict[str, Any] | None | NotSetT = NotSet, + output_schema: dict[str, Any] | None | NotSetT | Literal[False] = NotSet, serializer: Callable[[Any], str] | None = None, enabled: bool | None = None, ) -> FunctionTool: @@ -166,6 +191,7 @@ class Tool(FastMCPComponent): description: str | None = None, tags: set[str] | None = None, annotations: ToolAnnotations | None = None, + output_schema: dict[str, Any] | None | Literal[False] = None, serializer: Callable[[Any], str] | None = None, enabled: bool | None = None, ) -> TransformedTool: @@ -179,6 +205,7 @@ class Tool(FastMCPComponent): description=description, tags=tags, annotations=annotations, + output_schema=output_schema, serializer=serializer, enabled=enabled, ) @@ -196,7 +223,7 @@ class FunctionTool(Tool): tags: set[str] | None = None, annotations: ToolAnnotations | None = None, exclude_args: list[str] | None = None, - output_schema: dict[str, Any] | None | NotSetT = NotSet, + output_schema: dict[str, Any] | None | NotSetT | Literal[False] = NotSet, serializer: Callable[[Any], str] | None = None, enabled: bool | None = None, ) -> FunctionTool: @@ -208,14 +235,10 @@ class FunctionTool(Tool): raise ValueError("You must provide a name for lambda functions") if isinstance(output_schema, NotSetT): - output_schema = parsed_fn.output_schema - - if output_schema and output_schema.get("type") != "object": - output_schema = { - "type": "object", - "properties": {"result": output_schema}, - "x-fastmcp-wrap-result": True, - } + output_schema = _wrap_schema_if_needed(parsed_fn.output_schema) + elif output_schema is False: + output_schema = None + # Note: explicit schemas (dict) are used as-is without auto-wrapping return cls( fn=parsed_fn.fn, @@ -249,6 +272,7 @@ class FunctionTool(Tool): unstructured_result = _convert_to_content(result, serializer=self.serializer) + # Handle structured content based on output schema if self.output_schema is not None: if self.output_schema.get("x-fastmcp-wrap-result"): structured_output = {"result": result} diff --git a/src/fastmcp/tools/tool_transform.py b/src/fastmcp/tools/tool_transform.py index 38d074430..f52a54da3 100644 --- a/src/fastmcp/tools/tool_transform.py +++ b/src/fastmcp/tools/tool_transform.py @@ -9,7 +9,7 @@ from typing import Any, Literal from mcp.types import ToolAnnotations from pydantic import ConfigDict -from fastmcp.tools.tool import ParsedFunction, Tool, ToolResult +from fastmcp.tools.tool import ParsedFunction, Tool, ToolResult, _wrap_schema_if_needed from fastmcp.utilities.logging import get_logger from fastmcp.utilities.types import NotSet, NotSetT, get_cached_typeadapter @@ -52,7 +52,7 @@ async def forward(**kwargs) -> ToolResult: return await tool.forwarding_fn(**kwargs) -async def forward_raw(**kwargs) -> Any: +async def forward_raw(**kwargs) -> ToolResult: """Forward directly to parent tool without transformation. This function bypasses all argument transformation and validation, calling the parent @@ -66,7 +66,7 @@ async def forward_raw(**kwargs) -> Any: **kwargs: Arguments to pass directly to the parent tool (using original names). Returns: - The result from the parent tool execution. + The ToolResult from the parent tool execution. Raises: RuntimeError: If called outside a transformed tool context. @@ -268,15 +268,31 @@ class TransformedTool(Tool): token = _current_tool.set(self) try: result = await self.fn(**arguments) - + # If transform function returns ToolResult, use it directly if isinstance(result, ToolResult): return result - - # Otherwise convert to content and create basic ToolResult + + # Otherwise convert to content and create ToolResult with proper structured content from fastmcp.tools.tool import _convert_to_content - unstructured_result = _convert_to_content(result, serializer=self.serializer) - return ToolResult(content=unstructured_result) + + unstructured_result = _convert_to_content( + result, serializer=self.serializer + ) + + # Handle structured content based on output schema + if self.output_schema is not None: + if self.output_schema.get("x-fastmcp-wrap-result"): + structured_output = {"result": result} + else: + structured_output = result + else: + structured_output = None + + return ToolResult( + content=unstructured_result, + structured_content=structured_output, + ) finally: _current_tool.reset(token) @@ -290,6 +306,7 @@ class TransformedTool(Tool): transform_fn: Callable[..., Any] | None = None, transform_args: dict[str, ArgTransform] | None = None, annotations: ToolAnnotations | None = None, + output_schema: dict[str, Any] | None | Literal[False] = None, serializer: Callable[[Any], str] | None = None, enabled: bool | None = None, ) -> TransformedTool: @@ -352,13 +369,30 @@ class TransformedTool(Tool): # Always create the forwarding transform schema, forwarding_fn = cls._create_forwarding_transform(tool, transform_args) + # Handle output schema with smart fallback + if output_schema is False: + final_output_schema = None + elif output_schema is not None: + # Explicit schema provided - use as-is + final_output_schema = output_schema + else: + # Smart fallback: try custom function, then parent, then None + if transform_fn is not None: + parsed_fn = ParsedFunction.from_function(transform_fn, validate=False) + final_output_schema = _wrap_schema_if_needed(parsed_fn.output_schema) + if final_output_schema is None: + final_output_schema = tool.output_schema + else: + final_output_schema = tool.output_schema + if transform_fn is None: # User wants pure transformation - use forwarding_fn as the main function final_fn = forwarding_fn final_schema = schema else: # User provided custom function - merge schemas - parsed_fn = ParsedFunction.from_function(transform_fn, validate=False) + if "parsed_fn" not in locals(): + parsed_fn = ParsedFunction.from_function(transform_fn, validate=False) final_fn = transform_fn has_kwargs = cls._function_has_kwargs(transform_fn) @@ -426,6 +460,7 @@ class TransformedTool(Tool): name=name or tool.name, description=final_description, parameters=final_schema, + output_schema=final_output_schema, tags=tags or tool.tags, annotations=annotations or tool.annotations, serializer=serializer or tool.serializer, diff --git a/tests/tools/test_tool.py b/tests/tools/test_tool.py index eb5b045a0..b03559cdf 100644 --- a/tests/tools/test_tool.py +++ b/tests/tools/test_tool.py @@ -31,12 +31,10 @@ class TestToolFromFunction: assert len(tool.parameters["properties"]) == 2 assert tool.parameters["properties"]["a"]["type"] == "integer" assert tool.parameters["properties"]["b"]["type"] == "integer" - # With primitive wrapping, int return type becomes object with value property + # With primitive wrapping, int return type becomes object with result property expected_schema = { "type": "object", - "properties": {"value": {"title": "Value", "type": "integer"}}, - "required": ["value"], - "title": "Result", + "properties": {"result": {"type": "integer"}}, "x-fastmcp-wrap-result": True, } assert tool.output_schema == expected_schema @@ -251,7 +249,7 @@ class TestToolFromFunction: assert isinstance(result.content[0], TextContent) assert result.content[0].text == "Custom serializer: 15" # Structured output should have the raw value - assert result.structured_content == {"value": 15} + assert result.structured_content == {"result": 15} class TestToolFromFunctionOutputSchema: @@ -287,27 +285,20 @@ class TestToolFromFunctionOutputSchema: base_schema = TypeAdapter(annotation).json_schema() - # Only pure primitives (just type + optional title) get wrapped - primitive_types = {"string", "number", "integer", "boolean", "null"} + # Non-object types get wrapped schema_type = base_schema.get("type") - is_pure_primitive = ( - schema_type in primitive_types - and len(base_schema) <= 2 # Only 'type' and optionally 'title' - and all(key in {"type", "title"} for key in base_schema.keys()) - ) + is_object_type = schema_type == "object" - if is_pure_primitive: - # Pure primitives get wrapped + if not is_object_type: + # Non-object types get wrapped expected_schema = { "type": "object", - "properties": {"value": base_schema | {"title": "Value"}}, - "required": ["value"], - "title": "Result", + "properties": {"result": base_schema}, "x-fastmcp-wrap-result": True, } assert tool.output_schema == expected_schema else: - # Complex types (objects, unions, constrained types) remain unwrapped + # Object types remain unwrapped assert tool.output_schema == base_schema @pytest.mark.parametrize( @@ -326,8 +317,17 @@ class TestToolFromFunctionOutputSchema: tool = Tool.from_function(func) base_schema = TypeAdapter(annotation).json_schema() - # Complex types with constraints are not wrapped - they remain as-is - assert tool.output_schema == base_schema + # Special case for Any type - it generates an empty schema and doesn't get wrapped + if annotation is Any: + assert tool.output_schema == base_schema # Should be {} + else: + # All other non-object types get wrapped, including complex constrained types + expected_schema = { + "type": "object", + "properties": {"result": base_schema}, + "x-fastmcp-wrap-result": True, + } + assert tool.output_schema == expected_schema @pytest.mark.parametrize( "annotation, expected", diff --git a/tests/tools/test_tool_transform.py b/tests/tools/test_tool_transform.py index a616f556c..014422452 100644 --- a/tests/tools/test_tool_transform.py +++ b/tests/tools/test_tool_transform.py @@ -255,7 +255,7 @@ async def test_forward_raw_without_argument_mapping(add_tool): async def test_custom_fn_with_kwargs_and_no_transform_args(add_tool): async def custom_fn(extra: int, **kwargs) -> int: sum = await forward(**kwargs) - return int(sum[0].text) + extra # type: ignore[attr-defined] + return int(sum.content[0].text) + extra # type: ignore[attr-defined] new_tool = Tool.from_tool(add_tool, transform_fn=custom_fn) result = await new_tool.run(arguments={"extra": 1, "old_x": 2, "old_y": 3}) @@ -360,7 +360,7 @@ async def test_fn_with_kwargs_dropped_args_not_in_kwargs(add_tool): ) # drop 'old_y' result = await new_tool.run(arguments={"new_x": 8}) # 8 + 10 (default value of b in parent) - assert result[0].text == "18" # type: ignore[attr-defined] + assert result.content[0].text == "18" # type: ignore[attr-defined] async def test_forward_outside_context_raises_error(): @@ -480,18 +480,18 @@ async def test_tool_transform_chaining(add_tool): tool2 = Tool.from_tool(tool1, transform_args={"x": ArgTransform(name="final_x")}) result = await tool2.run(arguments={"final_x": 5}) - assert result[0].text == "15" # type: ignore[attr-defined] + assert result.content[0].text == "15" # type: ignore[attr-defined] # Transform tool1 with custom function that handles all parameters async def custom(final_x: int, **kwargs) -> str: result = await forward(final_x=final_x, **kwargs) - return f"custom {result[0].text}" # Extract text from content + return f"custom {result.content[0].text}" # Extract text from content tool3 = Tool.from_tool( tool1, transform_fn=custom, transform_args={"x": ArgTransform(name="final_x")} ) result = await tool3.run(arguments={"final_x": 3, "old_y": 5}) - assert result[0].text == "custom 8" # type: ignore[attr-defined] + assert result.content[0].text == "custom 8" # type: ignore[attr-defined] class MyModel(BaseModel): @@ -619,7 +619,7 @@ async def test_arg_transform_precedence_over_function_with_kwargs(): # 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}" + return f"custom: {result.content[0].text}" tool = Tool.from_tool( base, @@ -646,7 +646,7 @@ async def test_arg_transform_precedence_over_function_with_kwargs(): # 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[attr-defined] + assert "42: test" in result.content[0].text # type: ignore[attr-defined] def test_arg_transform_combined_attributes(): @@ -691,7 +691,7 @@ async def test_arg_transform_type_precedence_runtime(): # 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 + result_text = result.content[0].text return f"String input '{x}' converted to result: {result_text}" tool = Tool.from_tool( @@ -703,8 +703,8 @@ async def test_arg_transform_type_precedence_runtime(): # 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[attr-defined] - assert "result: 8" in result[0].text # type: ignore[attr-defined] + assert "String input '5'" in result.content[0].text # type: ignore[attr-defined] + assert "result: 8" in result.content[0].text # type: ignore[attr-defined] class TestProxy: @@ -739,7 +739,7 @@ class TestProxy: async with Client(proxy_server) as client: # The tool should be registered with its transformed name result = await client.call_tool("add_transformed", {"new_x": 1, "old_y": 2}) - assert result[0].text == "3" # type: ignore[attr-defined] + assert result.content[0].text == "3" # type: ignore[attr-defined] async def test_arg_transform_default_factory(): @@ -762,7 +762,7 @@ async def test_arg_transform_default_factory(): # Should work without providing timestamp (gets value from factory) result = await new_tool.run(arguments={"x": 42}) - assert result[0].text == "42_12345.0" # type: ignore[attr-defined] + assert result.content[0].text == "42_12345.0" # type: ignore[attr-defined] async def test_arg_transform_default_factory_called_each_time(): @@ -790,11 +790,11 @@ async def test_arg_transform_default_factory_called_each_time(): # First call result1 = await new_tool.run(arguments={"x": 1}) - assert result1[0].text == "1_1" # type: ignore[attr-defined] + assert result1.content[0].text == "1_1" # type: ignore[attr-defined] # Second call should get a different value result2 = await new_tool.run(arguments={"x": 2}) - assert result2[0].text == "2_2" # type: ignore[attr-defined] + assert result2.content[0].text == "2_2" # type: ignore[attr-defined] async def test_arg_transform_hidden_with_default_factory(): @@ -819,7 +819,7 @@ async def test_arg_transform_hidden_with_default_factory(): # Should pass hidden request_id with factory value result = await new_tool.run(arguments={"x": 42}) - assert result[0].text == "42_req_123" # type: ignore[attr-defined] + assert result.content[0].text == "42_req_123" # type: ignore[attr-defined] async def test_arg_transform_default_and_factory_raises_error(): @@ -856,7 +856,7 @@ async def test_arg_transform_required_true(): # Should work when parameter is provided result = await new_tool.run(arguments={"optional_param": 100}) - assert result[0].text == "value: 100" # type: ignore + assert result.content[0].text == "value: 100" # type: ignore # Should fail when parameter is not provided with pytest.raises(TypeError, match="Missing required argument"): @@ -903,7 +903,7 @@ async def test_arg_transform_required_with_rename(): # Should work with new name result = await new_tool.run(arguments={"new_param": 200}) - assert result[0].text == "value: 200" # type: ignore + assert result.content[0].text == "value: 200" # type: ignore async def test_arg_transform_required_true_with_default_raises_error(): @@ -945,7 +945,7 @@ async def test_arg_transform_required_no_change(): # Should work as expected result = await new_tool.run(arguments={"req": 1}) - assert result[0].text == "values: 1, 42" # type: ignore + assert result.content[0].text == "values: 1, 42" # type: ignore async def test_arg_transform_hide_and_required_raises_error(): @@ -977,7 +977,7 @@ class TestEnableDisable: assert {tool.name for tool in tools} == {"new_add"} result = await client.call_tool("new_add", {"x": 1, "y": 2}) - assert result[0].text == "3" # type: ignore[attr-defined] + assert result.content[0].text == "3" # type: ignore[attr-defined] with pytest.raises(ToolError): await client.call_tool("add", {"x": 1, "y": 2})