mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-24 06:24:18 +02:00
Narrow typed output alias handling
This commit is contained in:
Jeremiah Lowin
parent
d872f9de41
commit
827d444c48
3 changed files with 19 additions and 265 deletions
|
|
@ -174,87 +174,6 @@ class _ToolOutputSchemaGenerator(GenerateJsonSchema):
|
|||
self.by_alias = previous_by_alias
|
||||
|
||||
|
||||
def _has_contextual_alias_conflict(schema: core_schema.CoreSchema) -> bool:
|
||||
"""Whether one structural type is serialized under conflicting alias modes."""
|
||||
schema_refs: dict[str, dict[str, Any]] = {}
|
||||
|
||||
def collect(value: Any) -> None:
|
||||
if isinstance(value, dict):
|
||||
ref = value.get("ref")
|
||||
if isinstance(ref, str):
|
||||
schema_refs[ref] = value
|
||||
for child in value.values():
|
||||
collect(child)
|
||||
elif isinstance(value, list | tuple):
|
||||
for child in value:
|
||||
collect(child)
|
||||
|
||||
def has_alias(value: dict[str, Any]) -> bool:
|
||||
if value.get("type") == "dataclass":
|
||||
fields = value.get("schema", {}).get("fields", [])
|
||||
return any(
|
||||
field.get("serialization_alias") not in (None, field.get("name"))
|
||||
for field in fields
|
||||
)
|
||||
|
||||
fields = value.get("fields", {})
|
||||
return any(
|
||||
field.get("serialization_alias") not in (None, name)
|
||||
for name, field in fields.items()
|
||||
)
|
||||
|
||||
collect(schema)
|
||||
alias_contexts: dict[str, set[bool]] = {}
|
||||
|
||||
def visit(value: Any, by_alias: bool, active: set[str]) -> bool:
|
||||
if isinstance(value, list | tuple):
|
||||
return any(visit(child, by_alias, active) for child in value)
|
||||
if not isinstance(value, dict):
|
||||
return False
|
||||
|
||||
schema_type = value.get("type")
|
||||
if schema_type == "definitions":
|
||||
return visit(value.get("schema"), by_alias, active)
|
||||
if schema_type == "definition-ref":
|
||||
ref = value.get("schema_ref")
|
||||
target = schema_refs.get(ref)
|
||||
return target is not None and visit(target, by_alias, active)
|
||||
|
||||
next_by_alias = by_alias
|
||||
if schema_type == "model":
|
||||
configured = (value.get("config") or {}).get("serialize_by_alias")
|
||||
next_by_alias = False if configured is None else configured
|
||||
elif schema_type == "dataclass":
|
||||
configured = (value.get("config") or {}).get("serialize_by_alias")
|
||||
if configured is not None:
|
||||
next_by_alias = configured
|
||||
|
||||
ref = value.get("ref")
|
||||
if (
|
||||
isinstance(schema_type, str)
|
||||
and schema_type in {"dataclass", "typed-dict"}
|
||||
and isinstance(ref, str)
|
||||
and has_alias(value)
|
||||
):
|
||||
contexts = alias_contexts.setdefault(ref, set())
|
||||
contexts.add(next_by_alias)
|
||||
if len(contexts) > 1:
|
||||
return True
|
||||
|
||||
if isinstance(ref, str):
|
||||
if ref in active:
|
||||
return False
|
||||
active = active | {ref}
|
||||
|
||||
return any(
|
||||
visit(child, next_by_alias, active)
|
||||
for key, child in value.items()
|
||||
if key not in {"cls", "config", "metadata", "ref"}
|
||||
)
|
||||
|
||||
return visit(schema, False, set())
|
||||
|
||||
|
||||
T = TypeVarExt("T", default=Any)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
|
@ -518,43 +437,24 @@ class ParsedFunction:
|
|||
schema_generator=_ToolOutputSchemaGenerator,
|
||||
)
|
||||
|
||||
if _has_contextual_alias_conflict(type_adapter.core_schema):
|
||||
logger.debug(
|
||||
"Unable to generate one output schema for contextual "
|
||||
"serialization aliases in type %r",
|
||||
output_type,
|
||||
# Generate schema for wrapped type if it's non-object
|
||||
# because MCP requires that output schemas are objects
|
||||
# Check if schema is an object type, resolving $ref references
|
||||
# (self-referencing types use $ref at root level)
|
||||
if wrap_non_object_output_schema and not _is_object_schema(base_schema):
|
||||
# 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(
|
||||
mode="serialization",
|
||||
by_alias=False,
|
||||
schema_generator=_ToolOutputSchemaGenerator,
|
||||
)
|
||||
if _is_object_schema(base_schema):
|
||||
output_schema = {"type": "object"}
|
||||
elif wrap_non_object_output_schema:
|
||||
output_schema = {
|
||||
"type": "object",
|
||||
"properties": {"result": {}},
|
||||
"required": ["result"],
|
||||
"x-fastmcp-wrap-result": True,
|
||||
}
|
||||
output_schema["x-fastmcp-wrap-result"] = True
|
||||
else:
|
||||
# Generate schema for wrapped type if it's non-object
|
||||
# because MCP requires that output schemas are objects
|
||||
# Check if schema is an object type, resolving $ref references
|
||||
# (self-referencing types use $ref at root level)
|
||||
if wrap_non_object_output_schema and not _is_object_schema(
|
||||
base_schema
|
||||
):
|
||||
# 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(
|
||||
mode="serialization",
|
||||
by_alias=False,
|
||||
schema_generator=_ToolOutputSchemaGenerator,
|
||||
)
|
||||
output_schema["x-fastmcp-wrap-result"] = True
|
||||
else:
|
||||
output_schema = base_schema
|
||||
output_schema = base_schema
|
||||
|
||||
if output_schema is not None:
|
||||
output_schema = compress_schema(output_schema, prune_titles=True)
|
||||
output_schema = compress_schema(output_schema, prune_titles=True)
|
||||
|
||||
except PydanticSchemaGenerationError as e:
|
||||
if "_UnserializableType" not in str(e):
|
||||
|
|
|
|||
|
|
@ -339,28 +339,8 @@ class TestSerializeByAlias:
|
|||
"filepath",
|
||||
}
|
||||
|
||||
async def test_model_in_untyped_dict_respects_config(self):
|
||||
"""A model nested in an untyped dict keeps its serialization config."""
|
||||
|
||||
class Biofile(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=False)
|
||||
id: str = Field(alias="_id")
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_biofile() -> dict:
|
||||
return {"status": "success", "item": Biofile(_id="1")}
|
||||
|
||||
async with Client(mcp) as client:
|
||||
result = await client.call_tool("get_biofile", {})
|
||||
|
||||
expected = {"status": "success", "item": {"id": "1"}}
|
||||
assert result.structured_content == expected
|
||||
assert json.loads(result.content[0].text) == expected # type: ignore[union-attr]
|
||||
|
||||
async def test_model_in_typed_dict_respects_config(self):
|
||||
"""A typed dict's schema and result both use the model's field names."""
|
||||
async def test_model_in_typed_mapping_respects_config(self):
|
||||
"""A typed mapping's schema and result use the model's field names."""
|
||||
|
||||
class Biofile(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=False)
|
||||
|
|
@ -439,131 +419,6 @@ class TestSerializeByAlias:
|
|||
"aliased": {"aliasedValue": "aliased"},
|
||||
}
|
||||
|
||||
async def test_conflicting_structural_alias_contexts_use_permissive_schema(self):
|
||||
"""A reused structural type cannot have one schema for two alias modes."""
|
||||
|
||||
@dataclass
|
||||
class Child:
|
||||
value: Annotated[str, Field(serialization_alias="dataValue")]
|
||||
|
||||
class AliasedParent(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=True)
|
||||
child: Child
|
||||
|
||||
class NamedParent(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=False)
|
||||
child: Child
|
||||
|
||||
class Output(BaseModel):
|
||||
aliased: AliasedParent
|
||||
named: NamedParent
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_output() -> Output:
|
||||
return Output(
|
||||
aliased=AliasedParent(child=Child(value="aliased")),
|
||||
named=NamedParent(child=Child(value="named")),
|
||||
)
|
||||
|
||||
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 tools["get_output"].output_schema == {"type": "object"}
|
||||
assert result.structured_content == {
|
||||
"aliased": {"child": {"dataValue": "aliased"}},
|
||||
"named": {"child": {"value": "named"}},
|
||||
}
|
||||
|
||||
async def test_conflicting_typed_dict_alias_contexts_use_permissive_schema(self):
|
||||
"""TypedDict aliases can also depend on their containing model."""
|
||||
|
||||
class Child(TypedDict):
|
||||
value: Annotated[str, Field(serialization_alias="dataValue")]
|
||||
|
||||
class AliasedParent(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=True)
|
||||
child: Child
|
||||
|
||||
class NamedParent(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=False)
|
||||
child: Child
|
||||
|
||||
class Output(BaseModel):
|
||||
aliased: AliasedParent
|
||||
named: NamedParent
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_output() -> list[Output]:
|
||||
return [
|
||||
Output(
|
||||
aliased=AliasedParent(child={"value": "aliased"}),
|
||||
named=NamedParent(child={"value": "named"}),
|
||||
)
|
||||
]
|
||||
|
||||
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 tools["get_output"].output_schema == {
|
||||
"type": "object",
|
||||
"properties": {"result": {}},
|
||||
"required": ["result"],
|
||||
"x-fastmcp-wrap-result": True,
|
||||
}
|
||||
assert result.structured_content == {
|
||||
"result": [
|
||||
{
|
||||
"aliased": {"child": {"dataValue": "aliased"}},
|
||||
"named": {"child": {"value": "named"}},
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
async def test_configured_dataclass_has_one_alias_context(self):
|
||||
"""A dataclass's own alias config remains schema-safe when reused."""
|
||||
|
||||
@with_config(ConfigDict(serialize_by_alias=True))
|
||||
@dataclass
|
||||
class Child:
|
||||
value: Annotated[str, Field(serialization_alias="dataValue")]
|
||||
|
||||
class AliasedParent(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=True)
|
||||
child: Child
|
||||
|
||||
class NamedParent(BaseModel):
|
||||
model_config = ConfigDict(serialize_by_alias=False)
|
||||
child: Child
|
||||
|
||||
class Output(BaseModel):
|
||||
aliased: AliasedParent
|
||||
named: NamedParent
|
||||
|
||||
mcp = FastMCP()
|
||||
|
||||
@mcp.tool
|
||||
def get_output() -> Output:
|
||||
return Output(
|
||||
aliased=AliasedParent(child=Child(value="aliased")),
|
||||
named=NamedParent(child=Child(value="named")),
|
||||
)
|
||||
|
||||
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 tools["get_output"].output_schema is not None
|
||||
assert result.structured_content == {
|
||||
"aliased": {"child": {"dataValue": "aliased"}},
|
||||
"named": {"child": {"dataValue": "named"}},
|
||||
}
|
||||
|
||||
async def test_dataclass_uses_pydantic_alias_default(self):
|
||||
"""A dataclass schema uses field names when aliases are not enabled."""
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from typing import Annotated, Any
|
|||
|
||||
import pytest
|
||||
from mcp_types import TextContent
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field, with_config
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client.client import Client
|
||||
|
|
@ -726,7 +726,6 @@ async def test_transform_fn_wrapped_result_respects_serialize_by_alias():
|
|||
|
||||
async def test_transform_fn_configured_dataclass_respects_serialize_by_alias():
|
||||
"""A transform uses its return annotation for nested dataclass serialization."""
|
||||
from pydantic import ConfigDict, with_config
|
||||
|
||||
@with_config(ConfigDict(serialize_by_alias=True))
|
||||
@dataclass
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue