2
2
import json
3
3
from collections .abc import Awaitable , Callable , Sequence
4
4
from dataclasses import asdict , is_dataclass
5
- from typing import (
6
- Annotated ,
7
- Any ,
8
- ForwardRef ,
9
- Literal ,
10
- get_args ,
11
- get_origin ,
12
- get_type_hints ,
13
- )
5
+ from itertools import chain
6
+ from typing import Annotated , Any , ForwardRef , Literal , get_args , get_origin , get_type_hints
14
7
8
+ import pydantic_core
15
9
from pydantic import (
16
10
BaseModel ,
17
11
ConfigDict ,
24
18
from pydantic .fields import FieldInfo
25
19
from pydantic_core import PydanticUndefined
26
20
27
- from mcp .server .fastmcp .exceptions import InvalidSignature , ToolError
21
+ from mcp .server .fastmcp .exceptions import InvalidSignature
28
22
from mcp .server .fastmcp .utilities .logging import get_logger
23
+ from mcp .server .fastmcp .utilities .types import Image
24
+ from mcp .types import ContentBlock , TextContent
29
25
30
26
logger = get_logger (__name__ )
31
27
32
- OutputConversion = Literal ["none" , "wrapped" , "namedtuple" , "class" ]
28
+ OutputConversion = Literal ["none" , "basemodel" , " wrapped" , "namedtuple" , "class" ]
33
29
34
30
35
31
class ArgModelBase (BaseModel ):
@@ -54,16 +50,13 @@ class FuncMetadata(BaseModel):
54
50
arg_model : Annotated [type [ArgModelBase ], WithJsonSchema (None )]
55
51
output_model : Annotated [type [BaseModel ], WithJsonSchema (None )] | None = None
56
52
output_conversion : OutputConversion = "none"
57
- # We can add things in the future like
58
- # - Maybe some args are excluded from attempting to parse from JSON
59
- # - Maybe some args are special (like context) for dependency injection
60
53
61
54
async def call_fn_with_arg_validation (
62
55
self ,
63
- fn : Callable [..., Any ] | Awaitable [Any ],
56
+ fn : Callable [..., Any | Awaitable [Any ] ],
64
57
fn_is_async : bool ,
65
58
arguments_to_validate : dict [str , Any ],
66
- arguments_to_pass_directly : dict [str , Any ] | None ,
59
+ arguments_to_pass_directly : dict [str , Any ] | None = None ,
67
60
) -> Any :
68
61
"""Call the given function with arguments validated and injected.
69
62
@@ -75,36 +68,62 @@ async def call_fn_with_arg_validation(
75
68
arguments_parsed_dict = arguments_parsed_model .model_dump_one_level ()
76
69
77
70
arguments_parsed_dict |= arguments_to_pass_directly or {}
78
-
79
71
if fn_is_async :
80
- if isinstance (fn , Awaitable ):
81
- return await fn
82
72
return await fn (** arguments_parsed_dict )
83
- if isinstance (fn , Callable ):
84
- return fn (** arguments_parsed_dict )
85
- raise TypeError ("fn must be either Callable or Awaitable" )
73
+ return fn (** arguments_parsed_dict )
74
+
75
+ async def call_fn (
76
+ self ,
77
+ fn : Callable [..., Any | Awaitable [Any ]],
78
+ fn_is_async : bool ,
79
+ args : dict [str , Any ],
80
+ implicit_args : dict [str , Any ] | None = None ,
81
+ ) -> Any :
82
+ parsed_args = self .arg_model .model_construct (** args ).model_dump_one_level ()
83
+ kwargs = parsed_args | (implicit_args or {})
84
+ if fn_is_async :
85
+ return await fn (** kwargs )
86
+ else :
87
+ return fn (** kwargs )
86
88
87
- def to_validated_dict (self , result : Any ) -> dict [str , Any ]:
88
- """Validate and convert the result to a dict after validation."""
89
- if self .output_model is None :
90
- raise ValueError ("No output model to validate against" )
89
+ def convert_result (self , result : Any ) -> Sequence [ContentBlock ] | dict [str , Any ]:
90
+ """convert result to dict"""
91
+ if self .output_model :
92
+ return self ._convert_structured_result (result )
93
+ else :
94
+ return self ._convert_unstructured_result (result )
91
95
96
+ def _convert_structured_result (self , result : Any ) -> dict [str , Any ]:
92
97
match self .output_conversion :
98
+ case "none" :
99
+ return result
100
+ case "basemodel" :
101
+ return result .model_dump ()
93
102
case "wrapped" :
94
- converted = _convert_wrapped_result ( result )
103
+ return { "result" : result }
95
104
case "namedtuple" :
96
- converted = _convert_namedtuple_result ( result )
105
+ return result . _asdict ( )
97
106
case "class" :
98
- converted = _convert_class_result (result )
99
- case "none" :
100
- converted = result
107
+ if is_dataclass ( result ) and not isinstance (result , type ):
108
+ return asdict ( result )
109
+ return dict ( vars ( result ))
101
110
102
- try :
103
- validated = self .output_model .model_validate (converted )
104
- except Exception as e :
105
- raise ToolError (f"Output validation failed: { e } " ) from e
111
+ def _convert_unstructured_result (self , result : Any ) -> Sequence [ContentBlock ]:
112
+ if result is None :
113
+ return []
106
114
107
- return validated .model_dump ()
115
+ if isinstance (result , ContentBlock ):
116
+ return [result ]
117
+
118
+ if isinstance (result , Image ):
119
+ return [result .to_image_content ()]
120
+
121
+ if isinstance (result , list | tuple ):
122
+ return list (chain .from_iterable (self ._convert_unstructured_result (item ) for item in result )) # type: ignore
123
+
124
+ if not isinstance (result , str ):
125
+ result = pydantic_core .to_json (result , fallback = str , indent = 2 ).decode ()
126
+ return [TextContent (type = "text" , text = result )]
108
127
109
128
def pre_parse_json (self , data : dict [str , Any ]) -> dict [str , Any ]:
110
129
"""Pre-parse data from JSON.
@@ -177,6 +196,7 @@ def func_metadata(
177
196
A FuncMetadata object containing:
178
197
- arg_model: A pydantic model representing the function's arguments
179
198
- output_model: A pydantic model for the return type (if structured_output=True)
199
+ - output_conversion: Records how function output should be converted before returning.
180
200
"""
181
201
sig = _get_typed_signature (func )
182
202
params = sig .parameters
@@ -236,7 +256,7 @@ def func_metadata(
236
256
elif isinstance (annotation , type ):
237
257
if issubclass (annotation , BaseModel ):
238
258
output_model = annotation
239
- output_conversion = "none "
259
+ output_conversion = "basemodel "
240
260
elif _is_typeddict (annotation ):
241
261
output_model = _create_model_from_typeddict (annotation , globalns )
242
262
output_conversion = "none"
@@ -338,13 +358,6 @@ def _create_model_from_class(cls: type[Any], globalns: dict[str, Any]) -> type[B
338
358
return create_model (cls .__name__ , ** model_fields , __base__ = BaseModel )
339
359
340
360
341
- def _convert_class_result (result : Any ) -> dict [str , Any ]:
342
- if is_dataclass (result ) and not isinstance (result , type ):
343
- return asdict (result )
344
-
345
- return dict (vars (result ))
346
-
347
-
348
361
def _create_model_from_typeddict (td_type : type [Any ], globalns : dict [str , Any ]) -> type [BaseModel ]:
349
362
"""Create a Pydantic model from a TypedDict.
350
363
@@ -387,10 +400,6 @@ def _create_model_from_namedtuple(nt_type: type[Any], globalns: dict[str, Any])
387
400
return create_model (nt_type .__name__ , ** model_fields , __base__ = BaseModel )
388
401
389
402
390
- def _convert_namedtuple_result (result : Any ) -> dict [str , Any ]:
391
- return result ._asdict ()
392
-
393
-
394
403
def _create_wrapped_model (func_name : str , annotation : Any , field_info : FieldInfo ) -> type [BaseModel ]:
395
404
"""Create a model that wraps a type in a 'result' field.
396
405
@@ -405,10 +414,6 @@ def _create_wrapped_model(func_name: str, annotation: Any, field_info: FieldInfo
405
414
return create_model (model_name , result = (annotation , field_info ), __base__ = BaseModel )
406
415
407
416
408
- def _convert_wrapped_result (result : Any ) -> dict [str , Any ]:
409
- return {"result" : result }
410
-
411
-
412
417
def _create_dict_model (func_name : str , dict_annotation : Any ) -> type [BaseModel ]:
413
418
"""Create a RootModel for dict[str, T] types."""
414
419
0 commit comments