feat: Add structured output support for tool functions (#993)
This commit is contained in:
parent
c8bbfc034d
commit
43bb24f62b
5 changed files with 1774 additions and 61 deletions
|
|
@ -1,8 +1,10 @@
|
|||
import logging
|
||||
from datetime import timedelta
|
||||
from typing import Any, Protocol
|
||||
|
||||
import anyio.lowlevel
|
||||
from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
|
||||
from jsonschema import SchemaError, ValidationError, validate
|
||||
from pydantic import AnyUrl, TypeAdapter
|
||||
|
||||
import mcp.types as types
|
||||
|
|
@ -13,6 +15,9 @@ from mcp.shared.version import SUPPORTED_PROTOCOL_VERSIONS
|
|||
|
||||
DEFAULT_CLIENT_INFO = types.Implementation(name="mcp", version="0.1.0")
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logger = logging.getLogger("client")
|
||||
|
||||
|
||||
class SamplingFnT(Protocol):
|
||||
async def __call__(
|
||||
|
|
@ -128,6 +133,7 @@ class ClientSession(
|
|||
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
|
||||
self._logging_callback = logging_callback or _default_logging_callback
|
||||
self._message_handler = message_handler or _default_message_handler
|
||||
self._tool_output_schemas: dict[str, dict[str, Any] | None] = {}
|
||||
|
||||
async def initialize(self) -> types.InitializeResult:
|
||||
sampling = types.SamplingCapability() if self._sampling_callback is not _default_sampling_callback else None
|
||||
|
|
@ -285,7 +291,7 @@ class ClientSession(
|
|||
) -> types.CallToolResult:
|
||||
"""Send a tools/call request with optional progress callback support."""
|
||||
|
||||
return await self.send_request(
|
||||
result = await self.send_request(
|
||||
types.ClientRequest(
|
||||
types.CallToolRequest(
|
||||
method="tools/call",
|
||||
|
|
@ -300,6 +306,33 @@ class ClientSession(
|
|||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
if not result.isError:
|
||||
await self._validate_tool_result(name, result)
|
||||
|
||||
return result
|
||||
|
||||
async def _validate_tool_result(self, name: str, result: types.CallToolResult) -> None:
|
||||
"""Validate the structured content of a tool result against its output schema."""
|
||||
if name not in self._tool_output_schemas:
|
||||
# refresh output schema cache
|
||||
await self.list_tools()
|
||||
|
||||
output_schema = None
|
||||
if name in self._tool_output_schemas:
|
||||
output_schema = self._tool_output_schemas.get(name)
|
||||
else:
|
||||
logger.warning(f"Tool {name} not listed by server, cannot validate any structured content")
|
||||
|
||||
if output_schema is not None:
|
||||
if result.structuredContent is None:
|
||||
raise RuntimeError(f"Tool {name} has an output schema but did not return structured content")
|
||||
try:
|
||||
validate(result.structuredContent, output_schema)
|
||||
except ValidationError as e:
|
||||
raise RuntimeError(f"Invalid structured content returned by tool {name}: {e}")
|
||||
except SchemaError as e:
|
||||
raise RuntimeError(f"Invalid schema for tool {name}: {e}")
|
||||
|
||||
async def list_prompts(self, cursor: str | None = None) -> types.ListPromptsResult:
|
||||
"""Send a prompts/list request."""
|
||||
return await self.send_request(
|
||||
|
|
@ -351,7 +384,7 @@ class ClientSession(
|
|||
|
||||
async def list_tools(self, cursor: str | None = None) -> types.ListToolsResult:
|
||||
"""Send a tools/list request."""
|
||||
return await self.send_request(
|
||||
result = await self.send_request(
|
||||
types.ClientRequest(
|
||||
types.ListToolsRequest(
|
||||
method="tools/list",
|
||||
|
|
@ -361,6 +394,13 @@ class ClientSession(
|
|||
types.ListToolsResult,
|
||||
)
|
||||
|
||||
# Cache tool output schemas for future validation
|
||||
# Note: don't clear the cache, as we may be using a cursor
|
||||
for tool in result.tools:
|
||||
self._tool_output_schemas[tool.name] = tool.outputSchema
|
||||
|
||||
return result
|
||||
|
||||
async def send_roots_list_changed(self) -> None:
|
||||
"""Send a roots/list_changed notification."""
|
||||
await self.send_notification(
|
||||
|
|
|
|||
|
|
@ -9,7 +9,6 @@ from contextlib import (
|
|||
AbstractAsyncContextManager,
|
||||
asynccontextmanager,
|
||||
)
|
||||
from itertools import chain
|
||||
from typing import Any, Generic, Literal
|
||||
|
||||
import anyio
|
||||
|
|
@ -38,7 +37,6 @@ from mcp.server.fastmcp.prompts import Prompt, PromptManager
|
|||
from mcp.server.fastmcp.resources import FunctionResource, Resource, ResourceManager
|
||||
from mcp.server.fastmcp.tools import Tool, ToolManager
|
||||
from mcp.server.fastmcp.utilities.logging import configure_logging, get_logger
|
||||
from mcp.server.fastmcp.utilities.types import Image
|
||||
from mcp.server.lowlevel.helper_types import ReadResourceContents
|
||||
from mcp.server.lowlevel.server import LifespanResultT
|
||||
from mcp.server.lowlevel.server import Server as MCPServer
|
||||
|
|
@ -54,7 +52,6 @@ from mcp.types import (
|
|||
AnyFunction,
|
||||
ContentBlock,
|
||||
GetPromptResult,
|
||||
TextContent,
|
||||
ToolAnnotations,
|
||||
)
|
||||
from mcp.types import Prompt as MCPPrompt
|
||||
|
|
@ -235,7 +232,10 @@ class FastMCP:
|
|||
def _setup_handlers(self) -> None:
|
||||
"""Set up core MCP protocol handlers."""
|
||||
self._mcp_server.list_tools()(self.list_tools)
|
||||
self._mcp_server.call_tool()(self.call_tool)
|
||||
# Note: we disable the lowlevel server's input validation.
|
||||
# FastMCP does ad hoc conversion of incoming data before validating -
|
||||
# for now we preserve this for backwards compatibility.
|
||||
self._mcp_server.call_tool(validate_input=False)(self.call_tool)
|
||||
self._mcp_server.list_resources()(self.list_resources)
|
||||
self._mcp_server.read_resource()(self.read_resource)
|
||||
self._mcp_server.list_prompts()(self.list_prompts)
|
||||
|
|
@ -251,6 +251,7 @@ class FastMCP:
|
|||
title=info.title,
|
||||
description=info.description,
|
||||
inputSchema=info.parameters,
|
||||
outputSchema=info.output_schema,
|
||||
annotations=info.annotations,
|
||||
)
|
||||
for info in tools
|
||||
|
|
@ -267,12 +268,10 @@ class FastMCP:
|
|||
request_context = None
|
||||
return Context(request_context=request_context, fastmcp=self)
|
||||
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[ContentBlock]:
|
||||
async def call_tool(self, name: str, arguments: dict[str, Any]) -> Sequence[ContentBlock] | dict[str, Any]:
|
||||
"""Call a tool by name with arguments."""
|
||||
context = self.get_context()
|
||||
result = await self._tool_manager.call_tool(name, arguments, context=context)
|
||||
converted_result = _convert_to_content(result)
|
||||
return converted_result
|
||||
return await self._tool_manager.call_tool(name, arguments, context=context, convert_result=True)
|
||||
|
||||
async def list_resources(self) -> list[MCPResource]:
|
||||
"""List all available resources."""
|
||||
|
|
@ -322,6 +321,7 @@ class FastMCP:
|
|||
title: str | None = None,
|
||||
description: str | None = None,
|
||||
annotations: ToolAnnotations | None = None,
|
||||
structured_output: bool | None = None,
|
||||
) -> None:
|
||||
"""Add a tool to the server.
|
||||
|
||||
|
|
@ -334,8 +334,19 @@ class FastMCP:
|
|||
title: Optional human-readable title for the tool
|
||||
description: Optional description of what the tool does
|
||||
annotations: Optional ToolAnnotations providing additional tool information
|
||||
structured_output: Controls whether the tool's output is structured or unstructured
|
||||
- If None, auto-detects based on the function's return type annotation
|
||||
- If True, unconditionally creates a structured tool (return type annotation permitting)
|
||||
- If False, unconditionally creates an unstructured tool
|
||||
"""
|
||||
self._tool_manager.add_tool(fn, name=name, title=title, description=description, annotations=annotations)
|
||||
self._tool_manager.add_tool(
|
||||
fn,
|
||||
name=name,
|
||||
title=title,
|
||||
description=description,
|
||||
annotations=annotations,
|
||||
structured_output=structured_output,
|
||||
)
|
||||
|
||||
def tool(
|
||||
self,
|
||||
|
|
@ -343,6 +354,7 @@ class FastMCP:
|
|||
title: str | None = None,
|
||||
description: str | None = None,
|
||||
annotations: ToolAnnotations | None = None,
|
||||
structured_output: bool | None = None,
|
||||
) -> Callable[[AnyFunction], AnyFunction]:
|
||||
"""Decorator to register a tool.
|
||||
|
||||
|
|
@ -355,6 +367,10 @@ class FastMCP:
|
|||
title: Optional human-readable title for the tool
|
||||
description: Optional description of what the tool does
|
||||
annotations: Optional ToolAnnotations providing additional tool information
|
||||
structured_output: Controls whether the tool's output is structured or unstructured
|
||||
- If None, auto-detects based on the function's return type annotation
|
||||
- If True, unconditionally creates a structured tool (return type annotation permitting)
|
||||
- If False, unconditionally creates an unstructured tool
|
||||
|
||||
Example:
|
||||
@server.tool()
|
||||
|
|
@ -378,7 +394,14 @@ class FastMCP:
|
|||
)
|
||||
|
||||
def decorator(fn: AnyFunction) -> AnyFunction:
|
||||
self.add_tool(fn, name=name, title=title, description=description, annotations=annotations)
|
||||
self.add_tool(
|
||||
fn,
|
||||
name=name,
|
||||
title=title,
|
||||
description=description,
|
||||
annotations=annotations,
|
||||
structured_output=structured_output,
|
||||
)
|
||||
return fn
|
||||
|
||||
return decorator
|
||||
|
|
@ -942,28 +965,6 @@ class FastMCP:
|
|||
raise ValueError(str(e))
|
||||
|
||||
|
||||
def _convert_to_content(
|
||||
result: Any,
|
||||
) -> Sequence[ContentBlock]:
|
||||
"""Convert a result to a sequence of content objects."""
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
if isinstance(result, ContentBlock):
|
||||
return [result]
|
||||
|
||||
if isinstance(result, Image):
|
||||
return [result.to_image_content()]
|
||||
|
||||
if isinstance(result, list | tuple):
|
||||
return list(chain.from_iterable(_convert_to_content(item) for item in result)) # type: ignore[reportUnknownVariableType]
|
||||
|
||||
if not isinstance(result, str):
|
||||
result = pydantic_core.to_json(result, fallback=str, indent=2).decode()
|
||||
|
||||
return [TextContent(type="text", text=result)]
|
||||
|
||||
|
||||
class Context(BaseModel, Generic[ServerSessionT, LifespanContextT, RequestT]):
|
||||
"""Context object providing access to MCP capabilities.
|
||||
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ from __future__ import annotations as _annotations
|
|||
import functools
|
||||
import inspect
|
||||
from collections.abc import Callable
|
||||
from functools import cached_property
|
||||
from typing import TYPE_CHECKING, Any, get_origin
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
|
@ -32,6 +33,10 @@ class Tool(BaseModel):
|
|||
context_kwarg: str | None = Field(None, description="Name of the kwarg that should receive context")
|
||||
annotations: ToolAnnotations | None = Field(None, description="Optional annotations for the tool")
|
||||
|
||||
@cached_property
|
||||
def output_schema(self) -> dict[str, Any] | None:
|
||||
return self.fn_metadata.output_schema
|
||||
|
||||
@classmethod
|
||||
def from_function(
|
||||
cls,
|
||||
|
|
@ -41,6 +46,7 @@ class Tool(BaseModel):
|
|||
description: str | None = None,
|
||||
context_kwarg: str | None = None,
|
||||
annotations: ToolAnnotations | None = None,
|
||||
structured_output: bool | None = None,
|
||||
) -> Tool:
|
||||
"""Create a Tool from a function."""
|
||||
from mcp.server.fastmcp.server import Context
|
||||
|
|
@ -65,6 +71,7 @@ class Tool(BaseModel):
|
|||
func_arg_metadata = func_metadata(
|
||||
fn,
|
||||
skip_names=[context_kwarg] if context_kwarg is not None else [],
|
||||
structured_output=structured_output,
|
||||
)
|
||||
parameters = func_arg_metadata.arg_model.model_json_schema()
|
||||
|
||||
|
|
@ -84,15 +91,21 @@ class Tool(BaseModel):
|
|||
self,
|
||||
arguments: dict[str, Any],
|
||||
context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
|
||||
convert_result: bool = False,
|
||||
) -> Any:
|
||||
"""Run the tool with arguments."""
|
||||
try:
|
||||
return await self.fn_metadata.call_fn_with_arg_validation(
|
||||
result = await self.fn_metadata.call_fn_with_arg_validation(
|
||||
self.fn,
|
||||
self.is_async,
|
||||
arguments,
|
||||
{self.context_kwarg: context} if self.context_kwarg is not None else None,
|
||||
)
|
||||
|
||||
if convert_result:
|
||||
result = self.fn_metadata.convert_result(result)
|
||||
|
||||
return result
|
||||
except Exception as e:
|
||||
raise ToolError(f"Error executing tool {self.name}: {e}") from e
|
||||
|
||||
|
|
|
|||
|
|
@ -49,9 +49,17 @@ class ToolManager:
|
|||
title: str | None = None,
|
||||
description: str | None = None,
|
||||
annotations: ToolAnnotations | None = None,
|
||||
structured_output: bool | None = None,
|
||||
) -> Tool:
|
||||
"""Add a tool to the server."""
|
||||
tool = Tool.from_function(fn, name=name, title=title, description=description, annotations=annotations)
|
||||
tool = Tool.from_function(
|
||||
fn,
|
||||
name=name,
|
||||
title=title,
|
||||
description=description,
|
||||
annotations=annotations,
|
||||
structured_output=structured_output,
|
||||
)
|
||||
existing = self._tools.get(tool.name)
|
||||
if existing:
|
||||
if self.warn_on_duplicate_tools:
|
||||
|
|
@ -65,10 +73,11 @@ class ToolManager:
|
|||
name: str,
|
||||
arguments: dict[str, Any],
|
||||
context: Context[ServerSessionT, LifespanContextT, RequestT] | None = None,
|
||||
convert_result: bool = False,
|
||||
) -> Any:
|
||||
"""Call a tool by name with arguments."""
|
||||
tool = self.get_tool(name)
|
||||
if not tool:
|
||||
raise ToolError(f"Unknown tool: {name}")
|
||||
|
||||
return await tool.run(arguments, context=context)
|
||||
return await tool.run(arguments, context=context, convert_result=convert_result)
|
||||
|
|
|
|||
|
|
@ -1,23 +1,43 @@
|
|||
import inspect
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable, Sequence
|
||||
from typing import (
|
||||
Annotated,
|
||||
Any,
|
||||
ForwardRef,
|
||||
)
|
||||
from itertools import chain
|
||||
from types import GenericAlias
|
||||
from typing import Annotated, Any, ForwardRef, cast, get_args, get_origin, get_type_hints
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, Field, WithJsonSchema, create_model
|
||||
import pydantic_core
|
||||
from pydantic import (
|
||||
BaseModel,
|
||||
ConfigDict,
|
||||
Field,
|
||||
RootModel,
|
||||
WithJsonSchema,
|
||||
create_model,
|
||||
)
|
||||
from pydantic._internal._typing_extra import eval_type_backport
|
||||
from pydantic.fields import FieldInfo
|
||||
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaWarningKind
|
||||
from pydantic_core import PydanticUndefined
|
||||
|
||||
from mcp.server.fastmcp.exceptions import InvalidSignature
|
||||
from mcp.server.fastmcp.utilities.logging import get_logger
|
||||
from mcp.server.fastmcp.utilities.types import Image
|
||||
from mcp.types import ContentBlock, TextContent
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class StrictJsonSchema(GenerateJsonSchema):
|
||||
"""A JSON schema generator that raises exceptions instead of emitting warnings.
|
||||
|
||||
This is used to detect non-serializable types during schema generation.
|
||||
"""
|
||||
|
||||
def emit_warning(self, kind: JsonSchemaWarningKind, detail: str) -> None:
|
||||
# Raise an exception instead of emitting a warning
|
||||
raise ValueError(f"JSON schema warning: {kind} - {detail}")
|
||||
|
||||
|
||||
class ArgModelBase(BaseModel):
|
||||
"""A model representing the arguments to a function."""
|
||||
|
||||
|
|
@ -38,13 +58,13 @@ class ArgModelBase(BaseModel):
|
|||
|
||||
class FuncMetadata(BaseModel):
|
||||
arg_model: Annotated[type[ArgModelBase], WithJsonSchema(None)]
|
||||
# We can add things in the future like
|
||||
# - Maybe some args are excluded from attempting to parse from JSON
|
||||
# - Maybe some args are special (like context) for dependency injection
|
||||
output_schema: dict[str, Any] | None = None
|
||||
output_model: Annotated[type[BaseModel], WithJsonSchema(None)] | None = None
|
||||
wrap_output: bool = False
|
||||
|
||||
async def call_fn_with_arg_validation(
|
||||
self,
|
||||
fn: Callable[..., Any] | Awaitable[Any],
|
||||
fn: Callable[..., Any | Awaitable[Any]],
|
||||
fn_is_async: bool,
|
||||
arguments_to_validate: dict[str, Any],
|
||||
arguments_to_pass_directly: dict[str, Any] | None,
|
||||
|
|
@ -61,12 +81,39 @@ class FuncMetadata(BaseModel):
|
|||
arguments_parsed_dict |= arguments_to_pass_directly or {}
|
||||
|
||||
if fn_is_async:
|
||||
if isinstance(fn, Awaitable):
|
||||
return await fn
|
||||
return await fn(**arguments_parsed_dict)
|
||||
if isinstance(fn, Callable):
|
||||
else:
|
||||
return fn(**arguments_parsed_dict)
|
||||
raise TypeError("fn must be either Callable or Awaitable")
|
||||
|
||||
def convert_result(self, result: Any) -> Any:
|
||||
"""
|
||||
Convert the result of a function call to the appropriate format for
|
||||
the lowlevel server tool call handler:
|
||||
|
||||
- If output_model is None, return the unstructured content directly.
|
||||
- If output_model is not None, convert the result to structured output format
|
||||
(dict[str, Any]) and return both unstructured and structured content.
|
||||
|
||||
Note: we return unstructured content here **even though the lowlevel server
|
||||
tool call handler provides generic backwards compatibility serialization of
|
||||
structured content**. This is for FastMCP backwards compatibility: we need to
|
||||
retain FastMCP's ad hoc conversion logic for constructing unstructured output
|
||||
from function return values, whereas the lowlevel server simply serializes
|
||||
the structured output.
|
||||
"""
|
||||
unstructured_content = _convert_to_content(result)
|
||||
|
||||
if self.output_schema is None:
|
||||
return unstructured_content
|
||||
else:
|
||||
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")
|
||||
|
||||
return (unstructured_content, structured_content)
|
||||
|
||||
def pre_parse_json(self, data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Pre-parse data from JSON.
|
||||
|
|
@ -102,13 +149,17 @@ class FuncMetadata(BaseModel):
|
|||
)
|
||||
|
||||
|
||||
def func_metadata(func: Callable[..., Any], skip_names: Sequence[str] = ()) -> FuncMetadata:
|
||||
def func_metadata(
|
||||
func: Callable[..., Any],
|
||||
skip_names: Sequence[str] = (),
|
||||
structured_output: bool | None = None,
|
||||
) -> FuncMetadata:
|
||||
"""Given a function, return metadata including a pydantic model representing its
|
||||
signature.
|
||||
|
||||
The use case for this is
|
||||
```
|
||||
meta = func_to_pyd(func)
|
||||
meta = func_metadata(func)
|
||||
validated_args = meta.arg_model.model_validate(some_raw_data_dict)
|
||||
return func(**validated_args.model_dump_one_level())
|
||||
```
|
||||
|
|
@ -120,8 +171,25 @@ def func_metadata(func: Callable[..., Any], skip_names: Sequence[str] = ()) -> F
|
|||
func: The function to convert to a pydantic model
|
||||
skip_names: A list of parameter names to skip. These will not be included in
|
||||
the model.
|
||||
structured_output: Controls whether the tool's output is structured or unstructured
|
||||
- If None, auto-detects based on the function's return type annotation
|
||||
- If True, unconditionally creates a structured tool (return type annotation permitting)
|
||||
- If False, unconditionally creates an unstructured tool
|
||||
|
||||
If structured, creates a Pydantic model for the function's result based on its annotation.
|
||||
Supports various return types:
|
||||
- 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
|
||||
- Dataclasses and other annotated classes - converted to Pydantic models
|
||||
- Generic types (list, dict, Union, etc.) - wrapped in a model with a 'result' field
|
||||
|
||||
Returns:
|
||||
A pydantic model representing the function's signature.
|
||||
A FuncMetadata object containing:
|
||||
- arg_model: A pydantic model representing the function's arguments
|
||||
- output_model: A pydantic model for the return type if output is structured
|
||||
- output_conversion: Records how function output should be converted before returning.
|
||||
"""
|
||||
sig = _get_typed_signature(func)
|
||||
params = sig.parameters
|
||||
|
|
@ -162,8 +230,197 @@ def func_metadata(func: Callable[..., Any], skip_names: Sequence[str] = ()) -> F
|
|||
**dynamic_pydantic_model_params,
|
||||
__base__=ArgModelBase,
|
||||
)
|
||||
resp = FuncMetadata(arg_model=arguments_model)
|
||||
return resp
|
||||
|
||||
if structured_output is False:
|
||||
return FuncMetadata(arg_model=arguments_model)
|
||||
|
||||
# set up structured output support based on return type annotation
|
||||
|
||||
if sig.return_annotation is inspect.Parameter.empty and structured_output is True:
|
||||
raise InvalidSignature(f"Function {func.__name__}: return annotation required for structured output")
|
||||
|
||||
output_info = FieldInfo.from_annotation(_get_typed_annotation(sig.return_annotation, globalns))
|
||||
annotation = output_info.annotation
|
||||
|
||||
output_model, output_schema, wrap_output = _try_create_model_and_schema(annotation, func.__name__, output_info)
|
||||
|
||||
if output_model is None and structured_output is True:
|
||||
# Model creation failed or produced warnings - no structured output
|
||||
raise InvalidSignature(
|
||||
f"Function {func.__name__}: return type {annotation} is not serializable for structured output"
|
||||
)
|
||||
|
||||
return FuncMetadata(
|
||||
arg_model=arguments_model,
|
||||
output_schema=output_schema,
|
||||
output_model=output_model,
|
||||
wrap_output=wrap_output,
|
||||
)
|
||||
|
||||
|
||||
def _try_create_model_and_schema(
|
||||
annotation: Any, func_name: str, field_info: FieldInfo
|
||||
) -> tuple[type[BaseModel] | None, dict[str, Any] | None, bool]:
|
||||
"""Try to create a model and schema for the given annotation without warnings.
|
||||
|
||||
Returns:
|
||||
tuple of (model or None, schema or None, wrap_output)
|
||||
Model and schema are None if warnings occur or creation fails.
|
||||
wrap_output is True if the result needs to be wrapped in {"result": ...}
|
||||
"""
|
||||
model = None
|
||||
wrap_output = False
|
||||
|
||||
# First handle special case: None
|
||||
if annotation is None:
|
||||
model = _create_wrapped_model(func_name, annotation, field_info)
|
||||
wrap_output = True
|
||||
|
||||
# Handle GenericAlias types (list[str], dict[str, int], Union[str, int], etc.)
|
||||
elif isinstance(annotation, GenericAlias):
|
||||
origin = get_origin(annotation)
|
||||
|
||||
# Special case: dict with string keys can use RootModel
|
||||
if origin is dict:
|
||||
args = get_args(annotation)
|
||||
if len(args) == 2 and args[0] is str:
|
||||
model = _create_dict_model(func_name, annotation)
|
||||
else:
|
||||
# dict with non-str keys needs wrapping
|
||||
model = _create_wrapped_model(func_name, annotation, field_info)
|
||||
wrap_output = True
|
||||
else:
|
||||
# All other generic types need wrapping (list, tuple, Union, Optional, etc.)
|
||||
model = _create_wrapped_model(func_name, annotation, field_info)
|
||||
wrap_output = True
|
||||
|
||||
# Handle regular type objects
|
||||
elif isinstance(annotation, type):
|
||||
type_annotation: type[Any] = cast(type[Any], annotation)
|
||||
|
||||
# Case 1: BaseModel subclasses (can be used directly)
|
||||
if issubclass(annotation, BaseModel):
|
||||
model = annotation
|
||||
|
||||
# Case 2: TypedDict (special dict subclass with __annotations__)
|
||||
elif hasattr(type_annotation, "__annotations__") and issubclass(annotation, dict):
|
||||
model = _create_model_from_typeddict(type_annotation)
|
||||
|
||||
# Case 3: Primitive types that need wrapping
|
||||
elif annotation in (str, int, float, bool, bytes, type(None)):
|
||||
model = _create_wrapped_model(func_name, annotation, field_info)
|
||||
wrap_output = True
|
||||
|
||||
# Case 4: Other class types (dataclasses, regular classes with annotations)
|
||||
else:
|
||||
type_hints = get_type_hints(type_annotation)
|
||||
if type_hints:
|
||||
# Classes with type hints can be converted to Pydantic models
|
||||
model = _create_model_from_class(type_annotation)
|
||||
# Classes without type hints are not serializable - model remains None
|
||||
|
||||
# Handle any other types not covered above
|
||||
else:
|
||||
# This includes typing constructs that aren't GenericAlias in Python 3.10
|
||||
# (e.g., Union, Optional in some Python versions)
|
||||
model = _create_wrapped_model(func_name, annotation, field_info)
|
||||
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 (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
|
||||
# TypeError: When Pydantic can't handle the type
|
||||
# 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 {annotation} in {func_name}: {type(e).__name__}: {e}")
|
||||
return None, None, False
|
||||
|
||||
return model, schema, wrap_output
|
||||
|
||||
return None, None, False
|
||||
|
||||
|
||||
def _create_model_from_class(cls: type[Any]) -> type[BaseModel]:
|
||||
"""Create a Pydantic model from an ordinary class.
|
||||
|
||||
The created model will:
|
||||
- Have the same name as the class
|
||||
- Have fields with the same names and types as the class's fields
|
||||
- Include all fields whose type does not include None in the set of required fields
|
||||
|
||||
Precondition: cls must have type hints (i.e., get_type_hints(cls) is non-empty)
|
||||
"""
|
||||
type_hints = get_type_hints(cls)
|
||||
|
||||
model_fields: dict[str, Any] = {}
|
||||
for field_name, field_type in type_hints.items():
|
||||
if field_name.startswith("_"):
|
||||
continue
|
||||
|
||||
default = getattr(cls, field_name, PydanticUndefined)
|
||||
field_info = FieldInfo.from_annotated_attribute(field_type, default)
|
||||
model_fields[field_name] = (field_info.annotation, field_info)
|
||||
|
||||
# Create a base class with the config
|
||||
class BaseWithConfig(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True)
|
||||
|
||||
return create_model(cls.__name__, **model_fields, __base__=BaseWithConfig)
|
||||
|
||||
|
||||
def _create_model_from_typeddict(td_type: type[Any]) -> type[BaseModel]:
|
||||
"""Create a Pydantic model from a TypedDict.
|
||||
|
||||
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():
|
||||
field_info = FieldInfo.from_annotation(field_type)
|
||||
|
||||
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
|
||||
field_info.default = None
|
||||
|
||||
model_fields[field_name] = (field_info.annotation, field_info)
|
||||
|
||||
return create_model(td_type.__name__, **model_fields, __base__=BaseModel)
|
||||
|
||||
|
||||
def _create_wrapped_model(func_name: str, annotation: Any, field_info: FieldInfo) -> type[BaseModel]:
|
||||
"""Create a model that wraps a type in a 'result' field.
|
||||
|
||||
This is used for primitive types, generic types like list/dict, etc.
|
||||
"""
|
||||
model_name = f"{func_name}Output"
|
||||
|
||||
# Pydantic needs type(None) instead of None for the type annotation
|
||||
if annotation is None:
|
||||
annotation = type(None)
|
||||
|
||||
return create_model(model_name, result=(annotation, field_info), __base__=BaseModel)
|
||||
|
||||
|
||||
def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:
|
||||
"""Create a RootModel for dict[str, T] types."""
|
||||
|
||||
class DictModel(RootModel[dict_annotation]):
|
||||
pass
|
||||
|
||||
# Give it a meaningful name
|
||||
DictModel.__name__ = f"{func_name}DictOutput"
|
||||
DictModel.__qualname__ = f"{func_name}DictOutput"
|
||||
|
||||
return DictModel
|
||||
|
||||
|
||||
def _get_typed_annotation(annotation: Any, globalns: dict[str, Any]) -> Any:
|
||||
|
|
@ -198,5 +455,40 @@ def _get_typed_signature(call: Callable[..., Any]) -> inspect.Signature:
|
|||
)
|
||||
for param in signature.parameters.values()
|
||||
]
|
||||
typed_signature = inspect.Signature(typed_params)
|
||||
typed_return = _get_typed_annotation(signature.return_annotation, globalns)
|
||||
typed_signature = inspect.Signature(typed_params, return_annotation=typed_return)
|
||||
return typed_signature
|
||||
|
||||
|
||||
def _convert_to_content(
|
||||
result: Any,
|
||||
) -> Sequence[ContentBlock]:
|
||||
"""
|
||||
Convert a result to a sequence of content objects.
|
||||
|
||||
Note: This conversion logic comes from previous versions of FastMCP and is being
|
||||
retained for purposes of backwards compatibility. It produces different unstructured
|
||||
output than the lowlevel server tool call handler, which just serializes structured
|
||||
content verbatim.
|
||||
"""
|
||||
if result is None:
|
||||
return []
|
||||
|
||||
if isinstance(result, ContentBlock):
|
||||
return [result]
|
||||
|
||||
if isinstance(result, Image):
|
||||
return [result.to_image_content()]
|
||||
|
||||
if isinstance(result, list | tuple):
|
||||
return list(
|
||||
chain.from_iterable(
|
||||
_convert_to_content(item)
|
||||
for item in result # type: ignore
|
||||
)
|
||||
)
|
||||
|
||||
if not isinstance(result, str):
|
||||
result = pydantic_core.to_json(result, fallback=str, indent=2).decode()
|
||||
|
||||
return [TextContent(type="text", text=result)]
|
||||
|
|
|
|||
Loading…
Reference in a new issue