Jeremiah Lowin commited on
Commit
c903d7b
·
1 Parent(s): 4d797e0

Add output schema to tool decorator

Browse files
src/fastmcp/server/server.py CHANGED
@@ -768,6 +768,7 @@ class FastMCP(Generic[LifespanResultT]):
768
  name: str | None = None,
769
  description: str | None = None,
770
  tags: set[str] | None = None,
 
771
  annotations: ToolAnnotations | dict[str, Any] | None = None,
772
  exclude_args: list[str] | None = None,
773
  enabled: bool | None = None,
@@ -781,6 +782,7 @@ class FastMCP(Generic[LifespanResultT]):
781
  name: str | None = None,
782
  description: str | None = None,
783
  tags: set[str] | None = None,
 
784
  annotations: ToolAnnotations | dict[str, Any] | None = None,
785
  exclude_args: list[str] | None = None,
786
  enabled: bool | None = None,
@@ -793,6 +795,7 @@ class FastMCP(Generic[LifespanResultT]):
793
  name: str | None = None,
794
  description: str | None = None,
795
  tags: set[str] | None = None,
 
796
  annotations: ToolAnnotations | dict[str, Any] | None = None,
797
  exclude_args: list[str] | None = None,
798
  enabled: bool | None = None,
@@ -815,6 +818,7 @@ class FastMCP(Generic[LifespanResultT]):
815
  name: Optional name for the tool (keyword-only, alternative to name_or_fn)
816
  description: Optional description of what the tool does
817
  tags: Optional set of tags for categorizing the tool
 
818
  annotations: Optional annotations about the tool's behavior (e.g. {"is_async": True})
819
  exclude_args: Optional list of argument names to exclude from the tool schema
820
  enabled: Optional boolean to enable or disable the tool
@@ -867,6 +871,7 @@ class FastMCP(Generic[LifespanResultT]):
867
  name=tool_name,
868
  description=description,
869
  tags=tags,
 
870
  annotations=annotations,
871
  exclude_args=exclude_args,
872
  serializer=self._tool_serializer,
@@ -897,6 +902,7 @@ class FastMCP(Generic[LifespanResultT]):
897
  name=tool_name,
898
  description=description,
899
  tags=tags,
 
900
  annotations=annotations,
901
  exclude_args=exclude_args,
902
  enabled=enabled,
 
768
  name: str | None = None,
769
  description: str | None = None,
770
  tags: set[str] | None = None,
771
+ output_schema: dict[str, Any] | None = None,
772
  annotations: ToolAnnotations | dict[str, Any] | None = None,
773
  exclude_args: list[str] | None = None,
774
  enabled: bool | None = None,
 
782
  name: str | None = None,
783
  description: str | None = None,
784
  tags: set[str] | None = None,
785
+ output_schema: dict[str, Any] | None = None,
786
  annotations: ToolAnnotations | dict[str, Any] | None = None,
787
  exclude_args: list[str] | None = None,
788
  enabled: bool | None = None,
 
795
  name: str | None = None,
796
  description: str | None = None,
797
  tags: set[str] | None = None,
798
+ output_schema: dict[str, Any] | None = None,
799
  annotations: ToolAnnotations | dict[str, Any] | None = None,
800
  exclude_args: list[str] | None = None,
801
  enabled: bool | None = None,
 
818
  name: Optional name for the tool (keyword-only, alternative to name_or_fn)
819
  description: Optional description of what the tool does
820
  tags: Optional set of tags for categorizing the tool
821
+ output_schema: Optional JSON schema for the tool's output
822
  annotations: Optional annotations about the tool's behavior (e.g. {"is_async": True})
823
  exclude_args: Optional list of argument names to exclude from the tool schema
824
  enabled: Optional boolean to enable or disable the tool
 
871
  name=tool_name,
872
  description=description,
873
  tags=tags,
874
+ output_schema=output_schema,
875
  annotations=annotations,
876
  exclude_args=exclude_args,
877
  serializer=self._tool_serializer,
 
902
  name=tool_name,
903
  description=description,
904
  tags=tags,
905
+ output_schema=output_schema,
906
  annotations=annotations,
907
  exclude_args=exclude_args,
908
  enabled=enabled,
tests/server/test_server.py CHANGED
@@ -390,6 +390,15 @@ class TestToolDecorator:
390
  def my_function(x: int) -> str:
391
  return f"Result: {x}"
392
 
 
 
 
 
 
 
 
 
 
393
 
394
  class TestResourceDecorator:
395
  async def test_no_resources_before_decorator(self):
 
390
  def my_function(x: int) -> str:
391
  return f"Result: {x}"
392
 
393
+ async def test_tool_decorator_with_output_schema(self):
394
+ mcp = FastMCP()
395
+
396
+ @mcp.tool(output_schema={"type": "integer"})
397
+ def my_function(x: int) -> str:
398
+ return f"Result: {x}"
399
+
400
+ assert my_function.output_schema == {"type": "integer"}
401
+
402
 
