Hand TypedDict tool results to pydantic natively (#3331)

This commit is contained in:
Max 2026-08-20 15:36:21 +01:00 • committed by GitHub
commit 0cee6249ba
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 230 additions and 85 deletions

View file

@ -100,7 +100,7 @@ Not every shape deserves a class. A `TypedDict` produces the same schema:
--8<-- "docs_src/structured_output/tutorial003.py"
```
A `TypedDict` is a plain `dict` at runtime, so that is what you build and return. The schema, the validation, and `structured_content` are identical to the `BaseModel` version (minus the descriptions, which `TypedDict` has no place for).
A `TypedDict` is a plain `dict` at runtime, so that is what you build and return. The schema, the validation, and `structured_content` follow the same rules as the `BaseModel` version: add a class docstring or `Annotated[..., Field(description=...)]` and they become the descriptions, and a `NotRequired` key you leave out of the dict stays out of `structured_content`.
## A dataclass

View file

@ -1,19 +1,29 @@
import functools
import inspect
import json
import sys
from collections.abc import Awaitable, Callable, Sequence
from itertools import chain
from types import GenericAlias
from typing import Annotated, Any, Union, cast, get_args, get_origin, get_type_hints
from typing import Annotated, Any, Union, cast, get_args, get_origin
import anyio
import anyio.to_thread
import pydantic_core
from mcp_types import CallToolResult, ContentBlock, InputRequiredResult, TextContent
from pydantic import BaseModel, ConfigDict, Field, PydanticUserError, WithJsonSchema, create_model
from pydantic import (
BaseModel,
ConfigDict,
Field,
PrivateAttr,
PydanticUserError,
TypeAdapter,
WithJsonSchema,
create_model,
)
from pydantic.fields import FieldInfo
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind
from typing_extensions import is_typeddict
from typing_extensions import NotRequired, ReadOnly, TypedDict, get_type_hints, is_typeddict
from typing_inspection.introspection import (
UNKNOWN,
AnnotationSource,
@ -83,10 +93,27 @@ class ArgModelBase(BaseModel):
class FuncMetadata(BaseModel):
"""A tool function's argument model plus, for structured output, the published `output_schema` and the
`output_model` results are validated against. Constructing one with an `output_model` and no schema derives
the schema (and raises if pydantic can't); the fields are read live, so clearing or reassigning them later
takes effect on the next call."""
arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
output_schema: dict[str, Any] | None = None
output_model: Annotated[type[BaseModel], WithJsonSchema(None)] | None = None
output_model: Annotated[type[Any], WithJsonSchema(None)] | None = None
wrap_output: bool = False
_adapter: tuple[type[Any], TypeAdapter[Any]] | None = PrivateAttr(default=None)
def model_post_init(self, context: Any, /) -> None:
if self.output_model is not None and self.output_schema is None:
# StrictJsonSchema raises instead of warning, so an unserializable return type fails construction.
self.output_schema = self._output_adapter(self.output_model).json_schema(schema_generator=StrictJsonSchema)
def _output_adapter(self, output_model: type[Any]) -> TypeAdapter[Any]:
"""The validator/serializer for `output_model`, built once and rebuilt only if the field is reassigned."""
if self._adapter is None or self._adapter[0] is not output_model:
self._adapter = (output_model, TypeAdapter(_pydantic_readable_typeddict(output_model)))
return self._adapter[1]
def validate_arguments(self, arguments_to_validate: dict[str, Any]) -> dict[str, Any]:
"""Validate raw arguments into a one-level kwargs dict (no function call).
@ -142,23 +169,29 @@ class FuncMetadata(BaseModel):
"""
if isinstance(result, InputRequiredResult):
return result
# A schema published without a model (hand-built metadata) is advertised but not validated here.
output_model = self.output_model if self.output_schema is not None else None
if isinstance(result, CallToolResult):
if self.output_schema is not None:
assert self.output_model is not None, "Output model must be set if output schema is defined"
self.output_model.model_validate(result.structured_content)
if output_model is not None:
self._output_adapter(output_model).validate_python(result.structured_content)
return result
unstructured_content = _convert_to_content(result)
if self.output_schema is None:
if output_model is None:
return CallToolResult(content=unstructured_content)
if self.wrap_output:
result = {"result": result}
assert self.output_model is not None, "Output model must be set if output schema is defined"
validated = self.output_model.model_validate(result)
structured_content = validated.model_dump(mode="json", by_alias=True)
# The tool hands back Python-side names; the wire (and outputSchema) use aliases.
adapter = self._output_adapter(output_model)
validated = adapter.validate_python(result, by_alias=True, by_name=True)
if isinstance(validated, BaseModel):
# Dump via the instance so a returned subclass keeps its own fields.
structured_content = validated.model_dump(mode="json", by_alias=True)
else:
structured_content = adapter.dump_python(validated, mode="json", by_alias=True)
return CallToolResult(content=unstructured_content, structured_content=structured_content)
@ -238,7 +271,7 @@ def func_metadata(
- BaseModel subclasses (used directly)
- Primitive types (str, int, float, bool, bytes, None) - wrapped in a
model with a 'result' field
- TypedDict - converted to a Pydantic model with same fields
- TypedDict - used directly
- Dataclasses and other annotated classes - converted to Pydantic models
- Generic types (list, dict, Union, etc.) - wrapped in a model with a 'result' field
- Content blocks (TextContent, EmbeddedResource, ...), Image and Audio, bare or inside a
@ -248,7 +281,9 @@ def func_metadata(
Returns:
A FuncMetadata object containing:
- arg_model: A Pydantic model representing the function's arguments
- output_model: A Pydantic model for the return type if the output is structured
- output_schema: The published JSON schema for structured output, or None if the output is unstructured
- output_model: The type structured output is validated against: the declared BaseModel or TypedDict,
or a synthesized model for wrapped, `dict[str, T]` and annotated-class returns
- wrap_output: Whether the function result needs to be wrapped in `{"result": ...}` for structured output.
"""
try:
@ -374,30 +409,44 @@ def func_metadata(
# structured_output=True still forces one.
return FuncMetadata(arg_model=arguments_model)
output_model, output_schema, wrap_output = _try_create_model_and_schema(
original_annotation, return_type_expr, func.__name__
)
output_model, wrap_output = _create_output_model(original_annotation, return_type_expr, func.__name__)
if output_model is None and structured_output is True:
if output_model is not None:
try:
# FuncMetadata builds the validator and schema on construction, so an unsupported return type
# surfaces here, at registration, rather than on the first call.
return FuncMetadata(arg_model=arguments_model, output_model=output_model, wrap_output=wrap_output)
except (
PydanticUserError,
ForbiddenQualifier,
NameError,
TypeError,
ValueError,
pydantic_core.SchemaError,
pydantic_core.ValidationError,
) as e:
# These are expected errors when a type can't be converted to a Pydantic schema
# PydanticUserError: When Pydantic can't handle the type (e.g. PydanticInvalidForJsonSchema);
# subclasses TypeError on pydantic <2.13 and RuntimeError on pydantic >=2.13
# ForbiddenQualifier, NameError: an invalid qualifier or unresolvable annotation on a TypedDict key,
# met while rebuilding a stdlib TypedDict below 3.12 (pydantic reports both as PydanticUserError)
# ValueError: When there are issues with the type definition (including our custom warnings);
# arrives wrapped in a ValidationError when raised during FuncMetadata construction
# SchemaError: When Pydantic can't build a schema
# ValidationError: When validation fails
logger.info(f"Cannot create schema for type {return_type_expr} in {func.__name__}: {type(e).__name__}: {e}")
if structured_output is True:
# Model creation failed or produced warnings - no structured output
raise InvalidSignature(
f"Function {func.__name__}: return type {return_type_expr} is not serializable for structured output"
)
return FuncMetadata(
arg_model=arguments_model,
output_schema=output_schema,
output_model=output_model,
wrap_output=wrap_output,
)
return FuncMetadata(arg_model=arguments_model)
def _try_create_model_and_schema(
original_annotation: Any,
type_expr: Any,
func_name: str,
) -> tuple[type[BaseModel] | None, dict[str, Any] | None, bool]:
"""Try to create a model and schema for the given annotation without warnings.
def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[type[Any] | None, bool]:
"""Pick the type structured output is validated against for the given return annotation.
Args:
original_annotation: The original return annotation (may be wrapped in `Annotated`).
@ -406,11 +455,11 @@ def _try_create_model_and_schema(
func_name: The name of the function.
Returns:
tuple of (model or None, schema or None, wrap_output)
Model and schema are None if warnings occur or creation fails.
tuple of (model or None, wrap_output)
Model is None if the type cannot carry structured output.
wrap_output is True if the result needs to be wrapped in {"result": ...}
"""
model = None
model: type[Any] | None = None
wrap_output = False
# First handle special case: None
@ -446,9 +495,9 @@ def _try_create_model_and_schema(
if issubclass(type_annotation, BaseModel):
model = type_annotation
# Case 2: TypedDicts:
# Case 2: TypedDicts (pydantic reads qualifiers, totality, docstring and `Annotated` metadata natively)
elif is_typeddict(type_annotation):
model = _create_model_from_typeddict(type_annotation)
model = type_annotation
# Case 3: Primitive types that need wrapping
elif type_annotation in (str, int, float, bool, bytes, type(None)):
@ -470,30 +519,7 @@ def _try_create_model_and_schema(
model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True
if model:
# If we successfully created a model, try to get its schema
# Use StrictJsonSchema to raise exceptions instead of warnings
try:
schema = model.model_json_schema(schema_generator=StrictJsonSchema)
except (
PydanticUserError,
TypeError,
ValueError,
pydantic_core.SchemaError,
pydantic_core.ValidationError,
) as e:
# These are expected errors when a type can't be converted to a Pydantic schema
# PydanticUserError: When Pydantic can't handle the type (e.g. PydanticInvalidForJsonSchema);
# subclasses TypeError on pydantic <2.13 and RuntimeError on pydantic >=2.13
# ValueError: When there are issues with the type definition (including our custom warnings)
# SchemaError: When Pydantic can't build a schema
# ValidationError: When validation fails
logger.info(f"Cannot create schema for type {type_expr} in {func_name}: {type(e).__name__}: {e}")
return None, None, False
return model, schema, wrap_output
return None, None, False
return model, wrap_output
_no_default = object()
@ -523,25 +549,35 @@ def _create_model_from_class(cls: type[Any], type_hints: dict[str, Any]) -> type
return create_model(cls.__name__, __config__=ConfigDict(from_attributes=True), **model_fields)
def _create_model_from_typeddict(td_type: type[Any]) -> type[BaseModel]:
"""Create a Pydantic model from a TypedDict.
def _pydantic_readable_typeddict(output_model: type[Any]) -> type[Any]:
"""pydantic refuses `typing.TypedDict` below Python 3.12 (it needs `__orig_bases__`); rebuild such a return
type as an equivalent `typing_extensions.TypedDict` so tool authors don't have to know. Only the class itself
(its keys, docstring and own config) is rebuilt: stdlib TypedDicts nested inside it, or config inherited from
one, still need `typing_extensions` there. Delete with 3.11 support."""
if sys.version_info >= (3, 12) or not is_typeddict(output_model) or type(output_model).__module__ != "typing":
return output_model
return _as_typing_extensions_typeddict(output_model) # pragma: lax no cover
The created model will have the same name and fields as the TypedDict.
"""
type_hints = get_type_hints(td_type)
required_keys = getattr(td_type, "__required_keys__", set(type_hints.keys()))
model_fields: dict[str, Any] = {}
for field_name, field_type in type_hints.items():
if field_name not in required_keys:
# For optional TypedDict fields, set default=None
# This makes them not required in the Pydantic model
# The model should use exclude_unset=True when dumping to get TypedDict semantics
model_fields[field_name] = (field_type, None)
else:
model_fields[field_name] = field_type
return create_model(td_type.__name__, **model_fields)
def _as_typing_extensions_typeddict(td_type: type[Any]) -> type[Any]: # pragma: lax no cover
items: dict[str, Any] = {}
for name, hint in get_type_hints(td_type, include_extras=True).items():
key = inspect_annotation(hint, annotation_source=AnnotationSource.TYPED_DICT)
item: Any = Annotated[(key.type, *key.metadata)] if key.metadata else key.type
if "read_only" in key.qualifiers:
item = ReadOnly[item]
# pydantic's rule: an explicit qualifier wins over class totality. Needed because a stdlib TypedDict
# this old computes `__required_keys__` without seeing `typing_extensions` qualifiers.
required = (name in td_type.__required_keys__ or "required" in key.qualifiers) and (
"not_required" not in key.qualifiers
)
items[name] = item if required else NotRequired[item]
# The functional form, spelled so type checkers don't try to evaluate it statically.
rebuilt = cast("Callable[[str, dict[str, Any]], type[Any]]", TypedDict)(td_type.__name__, items)
for attr in ("__doc__", "__module__", "__qualname__", "__pydantic_config__"):
if hasattr(td_type, attr):
setattr(rebuilt, attr, getattr(td_type, attr))
return rebuilt
def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:

View file

@ -5,17 +5,21 @@
# pyright: reportUnknownLambdaType=false
from collections.abc import Callable
from dataclasses import dataclass
from typing import Annotated, Any, Final, NamedTuple, TypedDict
from typing import TYPE_CHECKING, Annotated, Any, Final, NamedTuple, TypedDict
import annotated_types
import pytest
from dirty_equals import IsPartialDict
from mcp_types import CallToolResult, ContentBlock, EmbeddedResource, InputRequiredResult, TextContent
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, ValidationError
from typing_extensions import NotRequired, ReadOnly, Required
from mcp.server.mcpserver import Audio, Image
from mcp.server.mcpserver.exceptions import InvalidSignature
from mcp.server.mcpserver.utilities.func_metadata import func_metadata
from mcp.server.mcpserver.utilities.func_metadata import ArgModelBase, FuncMetadata, func_metadata
if TYPE_CHECKING:
from decimal import Decimal
class SomeInputModelA(BaseModel):
@ -773,22 +777,28 @@ def test_structured_output_dataclass():
def test_structured_output_typeddict():
"""Test structured output with TypedDict return types"""
# stdlib TypedDict with a qualifier: exercises the typing_extensions rebuild below Python 3.12
class PersonTypedDictOptional(TypedDict, total=False):
name: str
name: Required[str]
age: int
def func_returning_typeddict_optional() -> PersonTypedDictOptional: # pragma: no cover
def func_returning_typeddict_optional() -> PersonTypedDictOptional:
return {"name": "Dave"} # Only returning one field to test partial dict
meta = func_metadata(func_returning_typeddict_optional)
assert meta.output_schema == {
"type": "object",
"properties": {
"name": {"title": "Name", "type": "string", "default": None},
"age": {"title": "Age", "type": "integer", "default": None},
"name": {"title": "Name", "type": "string"},
"age": {"title": "Age", "type": "integer"},
},
"required": ["name"],
"title": "PersonTypedDictOptional",
}
# An optional key the tool leaves out is absent, not null, so it validates against the schema above
result = meta.convert_result(func_returning_typeddict_optional())
assert isinstance(result, CallToolResult)
assert result.structured_content == {"name": "Dave"}
# Test with total=True (all required)
class PersonTypedDictRequired(TypedDict):
@ -812,6 +822,78 @@ def test_structured_output_typeddict():
}
def test_structured_output_typeddict_qualifiers_and_metadata():
"""PEP 655/705 qualifiers register on every supported Python and decide `required`; the docstring and
`Annotated` field metadata reach the schema like they do for a BaseModel, and an alias names the wire key.
A stdlib TypedDict, so below 3.12 all of this goes through the typing_extensions rebuild."""
class Forecast(TypedDict, total=False):
"""Tomorrow's weather."""
city: Required[Annotated[str, Field(alias="cityName", description="City name")]]
high: ReadOnly[Required[float]]
low: float
summary: NotRequired[Annotated[str, Field(max_length=80)]]
def forecast() -> Forecast:
return {"city": "Berlin", "high": 21.5}
with pytest.warns(UserWarning, match="ReadOnly"): # pydantic notes it won't enforce ReadOnly
meta = func_metadata(forecast)
result = meta.convert_result(forecast())
assert isinstance(result, CallToolResult)
assert result.structured_content == {"cityName": "Berlin", "high": 21.5}
assert meta.output_schema == {
"type": "object",
"title": "Forecast",
"description": "Tomorrow's weather.",
"properties": {
"cityName": {"title": "Cityname", "type": "string", "description": "City name"},
"high": {"title": "High", "type": "number"},
"low": {"title": "Low", "type": "number"},
"summary": {"title": "Summary", "type": "string", "maxLength": 80},
},
"required": ["cityName", "high"],
}
def test_func_metadata_built_by_hand_keeps_a_given_output_schema():
"""A hand-built FuncMetadata publishes the schema it was given but still validates results against output_model."""
meta = FuncMetadata(arg_model=ArgModelBase, output_model=SomeInputModelB.InnerModel, output_schema={"x": 1})
assert meta.output_schema == {"x": 1}
with pytest.raises(ValidationError):
meta.convert_result({"x": "not an int"})
def test_func_metadata_output_fields_are_read_live():
"""Code in the wild switches structured output off or on by assigning the fields after registration."""
def total() -> int:
return 3
meta = func_metadata(total)
meta.output_schema = None
off = meta.convert_result(total())
meta.output_schema, meta.output_model, meta.wrap_output = {"type": "object"}, SomeInputModelB.InnerModel, False
on = meta.convert_result({"x": 3})
assert isinstance(off, CallToolResult) and isinstance(on, CallToolResult)
assert (off.structured_content, on.structured_content) == (None, {"x": 3})
def test_structured_output_typeddict_with_unresolvable_annotation_is_unstructured():
"""An annotation only importable under TYPE_CHECKING degrades the same way on every supported Python."""
class Report(TypedDict):
total: "Decimal"
def report() -> Report: # pragma: no cover
raise NotImplementedError
assert func_metadata(report).output_schema is None
with pytest.raises(InvalidSignature):
func_metadata(report, structured_output=True)
def test_structured_output_ordinary_class():
"""Test structured output with ordinary annotated classes"""

View file

@ -44,6 +44,7 @@ from mcp_types import (
from pydantic import BaseModel
from starlette.applications import Starlette
from starlette.routing import Mount, Route
from typing_extensions import NotRequired, TypedDict
from mcp.client import Client
from mcp.server.context import ServerRequestContext
@ -715,6 +716,32 @@ class TestServerTools:
assert "Unknown tool" in content.text
async def test_typeddict_tool_omitting_optional_keys_passes_client_validation():
"""The client validates structured content against the tool's output schema, so a `NotRequired`
key the tool leaves out must be absent from `structured_content` rather than null."""
class Person(TypedDict):
name: str
age: NotRequired[int]
mcp = MCPServer()
@mcp.tool()
def get_person() -> Person:
return {"name": "Dave"}
async with Client(mcp) as client:
(tool,) = (await client.list_tools()).tools
assert tool.output_schema == {
"type": "object",
"title": "Person",
"properties": {"name": {"title": "Name", "type": "string"}, "age": {"title": "Age", "type": "integer"}},
"required": ["name"],
}
result = await client.call_tool("get_person", {})
assert result.structured_content == {"name": "Dave"}
class TestServerResources:
async def test_init_with_resources(self):
def get_text() -> str: