Ensure openapi defs are loaded

This commit is contained in:
Jeremiah Lowin 2025-05-14 14:20:00 -04:00
commit df92d288b6
2 changed files with 149 additions and 11 deletions

View file

@ -14,6 +14,7 @@ def _prune_param(schema: dict, param: str) -> dict:
removed = props.pop(param, None)
if removed is None: # nothing to do
return schema
# Keep empty properties object rather than removing it entirely
schema["properties"] = props
if param in schema.get("required", []):
@ -21,7 +22,12 @@ def _prune_param(schema: dict, param: str) -> dict:
if not schema["required"]:
schema.pop("required")
# ── 2. collect all remaining local $ref targets ───────────────────
return schema
def _prune_unused_defs(schema: dict) -> dict:
"""Remove unused definitions from the schema."""
# collect all remaining local $ref targets
used_defs: set[str] = set()
def walk(node: object) -> None: # depth-first traversal
@ -37,7 +43,8 @@ def _prune_param(schema: dict, param: str) -> dict:
walk(schema)
# ── 3. remove orphaned definitions ────────────────────────────────
# remove orphaned definitions
defs = schema.get("$defs", {})
for def_name in list(defs):
if def_name not in used_defs:
@ -48,12 +55,28 @@ def _prune_param(schema: dict, param: str) -> dict:
return schema
def prune_params(schema: dict, params: list[str]) -> dict:
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
def compress_schema(
schema: dict,
prune_params: list[str] | None = None,
prune_defs: bool = True,
prune_additional_properties: bool = True,
) -> dict:
"""
Remove the given parameters from the schema.
"""
schema = copy.deepcopy(schema)
for param in params:
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)
return schema

View file

