mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-20 12:34:17 +02:00
Fix tool output schema generation to respect Pydantic serialization aliases (#1148)
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
parent
64d469be98
commit
0977d552ff
2 changed files with 84 additions and 2 deletions
|
|
@ -399,7 +399,7 @@ class ParsedFunction:
|
|||
|
||||
try:
|
||||
type_adapter = get_cached_typeadapter(clean_output_type)
|
||||
base_schema = type_adapter.json_schema()
|
||||
base_schema = type_adapter.json_schema(mode="serialization")
|
||||
|
||||
# Generate schema for wrapped type if it's non-object
|
||||
# because MCP requires that output schemas are objects
|
||||
|
|
@ -410,7 +410,7 @@ class ParsedFunction:
|
|||
# Use the wrapped result schema directly
|
||||
wrapped_type = _WrappedResult[clean_output_type]
|
||||
wrapped_adapter = get_cached_typeadapter(wrapped_type)
|
||||
output_schema = wrapped_adapter.json_schema()
|
||||
output_schema = wrapped_adapter.json_schema(mode="serialization")
|
||||
output_schema["x-fastmcp-wrap-result"] = True
|
||||
else:
|
||||
output_schema = base_schema
|
||||
|
|
|
|||
|
|
@ -1291,6 +1291,88 @@ class TestUnionReturnTypes:
|
|||
assert result2.structured_content == {"result": "error occurred"}
|
||||
|
||||
|
||||
class TestSerializationAlias:
|
||||
"""Tests for Pydantic field serialization alias support in tool output schemas."""
|
||||
|
||||
def test_output_schema_respects_serialization_alias(self):
|
||||
"""Test that Tool.from_function generates output schema using serialization alias."""
|
||||
from pydantic import AliasChoices, BaseModel, Field
|
||||
|
||||
class Component(BaseModel):
|
||||
"""Model with multiple validation aliases but specific serialization alias."""
|
||||
|
||||
component_id: str = Field(
|
||||
validation_alias=AliasChoices("id", "componentId"),
|
||||
serialization_alias="componentId",
|
||||
description="The ID of the component",
|
||||
)
|
||||
|
||||
async def get_component(
|
||||
component_id: str,
|
||||
) -> Annotated[Component, Field(description="The component.")]:
|
||||
# API returns data with 'id' field
|
||||
api_data = {"id": component_id}
|
||||
return Component.model_validate(api_data)
|
||||
|
||||
tool = Tool.from_function(get_component, name="get-component")
|
||||
|
||||
# The output schema should use the serialization alias 'componentId'
|
||||
# not the first validation alias 'id'
|
||||
assert tool.output_schema is not None
|
||||
|
||||
# Check the wrapped result schema
|
||||
assert "properties" in tool.output_schema
|
||||
assert "result" in tool.output_schema["properties"]
|
||||
assert "$defs" in tool.output_schema
|
||||
|
||||
# Find the Component definition
|
||||
component_def = list(tool.output_schema["$defs"].values())[0]
|
||||
|
||||
# Should have 'componentId' not 'id' in properties
|
||||
assert "componentId" in component_def["properties"]
|
||||
assert "id" not in component_def["properties"]
|
||||
|
||||
# Should require 'componentId' not 'id'
|
||||
assert "componentId" in component_def["required"]
|
||||
assert "id" not in component_def.get("required", [])
|
||||
|
||||
async def test_tool_execution_with_serialization_alias(self):
|
||||
"""Test that tool execution works correctly with serialization aliases."""
|
||||
from pydantic import AliasChoices, BaseModel, Field
|
||||
|
||||
from fastmcp import Client, FastMCP
|
||||
|
||||
class Component(BaseModel):
|
||||
"""Model with multiple validation aliases but specific serialization alias."""
|
||||
|
||||
component_id: str = Field(
|
||||
validation_alias=AliasChoices("id", "componentId"),
|
||||
serialization_alias="componentId",
|
||||
description="The ID of the component",
|
||||
)
|
||||
|
||||
mcp = FastMCP("TestServer")
|
||||
|
||||
@mcp.tool
|
||||
async def get_component(
|
||||
component_id: str,
|
||||
) -> Annotated[Component, Field(description="The component.")]:
|
||||
# API returns data with 'id' field
|
||||
api_data = {"id": component_id}
|
||||
return Component.model_validate(api_data)
|
||||
|
||||
async with Client(mcp) as client:
|
||||
# Execute the tool - this should work without validation errors
|
||||
result = await client.call_tool(
|
||||
"get_component", {"component_id": "test123"}
|
||||
)
|
||||
|
||||
# The result should contain the serialized form with 'componentId'
|
||||
assert result.structured_content is not None
|
||||
assert result.structured_content["result"]["componentId"] == "test123"
|
||||
assert "id" not in result.structured_content["result"]
|
||||
|
||||
|
||||
class TestToolTitle:
|
||||
"""Tests for tool title functionality."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue