Skip to content

Commit 0cee624

Browse files
maxisbeymaxisbey
authored
Hand TypedDict tool results to pydantic natively (#3331)
1 parent 0d92192 commit 0cee624

4 files changed

Lines changed: 230 additions & 85 deletions

File tree

‎docs/servers/structured-output.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@ Not every shape deserves a class. A `TypedDict` produces the same schema:
100100
--8<-- "docs_src/structured_output/tutorial003.py"
101101
```
102102

103-
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).
103+
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`.
104104

105105
## A dataclass
106106

‎src/mcp/server/mcpserver/utilities/func_metadata.py‎

Lines changed: 113 additions & 77 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,29 @@
11
import functools
22
import inspect
33
import json
4+
import sys
45
from collections.abc import Awaitable, Callable, Sequence
56
from itertools import chain
67
from types import GenericAlias
7-
from typing import Annotated, Any, Union, cast, get_args, get_origin, get_type_hints
8+
from typing import Annotated, Any, Union, cast, get_args, get_origin
89

910
import anyio
1011
import anyio.to_thread
1112
import pydantic_core
1213
from mcp_types import CallToolResult, ContentBlock, InputRequiredResult, TextContent
13-
from pydantic import BaseModel, ConfigDict, Field, PydanticUserError, WithJsonSchema, create_model
14+
from pydantic import (
15+
BaseModel,
16+
ConfigDict,
17+
Field,
18+
PrivateAttr,
19+
PydanticUserError,
20+
TypeAdapter,
21+
WithJsonSchema,
22+
create_model,
23+
)
1424
from pydantic.fields import FieldInfo
1525
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind
16-
from typing_extensions import is_typeddict
26+
from typing_extensions import NotRequired, ReadOnly, TypedDict, get_type_hints, is_typeddict
1727
from typing_inspection.introspection import (
1828
UNKNOWN,
1929
AnnotationSource,
@@ -83,10 +93,27 @@ def model_dump_one_level(self) -> dict[str, Any]:
8393

8494

8595
class FuncMetadata(BaseModel):
96+
"""A tool function's argument model plus, for structured output, the published `output_schema` and the
97+
`output_model` results are validated against. Constructing one with an `output_model` and no schema derives
98+
the schema (and raises if pydantic can't); the fields are read live, so clearing or reassigning them later
99+
takes effect on the next call."""
100+
86101
arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
87102
output_schema: dict[str, Any] | None = None
88-
output_model: Annotated[type[BaseModel], WithJsonSchema(None)] | None = None
103+
output_model: Annotated[type[Any], WithJsonSchema(None)] | None = None
89104
wrap_output: bool = False
105+
_adapter: tuple[type[Any], TypeAdapter[Any]] | None = PrivateAttr(default=None)
106+
107+
def model_post_init(self, context: Any, /) -> None:
108+
if self.output_model is not None and self.output_schema is None:
109+
# StrictJsonSchema raises instead of warning, so an unserializable return type fails construction.
110+
self.output_schema = self._output_adapter(self.output_model).json_schema(schema_generator=StrictJsonSchema)
111+
112+
def _output_adapter(self, output_model: type[Any]) -> TypeAdapter[Any]:
113+
"""The validator/serializer for `output_model`, built once and rebuilt only if the field is reassigned."""
114+
if self._adapter is None or self._adapter[0] is not output_model:
115+
self._adapter = (output_model, TypeAdapter(_pydantic_readable_typeddict(output_model)))
116+
return self._adapter[1]
90117

91118
def validate_arguments(self, arguments_to_validate: dict[str, Any]) -> dict[str, Any]:
92119
"""Validate raw arguments into a one-level kwargs dict (no function call).
@@ -142,23 +169,29 @@ def convert_result(self, result: Any) -> CallToolResult | InputRequiredResult:
142169
"""
143170
if isinstance(result, InputRequiredResult):
144171
return result
172+
# A schema published without a model (hand-built metadata) is advertised but not validated here.
173+
output_model = self.output_model if self.output_schema is not None else None
145174
if isinstance(result, CallToolResult):
146-
if self.output_schema is not None:
147-
assert self.output_model is not None, "Output model must be set if output schema is defined"
148-
self.output_model.model_validate(result.structured_content)
175+
if output_model is not None:
176+
self._output_adapter(output_model).validate_python(result.structured_content)
149177
return result
150178

151179
unstructured_content = _convert_to_content(result)
152180

153-
if self.output_schema is None:
181+
if output_model is None:
154182
return CallToolResult(content=unstructured_content)
155183

156184
if self.wrap_output:
157185
result = {"result": result}
158186

159-
assert self.output_model is not None, "Output model must be set if output schema is defined"
160-
validated = self.output_model.model_validate(result)
161-
structured_content = validated.model_dump(mode="json", by_alias=True)
187+
# The tool hands back Python-side names; the wire (and outputSchema) use aliases.
188+
adapter = self._output_adapter(output_model)
189+
validated = adapter.validate_python(result, by_alias=True, by_name=True)
190+
if isinstance(validated, BaseModel):
191+
# Dump via the instance so a returned subclass keeps its own fields.
192+
structured_content = validated.model_dump(mode="json", by_alias=True)
193+
else:
194+
structured_content = adapter.dump_python(validated, mode="json", by_alias=True)
162195

163196
return CallToolResult(content=unstructured_content, structured_content=structured_content)
164197

@@ -238,7 +271,7 @@ def func_metadata(
238271
- BaseModel subclasses (used directly)
239272
- Primitive types (str, int, float, bool, bytes, None) - wrapped in a
240273
model with a 'result' field
241-
- TypedDict - converted to a Pydantic model with same fields
274+
- TypedDict - used directly
242275
- Dataclasses and other annotated classes - converted to Pydantic models
243276
- Generic types (list, dict, Union, etc.) - wrapped in a model with a 'result' field
244277
- Content blocks (TextContent, EmbeddedResource, ...), Image and Audio, bare or inside a
@@ -248,7 +281,9 @@ def func_metadata(
248281
Returns:
249282
A FuncMetadata object containing:
250283
- arg_model: A Pydantic model representing the function's arguments
251-
- output_model: A Pydantic model for the return type if the output is structured
284+
- output_schema: The published JSON schema for structured output, or None if the output is unstructured
285+
- output_model: The type structured output is validated against: the declared BaseModel or TypedDict,
286+
or a synthesized model for wrapped, `dict[str, T]` and annotated-class returns
252287
- wrap_output: Whether the function result needs to be wrapped in `{"result": ...}` for structured output.
253288
"""
254289
try:
@@ -374,30 +409,44 @@ def func_metadata(
374409
# structured_output=True still forces one.
375410
return FuncMetadata(arg_model=arguments_model)
376411

377-
output_model, output_schema, wrap_output = _try_create_model_and_schema(
378-
original_annotation, return_type_expr, func.__name__
379-
)
412+
output_model, wrap_output = _create_output_model(original_annotation, return_type_expr, func.__name__)
380413

381-
if output_model is None and structured_output is True:
414+
if output_model is not None:
415+
try:
416+
# FuncMetadata builds the validator and schema on construction, so an unsupported return type
417+
# surfaces here, at registration, rather than on the first call.
418+
return FuncMetadata(arg_model=arguments_model, output_model=output_model, wrap_output=wrap_output)
419+
except (
420+
PydanticUserError,
421+
ForbiddenQualifier,
422+
NameError,
423+
TypeError,
424+
ValueError,
425+
pydantic_core.SchemaError,
426+
pydantic_core.ValidationError,
427+
) as e:
428+
# These are expected errors when a type can't be converted to a Pydantic schema
429+
# PydanticUserError: When Pydantic can't handle the type (e.g. PydanticInvalidForJsonSchema);
430+
# subclasses TypeError on pydantic <2.13 and RuntimeError on pydantic >=2.13
431+
# ForbiddenQualifier, NameError: an invalid qualifier or unresolvable annotation on a TypedDict key,
432+
# met while rebuilding a stdlib TypedDict below 3.12 (pydantic reports both as PydanticUserError)
433+
# ValueError: When there are issues with the type definition (including our custom warnings);
434+
# arrives wrapped in a ValidationError when raised during FuncMetadata construction
435+
# SchemaError: When Pydantic can't build a schema
436+
# ValidationError: When validation fails
437+
logger.info(f"Cannot create schema for type {return_type_expr} in {func.__name__}: {type(e).__name__}: {e}")
438+
439+
if structured_output is True:
382440
# Model creation failed or produced warnings - no structured output
383441
raise InvalidSignature(
384442
f"Function {func.__name__}: return type {return_type_expr} is not serializable for structured output"
385443
)
386444

387-
return FuncMetadata(
388-
arg_model=arguments_model,
389-
output_schema=output_schema,
390-
output_model=output_model,
391-
wrap_output=wrap_output,
392-
)
445+
return FuncMetadata(arg_model=arguments_model)
393446

394447

395-
def _try_create_model_and_schema(
396-
original_annotation: Any,
397-
type_expr: Any,
398-
func_name: str,
399-
) -> tuple[type[BaseModel] | None, dict[str, Any] | None, bool]:
400-
"""Try to create a model and schema for the given annotation without warnings.
448+
def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[type[Any] | None, bool]:
449+
"""Pick the type structured output is validated against for the given return annotation.
401450
402451
Args:
403452
original_annotation: The original return annotation (may be wrapped in `Annotated`).
@@ -406,11 +455,11 @@ def _try_create_model_and_schema(
406455
func_name: The name of the function.
407456
408457
Returns:
409-
tuple of (model or None, schema or None, wrap_output)
410-
Model and schema are None if warnings occur or creation fails.
458+
tuple of (model or None, wrap_output)
459+
Model is None if the type cannot carry structured output.
411460
wrap_output is True if the result needs to be wrapped in {"result": ...}
412461
"""
413-
model = None
462+
model: type[Any] | None = None
414463
wrap_output = False
415464

416465
# First handle special case: None
@@ -446,9 +495,9 @@ def _try_create_model_and_schema(
446495
if issubclass(type_annotation, BaseModel):
447496
model = type_annotation
448497

449-
# Case 2: TypedDicts:
498+
# Case 2: TypedDicts (pydantic reads qualifiers, totality, docstring and `Annotated` metadata natively)
450499
elif is_typeddict(type_annotation):
451-
model = _create_model_from_typeddict(type_annotation)
500+
model = type_annotation
452501

453502
# Case 3: Primitive types that need wrapping
454503
elif type_annotation in (str, int, float, bool, bytes, type(None)):
@@ -470,30 +519,7 @@ def _try_create_model_and_schema(
470519
model = _create_wrapped_model(func_name, original_annotation)
471520
wrap_output = True
472521

473-
if model:
474-
# If we successfully created a model, try to get its schema
475-
# Use StrictJsonSchema to raise exceptions instead of warnings
476-
try:
477-
schema = model.model_json_schema(schema_generator=StrictJsonSchema)
478-
except (
479-
PydanticUserError,
480-
TypeError,
481-
ValueError,
482-
pydantic_core.SchemaError,
483-
pydantic_core.ValidationError,
484-
) as e:
485-
# These are expected errors when a type can't be converted to a Pydantic schema
486-
# PydanticUserError: When Pydantic can't handle the type (e.g. PydanticInvalidForJsonSchema);
487-
# subclasses TypeError on pydantic <2.13 and RuntimeError on pydantic >=2.13
488-
# ValueError: When there are issues with the type definition (including our custom warnings)
489-
# SchemaError: When Pydantic can't build a schema
490-
# ValidationError: When validation fails
491-
logger.info(f"Cannot create schema for type {type_expr} in {func_name}: {type(e).__name__}: {e}")
492-
return None, None, False
493-
494-
return model, schema, wrap_output
495-
496-
return None, None, False
522+
return model, wrap_output
497523

498524

499525
_no_default = object()
@@ -523,25 +549,35 @@ def _create_model_from_class(cls: type[Any], type_hints: dict[str, Any]) -> type
523549
return create_model(cls.__name__, __config__=ConfigDict(from_attributes=True), **model_fields)
524550

525551

526-
def _create_model_from_typeddict(td_type: type[Any]) -> type[BaseModel]:
527-
"""Create a Pydantic model from a TypedDict.
528-
529-
The created model will have the same name and fields as the TypedDict.
530-
"""
531-
type_hints = get_type_hints(td_type)
532-
required_keys = getattr(td_type, "__required_keys__", set(type_hints.keys()))
533-
534-
model_fields: dict[str, Any] = {}
535-
for field_name, field_type in type_hints.items():
536-
if field_name not in required_keys:
537-
# For optional TypedDict fields, set default=None
538-
# This makes them not required in the Pydantic model
539-
# The model should use exclude_unset=True when dumping to get TypedDict semantics
540-
model_fields[field_name] = (field_type, None)
541-
else:
542-
model_fields[field_name] = field_type
543-
544-
return create_model(td_type.__name__, **model_fields)
552+
def _pydantic_readable_typeddict(output_model: type[Any]) -> type[Any]:
553+
"""pydantic refuses `typing.TypedDict` below Python 3.12 (it needs `__orig_bases__`); rebuild such a return
554+
type as an equivalent `typing_extensions.TypedDict` so tool authors don't have to know. Only the class itself
555+
(its keys, docstring and own config) is rebuilt: stdlib TypedDicts nested inside it, or config inherited from
556+
one, still need `typing_extensions` there. Delete with 3.11 support."""
557+
if sys.version_info >= (3, 12) or not is_typeddict(output_model) or type(output_model).__module__ != "typing":
558+
return output_model
559+
return _as_typing_extensions_typeddict(output_model) # pragma: lax no cover
560+
561+
562+
def _as_typing_extensions_typeddict(td_type: type[Any]) -> type[Any]: # pragma: lax no cover
563+
items: dict[str, Any] = {}
564+
for name, hint in get_type_hints(td_type, include_extras=True).items():
565+
key = inspect_annotation(hint, annotation_source=AnnotationSource.TYPED_DICT)
566+
item: Any = Annotated[(key.type, *key.metadata)] if key.metadata else key.type
567+
if "read_only" in key.qualifiers:
568+
item = ReadOnly[item]
569+
# pydantic's rule: an explicit qualifier wins over class totality. Needed because a stdlib TypedDict
570+
# this old computes `__required_keys__` without seeing `typing_extensions` qualifiers.
571+
required = (name in td_type.__required_keys__ or "required" in key.qualifiers) and (
572+
"not_required" not in key.qualifiers
573+
)
574+
items[name] = item if required else NotRequired[item]
575+
# The functional form, spelled so type checkers don't try to evaluate it statically.
576+
rebuilt = cast("Callable[[str, dict[str, Any]], type[Any]]", TypedDict)(td_type.__name__, items)
577+
for attr in ("__doc__", "__module__", "__qualname__", "__pydantic_config__"):
578+
if hasattr(td_type, attr):
579+
setattr(rebuilt, attr, getattr(td_type, attr))
580+
return rebuilt
545581

546582

547583
def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:

0 commit comments

Comments
 (0)