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" --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 ## A dataclass

View file

@ -1,19 +1,29 @@
import functools import functools
import inspect import inspect
import json import json
import sys
from collections.abc import Awaitable, Callable, Sequence from collections.abc import Awaitable, Callable, Sequence
from itertools import chain from itertools import chain
from types import GenericAlias 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
import anyio.to_thread import anyio.to_thread
import pydantic_core import pydantic_core
from mcp_types import CallToolResult, ContentBlock, InputRequiredResult, TextContent 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.fields import FieldInfo
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind 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 ( from typing_inspection.introspection import (
UNKNOWN, UNKNOWN,
AnnotationSource, AnnotationSource,
@ -83,10 +93,27 @@ class ArgModelBase(BaseModel):
class FuncMetadata(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)] arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
output_schema: dict[str, Any] | None = 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 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]: 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). """Validate raw arguments into a one-level kwargs dict (no function call).
@ -142,23 +169,29 @@ class FuncMetadata(BaseModel):
""" """
if isinstance(result, InputRequiredResult): if isinstance(result, InputRequiredResult):
return result 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 isinstance(result, CallToolResult):
if self.output_schema is not None: if output_model is not None:
assert self.output_model is not None, "Output model must be set if output schema is defined" self._output_adapter(output_model).validate_python(result.structured_content)
self.output_model.model_validate(result.structured_content)
return result return result
unstructured_content = _convert_to_content(result) unstructured_content = _convert_to_content(result)
if self.output_schema is None: if output_model is None:
return CallToolResult(content=unstructured_content) return CallToolResult(content=unstructured_content)
if self.wrap_output: if self.wrap_output:
result = {"result": result} result = {"result": result}
assert self.output_model is not None, "Output model must be set if output schema is defined" # The tool hands back Python-side names; the wire (and outputSchema) use aliases.
validated = self.output_model.model_validate(result) adapter = self._output_adapter(output_model)
structured_content = validated.model_dump(mode="json", by_alias=True) 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) return CallToolResult(content=unstructured_content, structured_content=structured_content)
@ -238,7 +271,7 @@ def func_metadata(
- BaseModel subclasses (used directly) - BaseModel subclasses (used directly)
- Primitive types (str, int, float, bool, bytes, None) - wrapped in a - Primitive types (str, int, float, bool, bytes, None) - wrapped in a
model with a 'result' field 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 - Dataclasses and other annotated classes - converted to Pydantic models
- Generic types (list, dict, Union, etc.) - wrapped in a model with a 'result' field - 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 - Content blocks (TextContent, EmbeddedResource, ...), Image and Audio, bare or inside a
@ -248,7 +281,9 @@ def func_metadata(
Returns: Returns:
A FuncMetadata object containing: A FuncMetadata object containing:
- arg_model: A Pydantic model representing the function's arguments - 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. - wrap_output: Whether the function result needs to be wrapped in `{"result": ...}` for structured output.
""" """
try: try:
@ -374,30 +409,44 @@ def func_metadata(
# structured_output=True still forces one. # structured_output=True still forces one.
return FuncMetadata(arg_model=arguments_model) return FuncMetadata(arg_model=arguments_model)
output_model, output_schema, wrap_output = _try_create_model_and_schema( output_model, wrap_output = _create_output_model(original_annotation, return_type_expr, func.__name__)
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 # Model creation failed or produced warnings - no structured output
raise InvalidSignature( raise InvalidSignature(
f"Function {func.__name__}: return type {return_type_expr} is not serializable for structured output" f"Function {func.__name__}: return type {return_type_expr} is not serializable for structured output"
) )
return FuncMetadata( return FuncMetadata(arg_model=arguments_model)
arg_model=arguments_model,
output_schema=output_schema,
output_model=output_model,
wrap_output=wrap_output,
)
def _try_create_model_and_schema( def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[type[Any] | None, bool]:
original_annotation: Any, """Pick the type structured output is validated against for the given return annotation.
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.
Args: Args:
original_annotation: The original return annotation (may be wrapped in `Annotated`). 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. func_name: The name of the function.
Returns: Returns:
tuple of (model or None, schema or None, wrap_output) tuple of (model or None, wrap_output)
Model and schema are None if warnings occur or creation fails. Model is None if the type cannot carry structured output.
wrap_output is True if the result needs to be wrapped in {"result": ...} wrap_output is True if the result needs to be wrapped in {"result": ...}
""" """
model = None model: type[Any] | None = None
wrap_output = False wrap_output = False
# First handle special case: None # First handle special case: None
@ -446,9 +495,9 @@ def _try_create_model_and_schema(
if issubclass(type_annotation, BaseModel): if issubclass(type_annotation, BaseModel):
model = type_annotation model = type_annotation
# Case 2: TypedDicts: # Case 2: TypedDicts (pydantic reads qualifiers, totality, docstring and `Annotated` metadata natively)
elif is_typeddict(type_annotation): elif is_typeddict(type_annotation):
model = _create_model_from_typeddict(type_annotation) model = type_annotation
# Case 3: Primitive types that need wrapping # Case 3: Primitive types that need wrapping
elif type_annotation in (str, int, float, bool, bytes, type(None)): 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) model = _create_wrapped_model(func_name, original_annotation)
wrap_output = True wrap_output = True
if model: return model, wrap_output
# 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
_no_default = object() _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) return create_model(cls.__name__, __config__=ConfigDict(from_attributes=True), **model_fields)
def _create_model_from_typeddict(td_type: type[Any]) -> type[BaseModel]: def _pydantic_readable_typeddict(output_model: type[Any]) -> type[Any]:
"""Create a Pydantic model from a TypedDict. """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] = {} def _as_typing_extensions_typeddict(td_type: type[Any]) -> type[Any]: # pragma: lax no cover
for field_name, field_type in type_hints.items(): items: dict[str, Any] = {}
if field_name not in required_keys: for name, hint in get_type_hints(td_type, include_extras=True).items():
# For optional TypedDict fields, set default=None key = inspect_annotation(hint, annotation_source=AnnotationSource.TYPED_DICT)
# This makes them not required in the Pydantic model item: Any = Annotated[(key.type, *key.metadata)] if key.metadata else key.type
# The model should use exclude_unset=True when dumping to get TypedDict semantics if "read_only" in key.qualifiers:
model_fields[field_name] = (field_type, None) item = ReadOnly[item]
else: # pydantic's rule: an explicit qualifier wins over class totality. Needed because a stdlib TypedDict
model_fields[field_name] = field_type # this old computes `__required_keys__` without seeing `typing_extensions` qualifiers.
required = (name in td_type.__required_keys__ or "required" in key.qualifiers) and (
return create_model(td_type.__name__, **model_fields) "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]: def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:

View file

@ -5,17 +5,21 @@
# pyright: reportUnknownLambdaType=false # pyright: reportUnknownLambdaType=false
from collections.abc import Callable from collections.abc import Callable
from dataclasses import dataclass 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 annotated_types
import pytest import pytest
from dirty_equals import IsPartialDict from dirty_equals import IsPartialDict
from mcp_types import CallToolResult, ContentBlock, EmbeddedResource, InputRequiredResult, TextContent 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 import Audio, Image
from mcp.server.mcpserver.exceptions import InvalidSignature 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): class SomeInputModelA(BaseModel):
@ -773,22 +777,28 @@ def test_structured_output_dataclass():
def test_structured_output_typeddict(): def test_structured_output_typeddict():
"""Test structured output with TypedDict return types""" """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): class PersonTypedDictOptional(TypedDict, total=False):
name: str name: Required[str]
age: int 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 return {"name": "Dave"} # Only returning one field to test partial dict
meta = func_metadata(func_returning_typeddict_optional) meta = func_metadata(func_returning_typeddict_optional)
assert meta.output_schema == { assert meta.output_schema == {
"type": "object", "type": "object",
"properties": { "properties": {
"name": {"title": "Name", "type": "string", "default": None}, "name": {"title": "Name", "type": "string"},
"age": {"title": "Age", "type": "integer", "default": None}, "age": {"title": "Age", "type": "integer"},
}, },
"required": ["name"],
"title": "PersonTypedDictOptional", "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) # Test with total=True (all required)
class PersonTypedDictRequired(TypedDict): 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(): def test_structured_output_ordinary_class():
"""Test structured output with ordinary annotated classes""" """Test structured output with ordinary annotated classes"""

View file

@ -44,6 +44,7 @@ from mcp_types import (
from pydantic import BaseModel from pydantic import BaseModel
from starlette.applications import Starlette from starlette.applications import Starlette
from starlette.routing import Mount, Route from starlette.routing import Mount, Route
from typing_extensions import NotRequired, TypedDict
from mcp.client import Client from mcp.client import Client
from mcp.server.context import ServerRequestContext from mcp.server.context import ServerRequestContext
@ -715,6 +716,32 @@ class TestServerTools:
assert "Unknown tool" in content.text 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: class TestServerResources:
async def test_init_with_resources(self): async def test_init_with_resources(self):
def get_text() -> str: def get_text() -> str: