diff --git a/src/fastmcp/utilities/json_schema.py b/src/fastmcp/utilities/json_schema.py index 29574632d..09ecd6750 100644 --- a/src/fastmcp/utilities/json_schema.py +++ b/src/fastmcp/utilities/json_schema.py @@ -1,7 +1,6 @@ from __future__ import annotations import copy -from collections.abc import Mapping, Sequence def _prune_param(schema: dict, param: str) -> dict: @@ -25,32 +24,57 @@ def _prune_param(schema: dict, param: str) -> dict: return schema -def _prune_unused_defs(schema: dict) -> dict: - """Remove unused definitions from the schema.""" - # collect all remaining local $ref targets +def _walk_and_prune( + schema: dict, + prune_defs: bool = False, + prune_titles: bool = False, + prune_additional_properties: bool = False, +) -> dict: + """Walk the schema and optionally prune titles, unused definitions, and additionalProperties: false.""" + # Deep copy to avoid modifying the original + schema = copy.deepcopy(schema) + + # Will only be used if prune_defs is True used_defs: set[str] = set() - def walk(node: object) -> None: # depth-first traversal - if isinstance(node, Mapping): - ref = node.get("$ref") - if isinstance(ref, str) and ref.startswith("#/$defs/"): - used_defs.add(ref.split("/")[-1]) + def walk(node: object) -> None: + if isinstance(node, dict): + # Process $ref for definition tracking + if prune_defs: + ref = node.get("$ref") + if isinstance(ref, str) and ref.startswith("#/$defs/"): + used_defs.add(ref.split("/")[-1]) + + # Remove title if requested + if prune_titles and "title" in node: + node.pop("title") + + # Remove additionalProperties: false at any level if requested + if ( + prune_additional_properties + and node.get("additionalProperties", None) is False + ): + node.pop("additionalProperties") + + # Walk children for v in node.values(): walk(v) - elif isinstance(node, Sequence) and not isinstance(node, str | bytes): + + elif isinstance(node, list): for v in node: walk(v) + # Traverse the schema once walk(schema) - # remove orphaned definitions - - defs = schema.get("$defs", {}) - for def_name in list(defs): - if def_name not in used_defs: - defs.pop(def_name) - if not defs: - schema.pop("$defs", None) + # Remove orphaned definitions if requested + if prune_defs: + defs = schema.get("$defs", {}) + for def_name in list(defs): + if def_name not in used_defs: + defs.pop(def_name) + if not defs: + schema.pop("$defs", None) return schema @@ -67,16 +91,32 @@ def compress_schema( prune_params: list[str] | None = None, prune_defs: bool = True, prune_additional_properties: bool = True, + prune_titles: bool = False, ) -> dict: """ Remove the given parameters from the schema. + Args: + schema: The schema to compress + prune_params: List of parameter names to remove from properties + prune_defs: Whether to remove unused definitions + prune_additional_properties: Whether to remove additionalProperties: false + prune_titles: Whether to remove title fields from the schema """ + # Make a copy so we don't modify the original schema = copy.deepcopy(schema) + + # Remove specific parameters if requested for param in prune_params or []: schema = _prune_param(schema, param=param) - if prune_defs: - schema = _prune_unused_defs(schema) - if prune_additional_properties: - schema = _prune_additional_properties(schema) + + # Do a single walk to handle pruning operations + if prune_defs or prune_titles or prune_additional_properties: + schema = _walk_and_prune( + schema, + prune_defs=prune_defs, + prune_titles=prune_titles, + prune_additional_properties=prune_additional_properties, + ) + return schema diff --git a/tests/utilities/test_json_schema.py b/tests/utilities/test_json_schema.py index cc9cc156c..a1b5f1584 100644 --- a/tests/utilities/test_json_schema.py +++ b/tests/utilities/test_json_schema.py @@ -1,11 +1,21 @@ from fastmcp.utilities.json_schema import ( - _prune_additional_properties, _prune_param, - _prune_unused_defs, + _walk_and_prune, compress_schema, ) +# Create wrappers for backward compatibility with tests +def _prune_unused_defs(schema): + """Wrapper for _walk_and_prune that only prunes definitions.""" + return _walk_and_prune(schema, prune_defs=True) + + +def _prune_additional_properties(schema): + """Wrapper for _walk_and_prune that only prunes additionalProperties: false.""" + return _walk_and_prune(schema, prune_additional_properties=True) + + class TestPruneParam: """Tests for the _prune_param function.""" @@ -244,3 +254,51 @@ class TestCompressSchema: assert "$defs" not in result # Both defs should be gone # Check that additionalProperties was removed assert "additionalProperties" not in result + + def test_prune_titles(self): + """Test pruning title fields.""" + schema = { + "title": "Root Schema", + "type": "object", + "properties": { + "foo": {"title": "Foo Property", "type": "string"}, + "bar": { + "title": "Bar Property", + "type": "object", + "properties": { + "nested": {"title": "Nested Property", "type": "string"} + }, + }, + }, + } + result = compress_schema(schema, prune_titles=True) + assert "title" not in result + assert "title" not in result["properties"]["foo"] + assert "title" not in result["properties"]["bar"] + assert "title" not in result["properties"]["bar"]["properties"]["nested"] + + def test_prune_nested_additional_properties(self): + """Test pruning additionalProperties: false at all levels.""" + schema = { + "type": "object", + "additionalProperties": False, + "properties": { + "foo": { + "type": "object", + "additionalProperties": False, + "properties": { + "nested": { + "type": "object", + "additionalProperties": False, + } + }, + }, + }, + } + result = compress_schema(schema) + assert "additionalProperties" not in result + assert "additionalProperties" not in result["properties"]["foo"] + assert ( + "additionalProperties" + not in result["properties"]["foo"]["properties"]["nested"] + )