mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 14:04:18 +02:00
Align dataclass output schemas
This commit is contained in:
parent
330722aad6
commit
98bf8557d5
2 changed files with 74 additions and 0 deletions
|
|
@ -164,6 +164,15 @@ class _ToolOutputSchemaGenerator(GenerateJsonSchema):
|
|||
finally:
|
||||
self.by_alias = previous_by_alias
|
||||
|
||||
def dataclass_schema(self, schema: core_schema.DataclassSchema) -> JsonSchemaValue:
|
||||
previous_by_alias = self.by_alias
|
||||
configured = (schema.get("config") or {}).get("serialize_by_alias")
|
||||
self.by_alias = False if configured is None else configured
|
||||
try:
|
||||
return super().dataclass_schema(schema)
|
||||
finally:
|
||||
self.by_alias = previous_by_alias
|
||||
|
||||
|
||||
T = TypeVarExt("T", default=Any)
|
||||
|
||||
|
|
@ -424,6 +433,7 @@ class ParsedFunction:
|
|||
type_adapter = get_cached_typeadapter(clean_output_type)
|
||||
base_schema = type_adapter.json_schema(
|
||||
mode="serialization",
|
||||
by_alias=False,
|
||||
schema_generator=_ToolOutputSchemaGenerator,
|
||||
)
|
||||
|
||||
|
|
@ -437,6 +447,7 @@ class ParsedFunction:
|
|||
wrapped_adapter = get_cached_typeadapter(wrapped_type)
|
||||
output_schema = wrapped_adapter.json_schema(
|
||||
mode="serialization",
|
||||
by_alias=False,
|
||||
schema_generator=_ToolOutputSchemaGenerator,
|
||||
)
|
||||
output_schema["x-fastmcp-wrap-result"] = True
|
||||
|
|
|
|||
|
|
@ -5,6 +5,8 @@ from typing import Annotated, Any
|
|||
import pytest
|
||||
from mcp_types import CallToolResult, TextContent
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
from pydantic.dataclasses import dataclass as pydantic_dataclass
|
||||
from typing_extensions import TypedDict
|
||||
|
||||
from fastmcp import Client, FastMCP
|
||||
from fastmcp.tools.base import Tool, ToolResult
|
||||
|
|
@ -437,6 +439,67 @@ class TestSerializeByAlias:
|
|||
"aliased": {"aliasedValue": "aliased"},
|
||||
}
|
||||
|
||||
async def test_dataclass_uses_pydantic_alias_default(self):
|
||||
"""A dataclass schema uses field names when aliases are not enabled."""
|
||||
|
||||
@dataclass
|
||||
class Output:
|
||||
value: Annotated[str, Field(serialization_alias="dataValue")]
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_output() -> Output:
|
||||
return Output(value="data")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tools = {tool.name: tool for tool in await client.list_tools()}
|
||||
result = await client.call_tool("get_output", {})
|
||||
|
||||
assert set(tools["get_output"].output_schema["properties"]) == {"value"} # type: ignore[index]
|
||||
assert result.structured_content == {"value": "data"}
|
||||
|
||||
async def test_pydantic_dataclass_can_enable_aliases(self):
|
||||
"""A Pydantic dataclass can opt in to serialization aliases."""
|
||||
|
||||
@pydantic_dataclass(config=ConfigDict(serialize_by_alias=True))
|
||||
class Output:
|
||||
value: Annotated[str, Field(serialization_alias="dataValue")]
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_output() -> Output:
|
||||
return Output(value="data")
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tools = {tool.name: tool for tool in await client.list_tools()}
|
||||
result = await client.call_tool("get_output", {})
|
||||
|
||||
assert set(tools["get_output"].output_schema["properties"]) == { # type: ignore[index]
|
||||
"dataValue"
|
||||
}
|
||||
assert result.structured_content == {"dataValue": "data"}
|
||||
|
||||
async def test_typed_dict_uses_pydantic_alias_default(self):
|
||||
"""A TypedDict schema uses the field names emitted by Pydantic."""
|
||||
|
||||
class Output(TypedDict):
|
||||
value: Annotated[str, Field(serialization_alias="dataValue")]
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_output() -> Output:
|
||||
return {"value": "data"}
|
||||
|
||||
async with Client(mcp) as client:
|
||||
tools = {tool.name: tool for tool in await client.list_tools()}
|
||||
result = await client.call_tool("get_output", {})
|
||||
|
||||
assert set(tools["get_output"].output_schema["properties"]) == {"value"} # type: ignore[index]
|
||||
assert result.structured_content == {"value": "data"}
|
||||
|
||||
async def test_serialize_by_alias_true_uses_alias(self):
|
||||
"""serialize_by_alias=True emits aliases."""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue