diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py index d57dde0459..c448742f6d 100644 --- a/litellm/proxy/common_utils/custom_openapi_spec.py +++ b/litellm/proxy/common_utils/custom_openapi_spec.py @@ -93,6 +93,11 @@ class CustomOpenAPISpec: schema_properties = actual_schema.get("properties", {}) required_fields = actual_schema.get("required", []) + # Extract $defs and add them to components/schemas + # This fixes Pydantic v2 $defs not being resolvable in Swagger/OpenAPI + if "$defs" in actual_schema: + CustomOpenAPISpec._move_defs_to_components(openapi_schema, actual_schema["$defs"]) + # Create an expanded inline schema instead of just a $ref # This makes Swagger UI show all individual fields in the request body editor expanded_schema = { @@ -105,6 +110,9 @@ class CustomOpenAPISpec: for field_name, field_def in schema_properties.items(): expanded_field = CustomOpenAPISpec._expand_field_definition(field_def) + # Rewrite $defs references to use components/schemas instead + expanded_field = CustomOpenAPISpec._rewrite_defs_refs(expanded_field) + # Add a simple example for the messages field if field_name == "messages": expanded_field["example"] = [ @@ -113,11 +121,6 @@ class CustomOpenAPISpec: expanded_schema["properties"][field_name] = expanded_field - # Include $defs from the original schema to support complex types like AllMessageValues - # This ensures that message types and other complex union types work properly - if "$defs" in actual_schema: - expanded_schema["$defs"] = actual_schema["$defs"] - # Set the request body with the expanded schema openapi_schema["paths"][path]["post"]["requestBody"] = { "required": True, @@ -138,6 +141,66 @@ class CustomOpenAPISpec: ] openapi_schema["paths"][path]["post"]["parameters"] = filtered_params + @staticmethod + def _move_defs_to_components(openapi_schema: Dict[str, Any], defs: Dict[str, Any]) -> None: + """ + Move $defs from Pydantic v2 schema to OpenAPI components/schemas. + This makes the definitions resolvable in Swagger/OpenAPI viewers. + + Args: + openapi_schema: The OpenAPI schema dict to modify + defs: The $defs dictionary from Pydantic schema + """ + if not defs: + return + + # Ensure components/schemas exists + if "components" not in openapi_schema: + openapi_schema["components"] = {} + if "schemas" not in openapi_schema["components"]: + openapi_schema["components"]["schemas"] = {} + + # Add each definition to components/schemas + for def_name, def_schema in defs.items(): + # Recursively rewrite any nested $defs references within this definition + rewritten_def = CustomOpenAPISpec._rewrite_defs_refs(def_schema) + openapi_schema["components"]["schemas"][def_name] = rewritten_def + + # If this definition also has $defs, process them recursively + if "$defs" in def_schema: + CustomOpenAPISpec._move_defs_to_components(openapi_schema, def_schema["$defs"]) + + @staticmethod + def _rewrite_defs_refs(schema: Any) -> Any: + """ + Recursively rewrite $ref values from #/$defs/... to #/components/schemas/... + This converts Pydantic v2 references to OpenAPI-compatible references. + + Args: + schema: Schema object to process (can be dict, list, or primitive) + + Returns: + Schema with rewritten references + """ + if isinstance(schema, dict): + result = {} + for key, value in schema.items(): + if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"): + # Rewrite the reference to use components/schemas + def_name = value.replace("#/$defs/", "") + result[key] = f"#/components/schemas/{def_name}" + elif key == "$defs": + # Remove $defs from the schema since they're moved to components + continue + else: + # Recursively process nested structures + result[key] = CustomOpenAPISpec._rewrite_defs_refs(value) + return result + elif isinstance(schema, list): + return [CustomOpenAPISpec._rewrite_defs_refs(item) for item in schema] + else: + return schema + @staticmethod def _extract_field_schema(field_def: Dict[str, Any]) -> Dict[str, Any]: """ diff --git a/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py b/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py index 8dd99e2ae5..7549b4259a 100644 --- a/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py +++ b/tests/test_litellm/proxy/common_utils/test_custom_openapi_spec.py @@ -88,4 +88,66 @@ class TestCustomOpenAPISpec: paths=CustomOpenAPISpec.RESPONSES_API_PATHS, operation_name="responses API" ) - assert result == base_openapi_schema \ No newline at end of file + assert result == base_openapi_schema + + +def test_move_defs_to_components(): + """ + Test that $defs from Pydantic v2 schemas are moved to components/schemas. + """ + openapi_schema = {} + + defs = { + "UserMessage": { + "type": "object", + "properties": { + "role": {"type": "string"}, + "content": {"type": "string"} + } + }, + "AssistantMessage": { + "type": "object", + "properties": { + "role": {"type": "string"}, + "content": {"type": "string"} + } + } + } + + CustomOpenAPISpec._move_defs_to_components(openapi_schema=openapi_schema, defs=defs) + + assert "components" in openapi_schema + assert "schemas" in openapi_schema["components"] + assert "UserMessage" in openapi_schema["components"]["schemas"] + assert "AssistantMessage" in openapi_schema["components"]["schemas"] + assert openapi_schema["components"]["schemas"]["UserMessage"]["type"] == "object" + + +def test_rewrite_defs_refs(): + """ + Test that $ref values are rewritten from #/$defs/ to #/components/schemas/. + """ + schema = { + "type": "object", + "properties": { + "messages": { + "type": "array", + "items": { + "anyOf": [ + {"$ref": "#/$defs/UserMessage"}, + {"$ref": "#/$defs/AssistantMessage"} + ] + } + } + }, + "$defs": { + "UserMessage": {"type": "object"}, + "AssistantMessage": {"type": "object"} + } + } + + rewritten = CustomOpenAPISpec._rewrite_defs_refs(schema=schema) + + assert "$defs" not in rewritten + assert rewritten["properties"]["messages"]["items"]["anyOf"][0]["$ref"] == "#/components/schemas/UserMessage" + assert rewritten["properties"]["messages"]["items"]["anyOf"][1]["$ref"] == "#/components/schemas/AssistantMessage"