403
  class TestResourceDecorator:
404
  async def test_no_resources_before_decorator(self):
tests/tools/test_tool.py CHANGED
@@ -369,6 +369,111 @@ class TestToolFromFunctionOutputSchema:
369
  tool = Tool.from_function(func)
370
  assert tool.output_schema is None
371
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
372
 
373
  class TestLegacyToolJsonParsing:
374
  """Tests for Tool's JSON pre-parsing functionality."""
 
369
  tool = Tool.from_function(func)
370
  assert tool.output_schema is None
371
 
372
+ async def test_provided_output_schema_takes_precedence_over_json_compatible_annotation(
373
+ self,
374
+ ):
375
+ """Test that provided output_schema takes precedence over inferred schema from JSON-compatible annotation."""
376
+
377
+ def func() -> dict[str, int]:
378
+ return {"a": 1, "b": 2}
379
+
380
+ # Provide a custom output schema that differs from the inferred one
381
+ custom_schema = {"type": "string", "description": "Custom schema"}
382
+
383
+ tool = Tool.from_function(func, output_schema=custom_schema)
384
+ assert tool.output_schema == custom_schema
385
+
386
+ async def test_provided_output_schema_takes_precedence_over_complex_annotation(
387
+ self,
388
+ ):
389
+ """Test that provided output_schema takes precedence over inferred schema from complex annotation."""
390
+
391
+ def func() -> list[dict[str, int | float]]:
392
+ return [{"a": 1, "b": 2.5}]
393
+
394
+ # Provide a custom output schema that differs from the inferred one
395
+ custom_schema = {"type": "object", "properties": {"custom": {"type": "string"}}}
396
+
397
+ tool = Tool.from_function(func, output_schema=custom_schema)
398
+ assert tool.output_schema == custom_schema
399
+
400
+ async def test_provided_output_schema_takes_precedence_over_unserializable_annotation(
401
+ self,
402
+ ):
403
+ """Test that provided output_schema takes precedence over None schema from unserializable annotation."""
404
+
405
+ class Unserializable:
406
+ def __init__(self, data: Any):
407
+ self.data = data
408
+
409
+ def func() -> Unserializable:
410
+ return Unserializable(data="test")
411
+
412
+ # Provide a custom output schema even though the annotation is unserializable
413
+ custom_schema = {"type": "array", "items": {"type": "string"}}
414
+
415
+ tool = Tool.from_function(func, output_schema=custom_schema)
416
+ assert tool.output_schema == custom_schema
417
+
418
+ async def test_provided_output_schema_takes_precedence_over_no_annotation(self):
419
+ """Test that provided output_schema takes precedence over None schema from no annotation."""
420
+
421
+ def func():
422
+ return "hello"
423
+
424
+ # Provide a custom output schema even though there's no return annotation
425
+ custom_schema = {"type": "number", "minimum": 0}
426
+
427
+ tool = Tool.from_function(func, output_schema=custom_schema)
428
+ assert tool.output_schema == custom_schema
429
+
430
+ async def test_provided_output_schema_takes_precedence_over_converted_annotation(
431
+ self,
432
+ ):
433
+ """Test that provided output_schema takes precedence over converted schema from Image/Audio/File annotations."""
434
+
435
+ def func() -> Image:
436
+ return Image(data=b"test")
437
+
438
+ # Provide a custom output schema that differs from the converted ImageContent schema
439
+ custom_schema = {
440
+ "type": "object",
441
+ "properties": {"custom_image": {"type": "string"}},
442
+ }
443
+
444
+ tool = Tool.from_function(func, output_schema=custom_schema)
445
+ assert tool.output_schema == custom_schema
446
+
447
+ async def test_provided_output_schema_takes_precedence_over_union_annotation(self):
448
+ """Test that provided output_schema takes precedence over inferred schema from union annotation."""
449
+
450
+ def func() -> str | int | None:
451
+ return "hello"
452
+
453
+ # Provide a custom output schema that differs from the inferred union schema
454
+ custom_schema = {"type": "boolean"}
455
+
456
+ tool = Tool.from_function(func, output_schema=custom_schema)
457
+ assert tool.output_schema == custom_schema
458
+
459
+ async def test_provided_output_schema_takes_precedence_over_pydantic_annotation(
460
+ self,
461
+ ):
462
+ """Test that provided output_schema takes precedence over inferred schema from Pydantic model annotation."""
463
+
464
+ class Person(BaseModel):
465
+ name: str
466
+ age: int
467
+
468
+ def func() -> Person:
469
+ return Person(name="John", age=30)
470
+
471
+ # Provide a custom output schema that differs from the inferred Person schema
472
+ custom_schema = {"type": "array", "items": {"type": "number"}}
473
+
474
+ tool = Tool.from_function(func, output_schema=custom_schema)
475
+ assert tool.output_schema == custom_schema
476
+
477
 
478
  class TestLegacyToolJsonParsing:
479
  """Tests for Tool's JSON pre-parsing functionality."""