@ -84,6 +84,9 @@ class HTTPRoute(BaseModel):
responses: dict[str, ResponseInfo] = Field(
default_factory=dict
) # Key: status code str
schema_definitions: dict[str, JsonSchema] = Field(
default_factory=dict
) # Store component schemas
# Export public symbols
@ -221,6 +224,27 @@ class OpenAPI31Parser(BaseOpenAPIParser):
logger.warning("OpenAPI schema has no paths defined.")
return []
# Extract component schemas to add to each route
schema_definitions = {}
if hasattr(self.openapi, "components") and self.openapi.components:
components = self.openapi.components
if hasattr(components, "schemas") and components.schemas:
for name, schema in components.schemas.items():
try:
if isinstance(schema, Reference):
resolved_schema = self._resolve_ref(schema)
schema_definitions[name] = self._extract_schema_as_dict(
resolved_schema
)
else:
schema_definitions[name] = self._extract_schema_as_dict(
schema
)
except Exception as e:
logger.warning(
f"Failed to extract schema definition '{name}': {e}"
)
for path_str, path_item_obj in self.openapi.paths.items():
if not isinstance(path_item_obj, PathItem):
logger.warning(
@ -269,6 +293,7 @@ class OpenAPI31Parser(BaseOpenAPIParser):
parameters=parameters,
request_body=request_body_info,
responses=responses,
schema_definitions=schema_definitions,
)
routes.append(route)
logger.info(
@ -386,16 +411,36 @@ class OpenAPI31Parser(BaseOpenAPIParser):
param_schema_dict = {}
if param_schema_obj: # Check if schema exists
# Resolve the schema if it's a reference
resolved_schema = self._resolve_ref(param_schema_obj)
param_schema_dict = self._extract_schema_as_dict(param_schema_obj)
# Ensure default value is preserved from resolved schema
if (
not isinstance(resolved_schema, Reference)
and hasattr(resolved_schema, "default")
and resolved_schema.default is not None
):
param_schema_dict["default"] = resolved_schema.default
elif parameter.content:
# Handle complex parameters with 'content'
first_media_type = next(iter(parameter.content.values()), None)
if (
first_media_type and first_media_type.media_type_schema
): # CORRECTED: Use 'media_type_schema'
param_schema_dict = self._extract_schema_as_dict(
first_media_type.media_type_schema
)
# Resolve the schema if it's a reference
media_schema = first_media_type.media_type_schema
resolved_media_schema = self._resolve_ref(media_schema)
param_schema_dict = self._extract_schema_as_dict(media_schema)
# Ensure default value is preserved from resolved schema
if (
not isinstance(resolved_media_schema, Reference)
and hasattr(resolved_media_schema, "default")
and resolved_media_schema.default is not None
):
param_schema_dict["default"] = resolved_media_schema.default
logger.debug(
f"Parameter '{parameter.name}' using schema from 'content' field."
)
@ -543,6 +588,27 @@ class OpenAPI30Parser(BaseOpenAPIParser):
logger.warning("OpenAPI schema has no paths defined.")
return []
# Extract component schemas to add to each route
schema_definitions = {}
if hasattr(self.openapi, "components") and self.openapi.components:
components = self.openapi.components
if hasattr(components, "schemas") and components.schemas:
for name, schema in components.schemas.items():
try:
if isinstance(schema, Reference_30):
resolved_schema = self._resolve_ref(schema)
schema_definitions[name] = self._extract_schema_as_dict(
resolved_schema
)
else:
schema_definitions[name] = self._extract_schema_as_dict(
schema
)
except Exception as e:
logger.warning(
f"Failed to extract schema definition '{name}': {e}"
)
for path_str, path_item_obj in self.openapi.paths.items():
if not isinstance(path_item_obj, PathItem_30):
logger.warning(
@ -593,6 +659,7 @@ class OpenAPI30Parser(BaseOpenAPIParser):
parameters=parameters,
request_body=request_body_info,
responses=responses,
schema_definitions=schema_definitions,
)
routes.append(route)
logger.info(
@ -711,14 +778,34 @@ class OpenAPI30Parser(BaseOpenAPIParser):
param_schema_dict = {}
if param_schema_obj: # Check if schema exists
# Resolve the schema if it's a reference
resolved_schema = self._resolve_ref(param_schema_obj)
param_schema_dict = self._extract_schema_as_dict(param_schema_obj)
# Ensure default value is preserved from resolved schema
if (
not isinstance(resolved_schema, Reference_30)
and hasattr(resolved_schema, "default")
and resolved_schema.default is not None
):
param_schema_dict["default"] = resolved_schema.default
elif parameter.content:
# Handle complex parameters with 'content'
first_media_type = next(iter(parameter.content.values()), None)
if first_media_type and first_media_type.media_type_schema:
param_schema_dict = self._extract_schema_as_dict(
first_media_type.media_type_schema
)
# Resolve the schema if it's a reference
media_schema = first_media_type.media_type_schema
resolved_media_schema = self._resolve_ref(media_schema)
param_schema_dict = self._extract_schema_as_dict(media_schema)
# Ensure default value is preserved from resolved schema
if (
not isinstance(resolved_media_schema, Reference_30)
and hasattr(resolved_media_schema, "default")
and resolved_media_schema.default is not None
):
param_schema_dict["default"] = resolved_media_schema.default
logger.debug(
f"Parameter '{parameter.name}' using schema from 'content' field."
)
@ -1173,6 +1260,23 @@ def _combine_schemas(route: openapi.HTTPRoute) -> dict[str, Any]:
# Copy the schema and add description if available
param_schema = param.schema_.copy() if isinstance(param.schema_, dict) else {}
# Convert #/components/schemas references to #/$defs references
if isinstance(param_schema, dict) and "$ref" in param_schema:
ref_path = param_schema["$ref"]
if ref_path.startswith("#/components/schemas/"):
schema_name = ref_path.split("/")[-1]
param_schema["$ref"] = f"#/$defs/{schema_name}"
# Also handle anyOf, allOf, oneOf references
for section in ["anyOf", "allOf", "oneOf"]:
if section in param_schema and isinstance(param_schema[section], list):
for i, item in enumerate(param_schema[section]):
if isinstance(item, dict) and "$ref" in item:
ref_path = item["$ref"]
if ref_path.startswith("#/components/schemas/"):
schema_name = ref_path.split("/")[-1]
param_schema[section][i]["$ref"] = f"#/$defs/{schema_name}"
# Add parameter description to schema if available and not already present
if param.description and not param_schema.get("description"):
param_schema["description"] = param.description
@ -1193,8 +1297,19 @@ def _combine_schemas(route: openapi.HTTPRoute) -> dict[str, Any]:
if route.request_body.required:
required.extend(body_schema.get("required", []))
return {
result = {
"type": "object",
"properties": properties,
"required": required,
}
# Add schema definitions if available
if route.schema_definitions:
result["$defs"] = route.schema_definitions
# Use compress_schema to remove unused definitions
from fastmcp.utilities.json_schema import compress_schema
result = compress_schema(result)
return result