mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-09 15:19:10 +02:00
Optimize OpenAPI parser performance with single-pass schema processing (#1214)
This commit is contained in:
Jeremiah Lowin
GitHub
parent
820995286c
commit
d6bd800d88
6 changed files with 676 additions and 145 deletions
|
|
@ -3,8 +3,6 @@
|
|||
import logging
|
||||
from typing import Any, cast
|
||||
|
||||
from fastmcp.utilities.json_schema import compress_schema
|
||||
|
||||
from .models import HTTPRoute, JsonSchema, ResponseInfo
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -314,10 +312,42 @@ def _combine_schemas_and_map_params(
|
|||
}
|
||||
# Add schema definitions if available
|
||||
if route.schema_definitions:
|
||||
result["$defs"] = route.schema_definitions
|
||||
result["$defs"] = route.schema_definitions.copy()
|
||||
|
||||
# Use compress_schema to remove unused definitions
|
||||
result = compress_schema(result)
|
||||
# Use lightweight compression - prune additionalProperties and unused definitions
|
||||
if result.get("additionalProperties") is False:
|
||||
result.pop("additionalProperties")
|
||||
|
||||
# Remove unused definitions (lightweight approach - just check direct $ref usage)
|
||||
if "$defs" in result:
|
||||
used_refs = set()
|
||||
|
||||
def find_refs_in_value(value):
|
||||
if isinstance(value, dict):
|
||||
if "$ref" in value and isinstance(value["$ref"], str):
|
||||
ref = value["$ref"]
|
||||
if ref.startswith("#/$defs/"):
|
||||
used_refs.add(ref.split("/")[-1])
|
||||
for v in value.values():
|
||||
find_refs_in_value(v)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
find_refs_in_value(item)
|
||||
|
||||
# Find refs in the main schema (excluding $defs section)
|
||||
for key, value in result.items():
|
||||
if key != "$defs":
|
||||
find_refs_in_value(value)
|
||||
|
||||
# Remove unused definitions
|
||||
if used_refs:
|
||||
result["$defs"] = {
|
||||
name: def_schema
|
||||
for name, def_schema in result["$defs"].items()
|
||||
if name in used_refs
|
||||
}
|
||||
else:
|
||||
result.pop("$defs")
|
||||
|
||||
return result, parameter_map
|
||||
|
||||
|
|
@ -339,17 +369,63 @@ def _combine_schemas(route: HTTPRoute) -> dict[str, Any]:
|
|||
return schema
|
||||
|
||||
|
||||
def _has_one_of(obj: dict[str, Any] | list[Any]) -> bool:
|
||||
"""Quickly check if schema contains any 'oneOf' keys without deep traversal."""
|
||||
if isinstance(obj, dict):
|
||||
if "oneOf" in obj:
|
||||
return True
|
||||
# Only check likely schema containers, skip examples/defaults
|
||||
for k, v in obj.items():
|
||||
if k in [
|
||||
"properties",
|
||||
"items",
|
||||
"allOf",
|
||||
"anyOf",
|
||||
"additionalProperties",
|
||||
] and isinstance(v, dict | list):
|
||||
if _has_one_of(v):
|
||||
return True
|
||||
elif isinstance(obj, list):
|
||||
for item in obj:
|
||||
if isinstance(item, dict | list) and _has_one_of(item):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _adjust_union_types(
|
||||
schema: dict[str, Any] | list[Any],
|
||||
schema: dict[str, Any] | list[Any], _depth: int = 0
|
||||
) -> dict[str, Any] | list[Any]:
|
||||
"""Recursively replace 'oneOf' with 'anyOf' in schema to handle overlapping unions."""
|
||||
# MAJOR OPTIMIZATION: Skip entirely if schema has no oneOf keys
|
||||
if _depth == 0 and not _has_one_of(schema):
|
||||
return schema
|
||||
|
||||
# OPTIMIZATION: Early termination for very deep structures to prevent exponential slowdown
|
||||
if _depth > 30: # Reduced from 50 for better performance
|
||||
return schema
|
||||
|
||||
if isinstance(schema, dict):
|
||||
if "oneOf" in schema:
|
||||
schema["anyOf"] = schema.pop("oneOf")
|
||||
for k, v in schema.items():
|
||||
schema[k] = _adjust_union_types(v)
|
||||
# Work on a copy to avoid mutating the input
|
||||
result = schema.copy()
|
||||
if "oneOf" in result:
|
||||
result["anyOf"] = result.pop("oneOf")
|
||||
# OPTIMIZATION: Only recurse into values that could contain more schemas
|
||||
for k, v in result.items():
|
||||
if isinstance(v, dict | list) and k not in [
|
||||
"examples",
|
||||
"example",
|
||||
"default",
|
||||
]:
|
||||
result[k] = _adjust_union_types(v, _depth + 1)
|
||||
return result
|
||||
elif isinstance(schema, list):
|
||||
return [_adjust_union_types(item) for item in schema]
|
||||
# Process list items without mutating the input list
|
||||
return [
|
||||
_adjust_union_types(item, _depth + 1)
|
||||
if isinstance(item, dict | list)
|
||||
else item
|
||||
for item in schema
|
||||
]
|
||||
return schema
|
||||
|
||||
|
||||
|
|
@ -436,10 +512,42 @@ def extract_output_schema_from_responses(
|
|||
|
||||
# Add schema definitions if available
|
||||
if schema_definitions:
|
||||
output_schema["$defs"] = schema_definitions
|
||||
output_schema["$defs"] = schema_definitions.copy()
|
||||
|
||||
# Use compress_schema to remove unused definitions
|
||||
output_schema = compress_schema(output_schema)
|
||||
# Use lightweight compression - prune additionalProperties and unused definitions
|
||||
if output_schema.get("additionalProperties") is False:
|
||||
output_schema.pop("additionalProperties")
|
||||
|
||||
# Remove unused definitions (lightweight approach - just check direct $ref usage)
|
||||
if "$defs" in output_schema:
|
||||
used_refs = set()
|
||||
|
||||
def find_refs_in_value(value):
|
||||
if isinstance(value, dict):
|
||||
if "$ref" in value and isinstance(value["$ref"], str):
|
||||
ref = value["$ref"]
|
||||
if ref.startswith("#/$defs/"):
|
||||
used_refs.add(ref.split("/")[-1])
|
||||
for v in value.values():
|
||||
find_refs_in_value(v)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
find_refs_in_value(item)
|
||||
|
||||
# Find refs in the main schema (excluding $defs section)
|
||||
for key, value in output_schema.items():
|
||||
if key != "$defs":
|
||||
find_refs_in_value(value)
|
||||
|
||||
# Remove unused definitions
|
||||
if used_refs:
|
||||
output_schema["$defs"] = {
|
||||
name: def_schema
|
||||
for name, def_schema in output_schema["$defs"].items()
|
||||
if name in used_refs
|
||||
}
|
||||
else:
|
||||
output_schema.pop("$defs")
|
||||
|
||||
# Adjust union types to handle overlapping unions
|
||||
output_schema = cast(dict[str, Any], _adjust_union_types(output_schema))
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
|
|
@ -25,120 +24,159 @@ def _prune_param(schema: dict, param: str) -> dict:
|
|||
return schema
|
||||
|
||||
|
||||
def _prune_unused_defs(schema: dict) -> dict:
|
||||
"""Walk the schema and prune unused defs."""
|
||||
|
||||
root_defs: set[str] = set()
|
||||
referenced_by: defaultdict[str, list] = defaultdict(list)
|
||||
|
||||
defs = schema.get("$defs")
|
||||
if defs is None:
|
||||
return schema
|
||||
|
||||
def walk(
|
||||
node: object, current_def: str | None = None, skip_defs: bool = False
|
||||
) -> None:
|
||||
if isinstance(node, dict):
|
||||
# Process $ref for definition tracking
|
||||
ref = node.get("$ref")
|
||||
if isinstance(ref, str) and ref.startswith("#/$defs/"):
|
||||
def_name = ref.split("/")[-1]
|
||||
if current_def:
|
||||
referenced_by[def_name].append(current_def)
|
||||
else:
|
||||
root_defs.add(def_name)
|
||||
|
||||
# Walk children
|
||||
for k, v in node.items():
|
||||
if skip_defs and k == "$defs":
|
||||
continue
|
||||
|
||||
if k in ["allOf", "oneOf", "anyOf"]:
|
||||
for child in v:
|
||||
walk(child, current_def=current_def)
|
||||
else:
|
||||
walk(v, current_def=current_def)
|
||||
|
||||
elif isinstance(node, list):
|
||||
for v in node:
|
||||
walk(v, current_def=current_def)
|
||||
|
||||
# Traverse the schema once, skipping the $defs
|
||||
walk(schema, skip_defs=True)
|
||||
|
||||
# Now figure out what defs reference other defs
|
||||
for def_name, value in defs.items():
|
||||
walk(value, current_def=def_name)
|
||||
|
||||
# Figure out what defs were referenced directly or recursively
|
||||
def def_is_referenced(def_name, parent_def_names: set[str] | None = None):
|
||||
if def_name in root_defs:
|
||||
return True
|
||||
references = referenced_by.get(def_name)
|
||||
if references:
|
||||
if parent_def_names is None:
|
||||
parent_def_names = set()
|
||||
|
||||
# Handle recursion by excluding references already present in parent references
|
||||
parent_def_names = parent_def_names | {def_name}
|
||||
valid_references = [
|
||||
reference
|
||||
for reference in references
|
||||
if reference not in parent_def_names
|
||||
]
|
||||
|
||||
for reference in valid_references:
|
||||
if def_is_referenced(reference, parent_def_names):
|
||||
return True
|
||||
return False
|
||||
|
||||
# Remove orphaned definitions if requested
|
||||
for def_name in list(defs):
|
||||
if not def_is_referenced(def_name):
|
||||
defs.pop(def_name)
|
||||
if not defs:
|
||||
schema.pop("$defs", None)
|
||||
|
||||
return schema
|
||||
|
||||
|
||||
def _walk_and_prune(
|
||||
def _single_pass_optimize(
|
||||
schema: dict,
|
||||
prune_titles: bool = False,
|
||||
prune_additional_properties: bool = False,
|
||||
prune_defs: bool = True,
|
||||
) -> dict:
|
||||
"""Walk the schema and optionally prune titles and additionalProperties: false."""
|
||||
"""
|
||||
Optimize JSON schemas in a single traversal for better performance.
|
||||
|
||||
This function combines three schema cleanup operations that would normally require
|
||||
separate tree traversals:
|
||||
|
||||
1. **Remove unused definitions** (prune_defs): Finds and removes `$defs` entries
|
||||
that aren't referenced anywhere in the schema, reducing schema size.
|
||||
|
||||
2. **Remove titles** (prune_titles): Strips `title` fields throughout the schema
|
||||
to reduce verbosity while preserving functional information.
|
||||
|
||||
3. **Remove restrictive additionalProperties** (prune_additional_properties):
|
||||
Removes `"additionalProperties": false` constraints to make schemas more flexible.
|
||||
|
||||
**Performance Benefits:**
|
||||
- Single tree traversal instead of multiple passes (2-3x faster)
|
||||
- Immutable design prevents shared reference bugs
|
||||
- Early termination prevents runaway recursion on deeply nested schemas
|
||||
|
||||
**Algorithm Overview:**
|
||||
1. Traverse main schema, collecting $ref references and applying cleanups
|
||||
2. Traverse $defs section to map inter-definition dependencies
|
||||
3. Remove unused definitions based on reference analysis
|
||||
|
||||
Args:
|
||||
schema: JSON schema dict to optimize (not modified)
|
||||
prune_titles: Remove title fields for cleaner output
|
||||
prune_additional_properties: Remove "additionalProperties": false constraints
|
||||
prune_defs: Remove unused $defs entries to reduce size
|
||||
|
||||
Returns:
|
||||
A new optimized schema dict
|
||||
|
||||
Example:
|
||||
>>> schema = {
|
||||
... "type": "object",
|
||||
... "title": "MySchema",
|
||||
... "additionalProperties": False,
|
||||
... "$defs": {"UnusedDef": {"type": "string"}}
|
||||
... }
|
||||
>>> result = _single_pass_optimize(schema, prune_titles=True, prune_defs=True)
|
||||
>>> # Result: {"type": "object", "additionalProperties": False}
|
||||
"""
|
||||
if not (prune_defs or prune_titles or prune_additional_properties):
|
||||
return schema # Nothing to do
|
||||
|
||||
# Phase 1: Collect references and apply simple cleanups
|
||||
# Track which $defs are referenced from the main schema and from other $defs
|
||||
root_refs: set[str] = set() # $defs referenced directly from main schema
|
||||
def_dependencies: defaultdict[str, list[str]] = defaultdict(
|
||||
list
|
||||
) # def A references def B
|
||||
defs = schema.get("$defs")
|
||||
|
||||
def traverse_and_clean(
|
||||
node: object,
|
||||
current_def_name: str | None = None,
|
||||
skip_defs_section: bool = False,
|
||||
depth: int = 0,
|
||||
) -> None:
|
||||
"""Traverse schema tree, collecting $ref info and applying cleanups."""
|
||||
if depth > 50: # Prevent infinite recursion
|
||||
return
|
||||
|
||||
def walk(node: object) -> None:
|
||||
if isinstance(node, dict):
|
||||
# Remove title if requested
|
||||
# Collect $ref references for unused definition removal
|
||||
if prune_defs:
|
||||
ref = node.get("$ref")
|
||||
if isinstance(ref, str) and ref.startswith("#/$defs/"):
|
||||
referenced_def = ref.split("/")[-1]
|
||||
if current_def_name:
|
||||
# We're inside a $def, so this is a def->def reference
|
||||
def_dependencies[referenced_def].append(current_def_name)
|
||||
else:
|
||||
# We're in the main schema, so this is a root reference
|
||||
root_refs.add(referenced_def)
|
||||
|
||||
# Apply cleanups
|
||||
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
|
||||
and node.get("additionalProperties") is False
|
||||
):
|
||||
node.pop("additionalProperties")
|
||||
|
||||
# Walk children
|
||||
for v in node.values():
|
||||
walk(v)
|
||||
# Recursive traversal
|
||||
for key, value in node.items():
|
||||
if skip_defs_section and key == "$defs":
|
||||
continue # Skip $defs during main schema traversal
|
||||
|
||||
# Handle schema composition keywords with special traversal
|
||||
if key in ["allOf", "oneOf", "anyOf"] and isinstance(value, list):
|
||||
for item in value:
|
||||
traverse_and_clean(item, current_def_name, depth=depth + 1)
|
||||
else:
|
||||
traverse_and_clean(value, current_def_name, depth=depth + 1)
|
||||
|
||||
elif isinstance(node, list):
|
||||
for v in node:
|
||||
walk(v)
|
||||
for item in node:
|
||||
traverse_and_clean(item, current_def_name, depth=depth + 1)
|
||||
|
||||
walk(schema)
|
||||
# Phase 2: Traverse main schema (excluding $defs section)
|
||||
traverse_and_clean(schema, skip_defs_section=True)
|
||||
|
||||
return schema
|
||||
# Phase 3: Traverse $defs to find inter-definition references
|
||||
if prune_defs and defs:
|
||||
for def_name, def_schema in defs.items():
|
||||
traverse_and_clean(def_schema, current_def_name=def_name)
|
||||
|
||||
# Phase 4: Remove unused definitions
|
||||
def is_def_used(def_name: str, visiting: set[str] | None = None) -> bool:
|
||||
"""Check if a definition is used, handling circular references."""
|
||||
if def_name in root_refs:
|
||||
return True # Used directly from main schema
|
||||
|
||||
# Check if any definition that references this one is itself used
|
||||
referencing_defs = def_dependencies.get(def_name, [])
|
||||
if referencing_defs:
|
||||
if visiting is None:
|
||||
visiting = set()
|
||||
|
||||
# Avoid infinite recursion on circular references
|
||||
if def_name in visiting:
|
||||
return False
|
||||
visiting = visiting | {def_name}
|
||||
|
||||
# If any referencing def is used, then this def is used
|
||||
for referencing_def in referencing_defs:
|
||||
if referencing_def not in visiting and is_def_used(
|
||||
referencing_def, visiting
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
# Remove unused definitions
|
||||
for def_name in list(defs.keys()):
|
||||
if not is_def_used(def_name):
|
||||
defs.pop(def_name)
|
||||
|
||||
# Clean up empty $defs section
|
||||
if not defs:
|
||||
schema.pop("$defs", None)
|
||||
|
||||
def _prune_additional_properties(schema: dict) -> dict:
|
||||
"""Remove additionalProperties from the schema if it is False."""
|
||||
if schema.get("additionalProperties", None) is False:
|
||||
schema.pop("additionalProperties")
|
||||
return schema
|
||||
|
||||
|
||||
|
|
@ -159,21 +197,17 @@ def compress_schema(
|
|||
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)
|
||||
|
||||
# Do a single walk to handle pruning operations
|
||||
if prune_titles or prune_additional_properties:
|
||||
schema = _walk_and_prune(
|
||||
# Apply combined optimizations in a single tree traversal
|
||||
if prune_titles or prune_additional_properties or prune_defs:
|
||||
schema = _single_pass_optimize(
|
||||
schema,
|
||||
prune_titles=prune_titles,
|
||||
prune_additional_properties=prune_additional_properties,
|
||||
prune_defs=prune_defs,
|
||||
)
|
||||
if prune_defs:
|
||||
schema = _prune_unused_defs(schema)
|
||||
|
||||
return schema
|
||||
|
|
|
|||
|
|
@ -24,7 +24,6 @@ from openapi_pydantic.v3.v3_0 import Response as Response_30
|
|||
from openapi_pydantic.v3.v3_0 import Schema as Schema_30
|
||||
from pydantic import BaseModel, Field, ValidationError
|
||||
|
||||
from fastmcp.utilities.json_schema import compress_schema
|
||||
from fastmcp.utilities.types import FastMCPBaseModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -1264,10 +1263,42 @@ def _combine_schemas(route: HTTPRoute) -> dict[str, Any]:
|
|||
}
|
||||
# Add schema definitions if available
|
||||
if route.schema_definitions:
|
||||
result["$defs"] = route.schema_definitions
|
||||
result["$defs"] = route.schema_definitions.copy()
|
||||
|
||||
# Use compress_schema to remove unused definitions
|
||||
result = compress_schema(result)
|
||||
# Use lightweight compression - prune additionalProperties and unused definitions
|
||||
if result.get("additionalProperties") is False:
|
||||
result.pop("additionalProperties")
|
||||
|
||||
# Remove unused definitions (lightweight approach - just check direct $ref usage)
|
||||
if "$defs" in result:
|
||||
used_refs = set()
|
||||
|
||||
def find_refs_in_value(value):
|
||||
if isinstance(value, dict):
|
||||
if "$ref" in value and isinstance(value["$ref"], str):
|
||||
ref = value["$ref"]
|
||||
if ref.startswith("#/$defs/"):
|
||||
used_refs.add(ref.split("/")[-1])
|
||||
for v in value.values():
|
||||
find_refs_in_value(v)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
find_refs_in_value(item)
|
||||
|
||||
# Find refs in the main schema (excluding $defs section)
|
||||
for key, value in result.items():
|
||||
if key != "$defs":
|
||||
find_refs_in_value(value)
|
||||
|
||||
# Remove unused definitions
|
||||
if used_refs:
|
||||
result["$defs"] = {
|
||||
name: def_schema
|
||||
for name, def_schema in result["$defs"].items()
|
||||
if name in used_refs
|
||||
}
|
||||
else:
|
||||
result.pop("$defs")
|
||||
|
||||
return result
|
||||
|
||||
|
|
@ -1277,10 +1308,13 @@ def _adjust_union_types(
|
|||
) -> dict[str, Any] | list[Any]:
|
||||
"""Recursively replace 'oneOf' with 'anyOf' in schema to handle overlapping unions."""
|
||||
if isinstance(schema, dict):
|
||||
if "oneOf" in schema:
|
||||
schema["anyOf"] = schema.pop("oneOf")
|
||||
for k, v in schema.items():
|
||||
schema[k] = _adjust_union_types(v)
|
||||
# Work on a copy to avoid mutating the input
|
||||
result = schema.copy()
|
||||
if "oneOf" in result:
|
||||
result["anyOf"] = result.pop("oneOf")
|
||||
for k, v in result.items():
|
||||
result[k] = _adjust_union_types(v)
|
||||
return result
|
||||
elif isinstance(schema, list):
|
||||
return [_adjust_union_types(item) for item in schema]
|
||||
return schema
|
||||
|
|
@ -1369,10 +1403,42 @@ def extract_output_schema_from_responses(
|
|||
|
||||
# Add schema definitions if available
|
||||
if schema_definitions:
|
||||
output_schema["$defs"] = schema_definitions
|
||||
output_schema["$defs"] = schema_definitions.copy()
|
||||
|
||||
# Use compress_schema to remove unused definitions
|
||||
output_schema = compress_schema(output_schema)
|
||||
# Use lightweight compression - prune additionalProperties and unused definitions
|
||||
if output_schema.get("additionalProperties") is False:
|
||||
output_schema.pop("additionalProperties")
|
||||
|
||||
# Remove unused definitions (lightweight approach - just check direct $ref usage)
|
||||
if "$defs" in output_schema:
|
||||
used_refs = set()
|
||||
|
||||
def find_refs_in_value(value):
|
||||
if isinstance(value, dict):
|
||||
if "$ref" in value and isinstance(value["$ref"], str):
|
||||
ref = value["$ref"]
|
||||
if ref.startswith("#/$defs/"):
|
||||
used_refs.add(ref.split("/")[-1])
|
||||
for v in value.values():
|
||||
find_refs_in_value(v)
|
||||
elif isinstance(value, list):
|
||||
for item in value:
|
||||
find_refs_in_value(item)
|
||||
|
||||
# Find refs in the main schema (excluding $defs section)
|
||||
for key, value in output_schema.items():
|
||||
if key != "$defs":
|
||||
find_refs_in_value(value)
|
||||
|
||||
# Remove unused definitions
|
||||
if used_refs:
|
||||
output_schema["$defs"] = {
|
||||
name: def_schema
|
||||
for name, def_schema in output_schema["$defs"].items()
|
||||
if name in used_refs
|
||||
}
|
||||
else:
|
||||
output_schema.pop("$defs")
|
||||
|
||||
# Adjust union types to handle overlapping unions
|
||||
output_schema = cast(dict[str, Any], _adjust_union_types(output_schema))
|
||||
|
|
|
|||
|
|
@ -202,3 +202,139 @@ class TestFastMCPOpenAPIBasicFunctionality:
|
|||
# Should handle empty paths gracefully
|
||||
assert hasattr(server, "_director")
|
||||
assert hasattr(server, "_spec")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_clean_schema_output_no_unused_defs(self):
|
||||
"""Test that unused schema definitions are removed from tool schemas."""
|
||||
# Create a spec with unused HTTPValidationError-like definitions
|
||||
spec_with_unused_defs = {
|
||||
"openapi": "3.0.0",
|
||||
"info": {"title": "Test API", "version": "1.0.0"},
|
||||
"servers": [{"url": "https://api.example.com"}],
|
||||
"paths": {
|
||||
"/users": {
|
||||
"post": {
|
||||
"operationId": "create_user",
|
||||
"summary": "Create a new user",
|
||||
"requestBody": {
|
||||
"required": True,
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "title": "Name"},
|
||||
"active": {
|
||||
"type": "boolean",
|
||||
"title": "Active",
|
||||
},
|
||||
},
|
||||
"required": ["name", "active"],
|
||||
}
|
||||
}
|
||||
},
|
||||
},
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "User created successfully",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {
|
||||
"type": "integer",
|
||||
"title": "Id",
|
||||
},
|
||||
"name": {
|
||||
"type": "string",
|
||||
"title": "Name",
|
||||
},
|
||||
"active": {
|
||||
"type": "boolean",
|
||||
"title": "Active",
|
||||
},
|
||||
},
|
||||
"required": ["id", "name", "active"],
|
||||
"title": "User",
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
},
|
||||
"components": {
|
||||
"schemas": {
|
||||
# This should be removed since it's not referenced
|
||||
"HTTPValidationError": {
|
||||
"properties": {
|
||||
"detail": {
|
||||
"items": {
|
||||
"$ref": "#/components/schemas/ValidationError"
|
||||
},
|
||||
"title": "Detail",
|
||||
"type": "array",
|
||||
}
|
||||
},
|
||||
"title": "HTTPValidationError",
|
||||
"type": "object",
|
||||
},
|
||||
"ValidationError": {
|
||||
"properties": {
|
||||
"loc": {
|
||||
"items": {
|
||||
"anyOf": [{"type": "string"}, {"type": "integer"}]
|
||||
},
|
||||
"title": "Location",
|
||||
"type": "array",
|
||||
},
|
||||
"msg": {"title": "Message", "type": "string"},
|
||||
"type": {"title": "Error Type", "type": "string"},
|
||||
},
|
||||
"required": ["loc", "msg", "type"],
|
||||
"title": "ValidationError",
|
||||
"type": "object",
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
async with httpx.AsyncClient(base_url="https://api.example.com") as client:
|
||||
server = FastMCPOpenAPI(
|
||||
openapi_spec=spec_with_unused_defs, client=client, name="Test Server"
|
||||
)
|
||||
|
||||
async with Client(server) as mcp_client:
|
||||
tools = await mcp_client.list_tools()
|
||||
|
||||
assert len(tools) == 1 # Only the POST operation
|
||||
tool = tools[0]
|
||||
|
||||
# Verify tool has clean schemas without unused $defs
|
||||
assert tool.name == "create_user"
|
||||
|
||||
# Input schema should not have $defs since no references are used
|
||||
expected_input_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"name": {"type": "string", "title": "Name"},
|
||||
"active": {"type": "boolean", "title": "Active"},
|
||||
},
|
||||
"required": ["name", "active"],
|
||||
}
|
||||
assert tool.inputSchema == expected_input_schema
|
||||
|
||||
# Output schema should not have $defs since no references are used
|
||||
expected_output_schema = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "integer", "title": "Id"},
|
||||
"name": {"type": "string", "title": "Name"},
|
||||
"active": {"type": "boolean", "title": "Active"},
|
||||
},
|
||||
"required": ["id", "name", "active"],
|
||||
"title": "User",
|
||||
}
|
||||
assert tool.outputSchema == expected_output_schema
|
||||
|
|
|
|||
142
tests/experimental/server/test_openapi_performance.py
Normal file
142
tests/experimental/server/test_openapi_performance.py
Normal file
|
|
@ -0,0 +1,142 @@
|
|||
"""Performance regression tests for OpenAPI parsing.
|
||||
|
||||
These tests ensure that large OpenAPI schemas (like GitHub's API) parse quickly
|
||||
and don't regress to the slow performance we had before optimization.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.utilities.tests import temporary_settings
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def use_new_openapi_parser():
|
||||
with temporary_settings(experimental__enable_new_openapi_parser=True):
|
||||
yield
|
||||
|
||||
|
||||
class TestOpenAPIPerformance:
|
||||
"""Performance tests for OpenAPI parsing with real-world large schemas."""
|
||||
|
||||
# 20 second maximum timeout for this test no matter what
|
||||
@pytest.mark.timeout(20)
|
||||
async def test_github_api_schema_performance(self):
|
||||
"""
|
||||
Test that GitHub's full API schema parses quickly.
|
||||
|
||||
This is a regression test to ensure our performance optimizations
|
||||
(eliminating deepcopy, single-pass optimization, smart union adjustment)
|
||||
continue to work. Without these optimizations, this test would take
|
||||
multiple minutes to parse.
|
||||
|
||||
On a local machine, this tests passes in ~2 seconds, but in GHA CI we see
|
||||
times as high as 6-7 seconds, so the test is asserted to pass in under
|
||||
10. Given that, this isn't intended to be a strict performance test, but
|
||||
rather a canary to ensure we don't regress significantly.
|
||||
"""
|
||||
|
||||
# Download the full GitHub API schema (typically ~10MB)
|
||||
response = httpx.get(
|
||||
"https://raw.githubusercontent.com/github/rest-api-description/refs/heads/main/descriptions-next/ghes-3.17/ghes-3.17.json",
|
||||
timeout=30.0, # Allow time for download
|
||||
)
|
||||
response.raise_for_status()
|
||||
schema = response.json()
|
||||
|
||||
# Time the parsing operation
|
||||
start_time = time.time()
|
||||
|
||||
# This should complete quickly with our optimizations
|
||||
mcp_server = FastMCP.from_openapi(schema, httpx.AsyncClient())
|
||||
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
print(f"OpenAPI parsing took {elapsed_time:.2f}s")
|
||||
|
||||
# Verify the server was created successfully
|
||||
assert mcp_server is not None
|
||||
|
||||
# Performance regression test: should complete in under 10 seconds
|
||||
assert elapsed_time < 10.0, (
|
||||
f"OpenAPI parsing took {elapsed_time:.2f}s, exceeding 10s limit. "
|
||||
f"This suggests a performance regression."
|
||||
)
|
||||
|
||||
# Verify server and tools were created successfully
|
||||
tools = await mcp_server.get_tools()
|
||||
assert len(tools) > 500
|
||||
|
||||
def test_medium_schema_performance(self):
|
||||
"""
|
||||
Test parsing performance with a smaller synthetic schema.
|
||||
|
||||
This test doesn't require network access and provides a baseline
|
||||
for performance testing in CI environments.
|
||||
"""
|
||||
# Create a medium-sized synthetic schema
|
||||
schema = {
|
||||
"openapi": "3.0.0",
|
||||
"info": {"title": "Test API", "version": "1.0.0"},
|
||||
"paths": {},
|
||||
}
|
||||
|
||||
# Generate multiple paths to create a reasonably sized schema
|
||||
for i in range(100):
|
||||
path = f"/test/{i}"
|
||||
schema["paths"][path] = {
|
||||
"get": {
|
||||
"operationId": f"test_{i}",
|
||||
"parameters": [
|
||||
{"name": "param1", "in": "query", "schema": {"type": "string"}}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"id": {"type": "integer"},
|
||||
"name": {"type": "string"},
|
||||
"data": {
|
||||
"type": "object",
|
||||
"additionalProperties": False,
|
||||
"properties": {
|
||||
"value": {"type": "string"},
|
||||
"metadata": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"created": {
|
||||
"type": "string"
|
||||
},
|
||||
"updated": {
|
||||
"type": "string"
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
# Time the parsing
|
||||
start_time = time.time()
|
||||
mcp_server = FastMCP.from_openapi(schema, httpx.AsyncClient())
|
||||
elapsed_time = time.time() - start_time
|
||||
|
||||
# Should be very fast for medium schemas (well under 1 second)
|
||||
assert elapsed_time < 1.0, (
|
||||
f"Medium schema parsing took {elapsed_time:.3f}s, expected <1s"
|
||||
)
|
||||
assert mcp_server is not None
|
||||
|
|
@ -1,7 +1,5 @@
|
|||
from fastmcp.utilities.json_schema import (
|
||||
_prune_param,
|
||||
_prune_unused_defs,
|
||||
_walk_and_prune,
|
||||
compress_schema,
|
||||
)
|
||||
|
||||
|
|
@ -9,8 +7,10 @@ from fastmcp.utilities.json_schema import (
|
|||
|
||||
|
||||
def _prune_additional_properties(schema):
|
||||
"""Wrapper for _walk_and_prune that only prunes additionalProperties: false."""
|
||||
return _walk_and_prune(schema, prune_additional_properties=True)
|
||||
"""Wrapper for compress_schema that only prunes additionalProperties: false."""
|
||||
return compress_schema(
|
||||
schema, prune_defs=False, prune_additional_properties=True, prune_titles=False
|
||||
)
|
||||
|
||||
|
||||
class TestPruneParam:
|
||||
|
|
@ -55,7 +55,7 @@ class TestPruneParam:
|
|||
|
||||
|
||||
class TestPruneUnusedDefs:
|
||||
"""Tests for the _prune_unused_defs function."""
|
||||
"""Tests for unused definition pruning (via compress_schema)."""
|
||||
|
||||
def test_removes_unreferenced_defs(self):
|
||||
"""Test that unreferenced definitions are removed."""
|
||||
|
|
@ -68,7 +68,12 @@ class TestPruneUnusedDefs:
|
|||
"unused_def": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "foo_def" in result["$defs"]
|
||||
assert "unused_def" not in result["$defs"]
|
||||
|
||||
|
|
@ -87,7 +92,12 @@ class TestPruneUnusedDefs:
|
|||
"unused_def": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "foo_def" in result["$defs"]
|
||||
assert "nested_def" in result["$defs"]
|
||||
assert "unused_def" not in result["$defs"]
|
||||
|
|
@ -104,7 +114,12 @@ class TestPruneUnusedDefs:
|
|||
"nested_def": {"type": "string"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "$defs" not in result
|
||||
|
||||
def test_nested_references_with_recursion_kept(self):
|
||||
|
|
@ -121,7 +136,12 @@ class TestPruneUnusedDefs:
|
|||
"unused_def": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "foo_def" in result["$defs"]
|
||||
assert "unused_def" not in result["$defs"]
|
||||
|
||||
|
|
@ -136,7 +156,12 @@ class TestPruneUnusedDefs:
|
|||
},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "$defs" not in result
|
||||
|
||||
def test_multiple_nested_references_with_recursion_kept(self):
|
||||
|
|
@ -157,7 +182,12 @@ class TestPruneUnusedDefs:
|
|||
"unused_def": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "foo_def" in result["$defs"]
|
||||
assert "nested_def" in result["$defs"]
|
||||
assert "unused_def" not in result["$defs"]
|
||||
|
|
@ -177,7 +207,12 @@ class TestPruneUnusedDefs:
|
|||
},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "$defs" not in result
|
||||
|
||||
def test_array_references_kept(self):
|
||||
|
|
@ -191,7 +226,12 @@ class TestPruneUnusedDefs:
|
|||
"unused_def": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "item_def" in result["$defs"]
|
||||
assert "unused_def" not in result["$defs"]
|
||||
|
||||
|
|
@ -205,7 +245,12 @@ class TestPruneUnusedDefs:
|
|||
"unused_def": {"type": "integer"},
|
||||
},
|
||||
}
|
||||
result = _prune_unused_defs(schema)
|
||||
result = compress_schema(
|
||||
schema,
|
||||
prune_defs=True,
|
||||
prune_additional_properties=False,
|
||||
prune_titles=False,
|
||||
)
|
||||
assert "$defs" not in result
|
||||
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue