Replace RootModel wrappers with type aliases and TypeAdapter validation (#3470)
This commit is contained in:
Marcelo Trylesinski
GitHub
parent
fd66270926
commit
9972c21aa4
6 changed files with 872 additions and 1127 deletions
|
|
@ -9,7 +9,7 @@ from __future__ import annotations
|
|||
from typing import Annotated, Any, Literal
|
||||
|
||||
from mcp_types._wire_base import WireModel
|
||||
from pydantic import ConfigDict, Field, RootModel
|
||||
from pydantic import ConfigDict, Field
|
||||
|
||||
|
||||
class BaseMetadata(WireModel):
|
||||
|
|
@ -285,11 +285,10 @@ class CompleteResult(WireModel):
|
|||
completion: Completion
|
||||
|
||||
|
||||
class Cursor(RootModel[str]):
|
||||
root: str
|
||||
"""
|
||||
An opaque token used to represent a cursor for pagination.
|
||||
"""
|
||||
Cursor = str
|
||||
"""
|
||||
An opaque token used to represent a cursor for pagination.
|
||||
"""
|
||||
|
||||
|
||||
class RequestedSchema(WireModel):
|
||||
|
|
@ -556,27 +555,13 @@ class LegacyTitledEnumSchema(WireModel):
|
|||
type: Literal["string"]
|
||||
|
||||
|
||||
class LoggingLevel(
|
||||
RootModel[
|
||||
Literal[
|
||||
"alert",
|
||||
"critical",
|
||||
"debug",
|
||||
"emergency",
|
||||
"error",
|
||||
"info",
|
||||
"notice",
|
||||
"warning",
|
||||
]
|
||||
]
|
||||
):
|
||||
root: Literal["alert", "critical", "debug", "emergency", "error", "info", "notice", "warning"]
|
||||
"""
|
||||
The severity of a log message.
|
||||
LoggingLevel = Literal["alert", "critical", "debug", "emergency", "error", "info", "notice", "warning"]
|
||||
"""
|
||||
The severity of a log message.
|
||||
|
||||
These map to syslog message severities, as specified in RFC-5424:
|
||||
https://datatracker.ietf.org/doc/html/rfc5424#section-6.2.1
|
||||
"""
|
||||
These map to syslog message severities, as specified in RFC-5424:
|
||||
https://datatracker.ietf.org/doc/html/rfc5424#section-6.2.1
|
||||
"""
|
||||
|
||||
|
||||
class LoggingMessageNotificationParams(WireModel):
|
||||
|
|
@ -723,11 +708,10 @@ class PaginatedResult(WireModel):
|
|||
"""
|
||||
|
||||
|
||||
class ProgressToken(RootModel[str | int]):
|
||||
root: str | int
|
||||
"""
|
||||
A progress token, used to associate progress notifications with the original request.
|
||||
"""
|
||||
ProgressToken = str | int
|
||||
"""
|
||||
A progress token, used to associate progress notifications with the original request.
|
||||
"""
|
||||
|
||||
|
||||
class PromptArgument(WireModel):
|
||||
|
|
@ -853,11 +837,10 @@ class Request(WireModel):
|
|||
params: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class RequestId(RootModel[str | int]):
|
||||
root: str | int
|
||||
"""
|
||||
A uniquely identifying ID for a request in JSON-RPC.
|
||||
"""
|
||||
RequestId = str | int
|
||||
"""
|
||||
A uniquely identifying ID for a request in JSON-RPC.
|
||||
"""
|
||||
|
||||
|
||||
class RequestParams(WireModel):
|
||||
|
|
@ -970,11 +953,10 @@ class Result(WireModel):
|
|||
"""
|
||||
|
||||
|
||||
class Role(RootModel[Literal["assistant", "user"]]):
|
||||
root: Literal["assistant", "user"]
|
||||
"""
|
||||
The sender or recipient of messages and data in a conversation.
|
||||
"""
|
||||
Role = Literal["assistant", "user"]
|
||||
"""
|
||||
The sender or recipient of messages and data in a conversation.
|
||||
"""
|
||||
|
||||
|
||||
class Root(WireModel):
|
||||
|
|
@ -1216,11 +1198,10 @@ class TaskMetadata(WireModel):
|
|||
"""
|
||||
|
||||
|
||||
class TaskStatus(RootModel[Literal["cancelled", "completed", "failed", "input_required", "working"]]):
|
||||
root: Literal["cancelled", "completed", "failed", "input_required", "working"]
|
||||
"""
|
||||
The status of a task.
|
||||
"""
|
||||
TaskStatus = Literal["cancelled", "completed", "failed", "input_required", "working"]
|
||||
"""
|
||||
The status of a task.
|
||||
"""
|
||||
|
||||
|
||||
class TextResourceContents(WireModel):
|
||||
|
|
@ -1867,26 +1848,16 @@ class EmbeddedResource(WireModel):
|
|||
type: Literal["resource"]
|
||||
|
||||
|
||||
class EmptyResult(RootModel[Result]):
|
||||
root: Result
|
||||
EmptyResult = Result
|
||||
|
||||
|
||||
class EnumSchema(
|
||||
RootModel[
|
||||
UntitledSingleSelectEnumSchema
|
||||
| TitledSingleSelectEnumSchema
|
||||
| UntitledMultiSelectEnumSchema
|
||||
| TitledMultiSelectEnumSchema
|
||||
| LegacyTitledEnumSchema
|
||||
]
|
||||
):
|
||||
root: (
|
||||
UntitledSingleSelectEnumSchema
|
||||
| TitledSingleSelectEnumSchema
|
||||
| UntitledMultiSelectEnumSchema
|
||||
| TitledMultiSelectEnumSchema
|
||||
| LegacyTitledEnumSchema
|
||||
)
|
||||
EnumSchema = (
|
||||
UntitledSingleSelectEnumSchema
|
||||
| TitledSingleSelectEnumSchema
|
||||
| UntitledMultiSelectEnumSchema
|
||||
| TitledMultiSelectEnumSchema
|
||||
| LegacyTitledEnumSchema
|
||||
)
|
||||
|
||||
|
||||
class GetPromptRequestParams(WireModel):
|
||||
|
|
@ -2115,8 +2086,7 @@ class LoggingMessageNotification(WireModel):
|
|||
params: LoggingMessageNotificationParams
|
||||
|
||||
|
||||
class MultiSelectEnumSchema(RootModel[UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema]):
|
||||
root: UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
|
||||
MultiSelectEnumSchema = UntitledMultiSelectEnumSchema | TitledMultiSelectEnumSchema
|
||||
|
||||
|
||||
class PaginatedRequestParams(WireModel):
|
||||
|
|
@ -2152,32 +2122,20 @@ class PingRequest(WireModel):
|
|||
params: RequestParams | None = None
|
||||
|
||||
|
||||
class PrimitiveSchemaDefinition(
|
||||
RootModel[
|
||||
StringSchema
|
||||
| NumberSchema
|
||||
| BooleanSchema
|
||||
| UntitledSingleSelectEnumSchema
|
||||
| TitledSingleSelectEnumSchema
|
||||
| UntitledMultiSelectEnumSchema
|
||||
| TitledMultiSelectEnumSchema
|
||||
| LegacyTitledEnumSchema
|
||||
]
|
||||
):
|
||||
root: (
|
||||
StringSchema
|
||||
| NumberSchema
|
||||
| BooleanSchema
|
||||
| UntitledSingleSelectEnumSchema
|
||||
| TitledSingleSelectEnumSchema
|
||||
| UntitledMultiSelectEnumSchema
|
||||
| TitledMultiSelectEnumSchema
|
||||
| LegacyTitledEnumSchema
|
||||
)
|
||||
"""
|
||||
Restricted schema definitions that only allow primitive types
|
||||
without nested objects or arrays.
|
||||
"""
|
||||
PrimitiveSchemaDefinition = (
|
||||
StringSchema
|
||||
| NumberSchema
|
||||
| BooleanSchema
|
||||
| UntitledSingleSelectEnumSchema
|
||||
| TitledSingleSelectEnumSchema
|
||||
| UntitledMultiSelectEnumSchema
|
||||
| TitledMultiSelectEnumSchema
|
||||
| LegacyTitledEnumSchema
|
||||
)
|
||||
"""
|
||||
Restricted schema definitions that only allow primitive types
|
||||
without nested objects or arrays.
|
||||
"""
|
||||
|
||||
|
||||
class ProgressNotificationParams(WireModel):
|
||||
|
|
@ -2499,8 +2457,7 @@ class SetLevelRequest(WireModel):
|
|||
params: SetLevelRequestParams
|
||||
|
||||
|
||||
class SingleSelectEnumSchema(RootModel[UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema]):
|
||||
root: UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
|
||||
SingleSelectEnumSchema = UntitledSingleSelectEnumSchema | TitledSingleSelectEnumSchema
|
||||
|
||||
|
||||
class SubscribeRequest(WireModel):
|
||||
|
|
@ -2793,8 +2750,7 @@ class CompleteRequest(WireModel):
|
|||
params: CompleteRequestParams
|
||||
|
||||
|
||||
class ContentBlock(RootModel[TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource]):
|
||||
root: TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
|
||||
ContentBlock = TextContent | ImageContent | AudioContent | ResourceLink | EmbeddedResource
|
||||
|
||||
|
||||
class CreateTaskResult(WireModel):
|
||||
|
|
@ -2812,11 +2768,10 @@ class CreateTaskResult(WireModel):
|
|||
task: Task
|
||||
|
||||
|
||||
class ElicitRequestParams(RootModel[ElicitRequestURLParams | ElicitRequestFormParams]):
|
||||
root: ElicitRequestURLParams | ElicitRequestFormParams
|
||||
"""
|
||||
The parameters for a request to elicit additional information from the user via the client.
|
||||
"""
|
||||
ElicitRequestParams = ElicitRequestURLParams | ElicitRequestFormParams
|
||||
"""
|
||||
The parameters for a request to elicit additional information from the user via the client.
|
||||
"""
|
||||
|
||||
|
||||
class GetPromptRequest(WireModel):
|
||||
|
|
@ -2857,18 +2812,16 @@ class InitializeRequest(WireModel):
|
|||
params: InitializeRequestParams
|
||||
|
||||
|
||||
class JSONRPCMessage(RootModel[JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse]):
|
||||
root: JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse
|
||||
"""
|
||||
Refers to any valid JSON-RPC object that can be decoded off the wire, or encoded to be sent.
|
||||
"""
|
||||
JSONRPCMessage = JSONRPCRequest | JSONRPCNotification | JSONRPCResultResponse | JSONRPCErrorResponse
|
||||
"""
|
||||
Refers to any valid JSON-RPC object that can be decoded off the wire, or encoded to be sent.
|
||||
"""
|
||||
|
||||
|
||||
class JSONRPCResponse(RootModel[JSONRPCResultResponse | JSONRPCErrorResponse]):
|
||||
root: JSONRPCResultResponse | JSONRPCErrorResponse
|
||||
"""
|
||||
A response to a request, containing either the result or error.
|
||||
"""
|
||||
JSONRPCResponse = JSONRPCResultResponse | JSONRPCErrorResponse
|
||||
"""
|
||||
A response to a request, containing either the result or error.
|
||||
"""
|
||||
|
||||
|
||||
class ListPromptsRequest(WireModel):
|
||||
|
|
@ -3173,64 +3126,34 @@ class CallToolResult(WireModel):
|
|||
"""
|
||||
|
||||
|
||||
class ClientNotification(
|
||||
RootModel[
|
||||
CancelledNotification
|
||||
| InitializedNotification
|
||||
| ProgressNotification
|
||||
| TaskStatusNotification
|
||||
| RootsListChangedNotification
|
||||
]
|
||||
):
|
||||
root: (
|
||||
CancelledNotification
|
||||
| InitializedNotification
|
||||
| ProgressNotification
|
||||
| TaskStatusNotification
|
||||
| RootsListChangedNotification
|
||||
)
|
||||
ClientNotification = (
|
||||
CancelledNotification
|
||||
| InitializedNotification
|
||||
| ProgressNotification
|
||||
| TaskStatusNotification
|
||||
| RootsListChangedNotification
|
||||
)
|
||||
|
||||
|
||||
class ClientRequest(
|
||||
RootModel[
|
||||
InitializeRequest
|
||||
| PingRequest
|
||||
| ListResourcesRequest
|
||||
| ListResourceTemplatesRequest
|
||||
| ReadResourceRequest
|
||||
| SubscribeRequest
|
||||
| UnsubscribeRequest
|
||||
| ListPromptsRequest
|
||||
| GetPromptRequest
|
||||
| ListToolsRequest
|
||||
| CallToolRequest
|
||||
| GetTaskRequest
|
||||
| GetTaskPayloadRequest
|
||||
| CancelTaskRequest
|
||||
| ListTasksRequest
|
||||
| SetLevelRequest
|
||||
| CompleteRequest
|
||||
]
|
||||
):
|
||||
root: (
|
||||
InitializeRequest
|
||||
| PingRequest
|
||||
| ListResourcesRequest
|
||||
| ListResourceTemplatesRequest
|
||||
| ReadResourceRequest
|
||||
| SubscribeRequest
|
||||
| UnsubscribeRequest
|
||||
| ListPromptsRequest
|
||||
| GetPromptRequest
|
||||
| ListToolsRequest
|
||||
| CallToolRequest
|
||||
| GetTaskRequest
|
||||
| GetTaskPayloadRequest
|
||||
| CancelTaskRequest
|
||||
| ListTasksRequest
|
||||
| SetLevelRequest
|
||||
| CompleteRequest
|
||||
)
|
||||
ClientRequest = (
|
||||
InitializeRequest
|
||||
| PingRequest
|
||||
| ListResourcesRequest
|
||||
| ListResourceTemplatesRequest
|
||||
| ReadResourceRequest
|
||||
| SubscribeRequest
|
||||
| UnsubscribeRequest
|
||||
| ListPromptsRequest
|
||||
| GetPromptRequest
|
||||
| ListToolsRequest
|
||||
| CallToolRequest
|
||||
| GetTaskRequest
|
||||
| GetTaskPayloadRequest
|
||||
| CancelTaskRequest
|
||||
| ListTasksRequest
|
||||
| SetLevelRequest
|
||||
| CompleteRequest
|
||||
)
|
||||
|
||||
|
||||
class ElicitRequest(WireModel):
|
||||
|
|
@ -3266,72 +3189,38 @@ class GetPromptResult(WireModel):
|
|||
messages: list[PromptMessage]
|
||||
|
||||
|
||||
class SamplingMessageContentBlock(
|
||||
RootModel[TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent]
|
||||
):
|
||||
root: TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
|
||||
SamplingMessageContentBlock = TextContent | ImageContent | AudioContent | ToolUseContent | ToolResultContent
|
||||
|
||||
|
||||
class ServerNotification(
|
||||
RootModel[
|
||||
CancelledNotification
|
||||
| ProgressNotification
|
||||
| ResourceListChangedNotification
|
||||
| ResourceUpdatedNotification
|
||||
| PromptListChangedNotification
|
||||
| ToolListChangedNotification
|
||||
| TaskStatusNotification
|
||||
| LoggingMessageNotification
|
||||
| ElicitationCompleteNotification
|
||||
]
|
||||
):
|
||||
root: (
|
||||
CancelledNotification
|
||||
| ProgressNotification
|
||||
| ResourceListChangedNotification
|
||||
| ResourceUpdatedNotification
|
||||
| PromptListChangedNotification
|
||||
| ToolListChangedNotification
|
||||
| TaskStatusNotification
|
||||
| LoggingMessageNotification
|
||||
| ElicitationCompleteNotification
|
||||
)
|
||||
ServerNotification = (
|
||||
CancelledNotification
|
||||
| ProgressNotification
|
||||
| ResourceListChangedNotification
|
||||
| ResourceUpdatedNotification
|
||||
| PromptListChangedNotification
|
||||
| ToolListChangedNotification
|
||||
| TaskStatusNotification
|
||||
| LoggingMessageNotification
|
||||
| ElicitationCompleteNotification
|
||||
)
|
||||
|
||||
|
||||
class ServerResult(
|
||||
RootModel[
|
||||
Result
|
||||
| InitializeResult
|
||||
| ListResourcesResult
|
||||
| ListResourceTemplatesResult
|
||||
| ReadResourceResult
|
||||
| ListPromptsResult
|
||||
| GetPromptResult
|
||||
| ListToolsResult
|
||||
| CallToolResult
|
||||
| GetTaskResult
|
||||
| GetTaskPayloadResult
|
||||
| CancelTaskResult
|
||||
| ListTasksResult
|
||||
| CompleteResult
|
||||
]
|
||||
):
|
||||
root: (
|
||||
Result
|
||||
| InitializeResult
|
||||
| ListResourcesResult
|
||||
| ListResourceTemplatesResult
|
||||
| ReadResourceResult
|
||||
| ListPromptsResult
|
||||
| GetPromptResult
|
||||
| ListToolsResult
|
||||
| CallToolResult
|
||||
| GetTaskResult
|
||||
| GetTaskPayloadResult
|
||||
| CancelTaskResult
|
||||
| ListTasksResult
|
||||
| CompleteResult
|
||||
)
|
||||
ServerResult = (
|
||||
Result
|
||||
| InitializeResult
|
||||
| ListResourcesResult
|
||||
| ListResourceTemplatesResult
|
||||
| ReadResourceResult
|
||||
| ListPromptsResult
|
||||
| GetPromptResult
|
||||
| ListToolsResult
|
||||
| CallToolResult
|
||||
| GetTaskResult
|
||||
| GetTaskPayloadResult
|
||||
| CancelTaskResult
|
||||
| ListTasksResult
|
||||
| CompleteResult
|
||||
)
|
||||
|
||||
|
||||
class CreateMessageResult(WireModel):
|
||||
|
|
@ -3398,28 +3287,16 @@ class SamplingMessage(WireModel):
|
|||
role: Role
|
||||
|
||||
|
||||
class ClientResult(
|
||||
RootModel[
|
||||
Result
|
||||
| GetTaskResult
|
||||
| GetTaskPayloadResult
|
||||
| CancelTaskResult
|
||||
| ListTasksResult
|
||||
| CreateMessageResult
|
||||
| ListRootsResult
|
||||
| ElicitResult
|
||||
]
|
||||
):
|
||||
root: (
|
||||
Result
|
||||
| GetTaskResult
|
||||
| GetTaskPayloadResult
|
||||
| CancelTaskResult
|
||||
| ListTasksResult
|
||||
| CreateMessageResult
|
||||
| ListRootsResult
|
||||
| ElicitResult
|
||||
)
|
||||
ClientResult = (
|
||||
Result
|
||||
| GetTaskResult
|
||||
| GetTaskPayloadResult
|
||||
| CancelTaskResult
|
||||
| ListTasksResult
|
||||
| CreateMessageResult
|
||||
| ListRootsResult
|
||||
| ElicitResult
|
||||
)
|
||||
|
||||
|
||||
class CreateMessageRequestParams(WireModel):
|
||||
|
|
@ -3502,25 +3379,13 @@ class CreateMessageRequest(WireModel):
|
|||
params: CreateMessageRequestParams
|
||||
|
||||
|
||||
class ServerRequest(
|
||||
RootModel[
|
||||
PingRequest
|
||||
| GetTaskRequest
|
||||
| GetTaskPayloadRequest
|
||||
| CancelTaskRequest
|
||||
| ListTasksRequest
|
||||
| CreateMessageRequest
|
||||
| ListRootsRequest
|
||||
| ElicitRequest
|
||||
]
|
||||
):
|
||||
root: (
|
||||
PingRequest
|
||||
| GetTaskRequest
|
||||
| GetTaskPayloadRequest
|
||||
| CancelTaskRequest
|
||||
| ListTasksRequest
|
||||
| CreateMessageRequest
|
||||
| ListRootsRequest
|
||||
| ElicitRequest
|
||||
)
|
||||
ServerRequest = (
|
||||
PingRequest
|
||||
| GetTaskRequest
|
||||
| GetTaskPayloadRequest
|
||||
| CancelTaskRequest
|
||||
| ListTasksRequest
|
||||
| CreateMessageRequest
|
||||
| ListRootsRequest
|
||||
| ElicitRequest
|
||||
)
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -80,19 +80,12 @@ def _clamp_inbound_ttl(raw: dict[str, Any]) -> None:
|
|||
|
||||
@cache
|
||||
def _wire_fields(target: type[BaseModel] | UnionType) -> frozenset[str]:
|
||||
"""Top-level wire keys `target` declares (its members', for a union).
|
||||
|
||||
A `RootModel` row (e.g. an empty result carried as `RootModel[Result]`)
|
||||
reports its wrapped type's keys, not the pydantic-internal `root`.
|
||||
"""
|
||||
"""Top-level wire keys `target` declares (its members', for a union)."""
|
||||
members: tuple[Any, ...] = get_args(target) if isinstance(target, UnionType) else (target,)
|
||||
models = [m for m in members if isinstance(m, type) and issubclass(m, BaseModel)]
|
||||
fields: set[str] = set()
|
||||
for model in models:
|
||||
if getattr(model, "__pydantic_root_model__", False): # a RootModel wrapper row
|
||||
fields |= _wire_fields(model.model_fields["root"].annotation)
|
||||
else:
|
||||
fields.update(field.alias or name for name, field in model.model_fields.items())
|
||||
fields.update(field.alias or name for name, field in model.model_fields.items())
|
||||
return frozenset(fields)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2,8 +2,7 @@ import logging
|
|||
from dataclasses import dataclass
|
||||
from typing import Any, Literal
|
||||
|
||||
# TODO(Marcelo): We should drop the `RootModel`.
|
||||
from pydantic import AnyUrl, BaseModel, Field, RootModel, ValidationError # noqa: TID251
|
||||
from pydantic import AnyUrl, BaseModel, Field, TypeAdapter, ValidationError
|
||||
from starlette.datastructures import FormData, QueryParams
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import RedirectResponse, Response
|
||||
|
|
@ -20,6 +19,7 @@ from mcp.server.auth.provider import (
|
|||
from mcp.shared.auth import InvalidRedirectUriError, InvalidScopeError
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_ANY_URL_ADAPTER = TypeAdapter(AnyUrl)
|
||||
|
||||
|
||||
class AuthorizationRequest(BaseModel):
|
||||
|
|
@ -59,10 +59,6 @@ def best_effort_extract_string(key: str, params: None | FormData | QueryParams)
|
|||
return None
|
||||
|
||||
|
||||
class AnyUrlModel(RootModel[AnyUrl]):
|
||||
root: AnyUrl
|
||||
|
||||
|
||||
@dataclass
|
||||
class AuthorizationHandler:
|
||||
provider: OAuthAuthorizationServerProvider[Any, Any, Any]
|
||||
|
|
@ -107,9 +103,9 @@ class AuthorizationHandler:
|
|||
if params is not None and "redirect_uri" not in params:
|
||||
raw_redirect_uri = None
|
||||
else:
|
||||
raw_redirect_uri = AnyUrlModel.model_validate(
|
||||
raw_redirect_uri = _ANY_URL_ADAPTER.validate_python(
|
||||
best_effort_extract_string("redirect_uri", params)
|
||||
).root
|
||||
)
|
||||
redirect_uri = client.validate_redirect_uri(raw_redirect_uri)
|
||||
except (ValidationError, InvalidRedirectUriError):
|
||||
# if the redirect URI is invalid, ignore it & just return the
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ from mcp_types import RequestId
|
|||
|
||||
# Internal surface package; imported as the gate's source of truth for spec-valid property schemas.
|
||||
from mcp_types._v2025_11_25 import PrimitiveSchemaDefinition
|
||||
from pydantic import BaseModel, ValidationError
|
||||
from pydantic import BaseModel, TypeAdapter, ValidationError
|
||||
from pydantic.json_schema import GenerateJsonSchema, JsonSchemaValue
|
||||
from pydantic_core import core_schema
|
||||
from typing_extensions import TypeAliasType
|
||||
|
|
@ -16,6 +16,7 @@ from typing_extensions import TypeAliasType
|
|||
from mcp.server.session import ServerSession
|
||||
|
||||
ElicitSchemaModelT = TypeVar("ElicitSchemaModelT", bound=BaseModel)
|
||||
_PRIMITIVE_SCHEMA_ADAPTER = TypeAdapter[PrimitiveSchemaDefinition](PrimitiveSchemaDefinition)
|
||||
|
||||
|
||||
class AcceptedElicitation(BaseModel, Generic[ElicitSchemaModelT]):
|
||||
|
|
@ -79,7 +80,7 @@ def _validate_rendered_properties(json_schema: dict[str, Any]) -> None:
|
|||
"""
|
||||
for field_name, prop in json_schema.get("properties", {}).items():
|
||||
try:
|
||||
PrimitiveSchemaDefinition.model_validate(prop)
|
||||
_PRIMITIVE_SCHEMA_ADAPTER.validate_python(prop)
|
||||
except ValidationError:
|
||||
raise TypeError(
|
||||
f"Elicitation schema field {field_name!r} rendered as {prop!r}, "
|
||||
|
|
|
|||
|
|
@ -114,15 +114,15 @@ 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
|
||||
`output_model` type annotation 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 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[Any], WithJsonSchema(None)] | None = None
|
||||
output_model: Annotated[Any, WithJsonSchema(None)] = None
|
||||
wrap_output: bool = False
|
||||
_adapter: tuple[type[Any], TypeAdapter[Any]] | None = PrivateAttr(default=None)
|
||||
_adapter: tuple[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:
|
||||
|
|
@ -130,7 +130,7 @@ class FuncMetadata(BaseModel):
|
|||
schema = self._output_adapter(self.output_model).json_schema(schema_generator=StrictJsonSchema)
|
||||
self.output_schema = _inline_root_ref(schema)
|
||||
|
||||
def _output_adapter(self, output_model: type[Any]) -> TypeAdapter[Any]:
|
||||
def _output_adapter(self, output_model: 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)))
|
||||
|
|
@ -477,7 +477,7 @@ def func_metadata(
|
|||
return FuncMetadata(arg_model=arguments_model)
|
||||
|
||||
|
||||
def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[type[Any] | None, bool]:
|
||||
def _create_output_model(original_annotation: Any, type_expr: Any, func_name: str) -> tuple[Any, bool]:
|
||||
"""Pick the type structured output is validated against for the given return annotation.
|
||||
|
||||
Args:
|
||||
|
|
@ -491,7 +491,7 @@ def _create_output_model(original_annotation: Any, type_expr: Any, func_name: st
|
|||
Model is None if the type cannot carry structured output.
|
||||
wrap_output is True if the result needs to be wrapped in {"result": ...}
|
||||
"""
|
||||
model: type[Any] | None = None
|
||||
model: Any = None
|
||||
wrap_output = False
|
||||
|
||||
# First handle special case: None
|
||||
|
|
@ -503,13 +503,12 @@ def _create_output_model(original_annotation: Any, type_expr: Any, func_name: st
|
|||
elif isinstance(type_expr, GenericAlias):
|
||||
origin = get_origin(type_expr)
|
||||
|
||||
# Special case: dict with string keys can use RootModel
|
||||
if origin is dict:
|
||||
args = get_args(type_expr)
|
||||
if len(args) == 2 and args[0] is str:
|
||||
# TODO: should we use the original annotation? We are losing any potential `Annotated`
|
||||
# metadata for Pydantic here:
|
||||
model = _create_dict_model(func_name, type_expr)
|
||||
model = Annotated[type_expr, Field(title=f"{func_name}DictOutput")]
|
||||
else:
|
||||
# dict with non-str keys needs wrapping
|
||||
model = _create_wrapped_model(func_name, original_annotation)
|
||||
|
|
@ -622,21 +621,6 @@ def _create_wrapped_model(func_name: str, annotation: Any) -> type[BaseModel]:
|
|||
return create_model(model_name, result=annotation)
|
||||
|
||||
|
||||
def _create_dict_model(func_name: str, dict_annotation: Any) -> type[BaseModel]:
|
||||
"""Create a RootModel for dict[str, T] types."""
|
||||
# TODO(Marcelo): We should not rely on RootModel for this.
|
||||
from pydantic import RootModel # noqa: TID251
|
||||
|
||||
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 _convert_to_content(result: Any) -> list[ContentBlock]:
|
||||
"""Convert a result to a sequence of content objects.
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue