Skip to content

Commit b16e716

Browse files
committed
enable tool to explicitly state output schema to override automatic output schema detection
1 parent 3261cbc commit b16e716

File tree

6 files changed

+35
-12
lines changed

6 files changed

+35
-12
lines changed

src/mcp/client/session.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -142,7 +142,9 @@ def __init__(
142142
self._list_roots_callback = list_roots_callback or _default_list_roots_callback
143143
self._logging_callback = logging_callback or _default_logging_callback
144144
self._message_handler = message_handler or _default_message_handler
145-
self._tool_output_validator_provider = tool_output_validator_provider or _default_tool_output_validator
145+
self._tool_output_validator_provider = (
146+
tool_output_validator_provider or _default_tool_output_validator
147+
)
146148

147149
async def initialize(self) -> types.InitializeResult:
148150
sampling = types.SamplingCapability()

src/mcp/server/fastmcp/server.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -324,6 +324,7 @@ def add_tool(
324324
name: str | None = None,
325325
description: str | None = None,
326326
annotations: ToolAnnotations | None = None,
327+
output_schema: dict[str, Any] | None = None,
327328
) -> None:
328329
"""Add a tool to the server.
329330
@@ -335,9 +336,15 @@ def add_tool(
335336
name: Optional name for the tool (defaults to function name)
336337
description: Optional description of what the tool does
337338
annotations: Optional ToolAnnotations providing additional tool information
339+
output_schema: Optional json schema that the tool should output. If
340+
not specified the schema will be inferred automatically
338341
"""
339342
self._tool_manager.add_tool(
340-
fn, name=name, description=description, annotations=annotations
343+
fn,
344+
name=name,
345+
description=description,
346+
annotations=annotations,
347+
output_schema=output_schema,
341348
)
342349

343350
def tool(

src/mcp/server/fastmcp/tools/base.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ def from_function(
4747
description: str | None = None,
4848
context_kwarg: str | None = None,
4949
annotations: ToolAnnotations | None = None,
50+
output_schema: dict[str, Any] | None = None,
5051
) -> Tool:
5152
"""Create a Tool from a function."""
5253
from mcp.server.fastmcp.server import Context
@@ -71,6 +72,7 @@ def from_function(
7172
func_arg_metadata = func_metadata(
7273
fn,
7374
skip_names=[context_kwarg] if context_kwarg is not None else [],
75+
output_schema=output_schema,
7476
)
7577
parameters = func_arg_metadata.arg_model.model_json_schema()
7678
output = func_arg_metadata.output_schema

src/mcp/server/fastmcp/tools/tool_manager.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,10 +37,15 @@ def add_tool(
3737
name: str | None = None,
3838
description: str | None = None,
3939
annotations: ToolAnnotations | None = None,
40+
output_schema: dict[str, Any] | None = None,
4041
) -> Tool:
4142
"""Add a tool to the server."""
4243
tool = Tool.from_function(
43-
fn, name=name, description=description, annotations=annotations
44+
fn,
45+
name=name,
46+
description=description,
47+
annotations=annotations,
48+
output_schema=output_schema,
4449
)
4550
existing = self._tools.get(tool.name)
4651
if existing:

src/mcp/server/fastmcp/utilities/func_metadata.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,9 @@ def pre_parse_json(self, data: dict[str, Any]) -> dict[str, Any]:
112112

113113

114114
def func_metadata(
115-
func: Callable[..., Any], skip_names: Sequence[str] = ()
115+
func: Callable[..., Any],
116+
skip_names: Sequence[str] = (),
117+
output_schema: dict[str, Any] | None = None,
116118
) -> FuncMetadata:
117119
"""Given a function, return metadata including a pydantic model representing its
118120
signature.
@@ -137,7 +139,6 @@ def func_metadata(
137139
sig = _get_typed_signature(func)
138140
params = sig.parameters
139141
dynamic_pydantic_model_params: dict[str, Any] = {}
140-
output_schema: dict[str, Any] | None = None
141142
globalns = getattr(func, "__globals__", {})
142143
for param in params.values():
143144
if param.name.startswith("_"):
@@ -183,14 +184,17 @@ def func_metadata(
183184
__base__=ArgModelBase,
184185
)
185186

186-
# TODO this could be moved to a constant or passed in as param as per skip_names
187-
ignore = [inspect.Parameter.empty, None, types.Image]
188-
if sig.return_annotation not in ignore:
189-
type_schema = TypeAdapter(sig.return_annotation).json_schema()
190-
if type_schema.get("type", None) == "object":
191-
output_schema = type_schema
187+
return_output_schema = output_schema
192188

193-
return FuncMetadata(arg_model=arguments_model, output_schema=output_schema)
189+
if return_output_schema is None:
190+
# TODO this could be moved to a constant or passed in as param as per skip_names
191+
ignore = [inspect.Parameter.empty, None, types.Image]
192+
if sig.return_annotation not in ignore:
193+
type_schema = TypeAdapter(sig.return_annotation).json_schema()
194+
if type_schema.get("type", None) == "object":
195+
return_output_schema = type_schema
196+
197+
return FuncMetadata(arg_model=arguments_model, output_schema=return_output_schema)
194198

195199

196200
def _get_typed_annotation(annotation: Any, globalns: dict[str, Any]) -> Any:

tests/server/fastmcp/test_func_metadata.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -449,6 +449,9 @@ def test_simple_function_output_schema():
449449

450450
assert func_metadata(simple_no_annotation_fun).output_schema is None
451451
assert func_metadata(simple_str_fun).output_schema is None
452+
assert func_metadata(
453+
simple_str_fun, output_schema={"type": "string"}
454+
).output_schema == {"type": "string"}
452455
assert func_metadata(simple_bool_fun).output_schema is None
453456
assert func_metadata(simple_int_fun).output_schema is None
454457
assert func_metadata(simple_float_fun).output_schema is None

0 commit comments

Comments
 (0)