From 099ccf56a746c9bca8afa4c0c62fdb52fa5abfa9 Mon Sep 17 00:00:00 2001
From: Richard Song <9144514+richardmcsong@users.noreply.github.com>
Date: Wed, 3 Dec 2025 00:57:07 -0500
Subject: [PATCH] Refactor add_schema_to_components to move definitions to
 components/schemas and add corresponding unit test (#17389)

---
 .../proxy/common_utils/custom_openapi_spec.py |  2 +-
 .../common_utils/test_custom_openapi_spec.py  | 29 +++++++++++++++++++
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/litellm/proxy/common_utils/custom_openapi_spec.py b/litellm/proxy/common_utils/custom_openapi_spec.py
index c448742f6d..69472c2cda 100644
--- a/litellm/proxy/common_utils/custom_openapi_spec.py
+++ b/litellm/proxy/common_utils/custom_openapi_spec.py
@@ -72,7 +72,7 @@ class CustomOpenAPISpec:
             openapi_schema["components"]["schemas"] = {}
         
         # Add the schema
-        openapi_schema["components"]["schemas"][schema_name] = schema_def
+        CustomOpenAPISpec._move_defs_to_components(openapi_schema, {schema_name: schema_def})
     
     @staticmethod
     def add_request_body_to_paths(openapi_schema: Dict[str, Any], paths: List[str], schema_ref: str) -> None:
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 7549b4259a..5fef35eb82 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
@@ -90,6 +90,35 @@ class TestCustomOpenAPISpec:
             )
             assert result == base_openapi_schema 
 
+def test_defs_rewritten_in_add_schema_to_components():
+    """
+    Test that defs are rewritten to components/schemas in add_schema_to_components.
+    """
+
+    openapi_schema = {}
+    schema_name = "SchemaName"
+    schema_def = {
+        "type": "object",
+        "properties": {
+            "messages": {
+                "type": "array",
+                "items": {
+                    "anyOf": [
+                        {"$ref": "#/$defs/UserMessage"},
+                        {"$ref": "#/$defs/AssistantMessage"}
+                    ]
+                }
+            }
+        },
+        "$defs": {
+            "UserMessage": {"type": "object"},
+            "AssistantMessage": {"type": "object"}
+        }
+    }
+    CustomOpenAPISpec.add_schema_to_components(openapi_schema=openapi_schema, schema_name=schema_name, schema_def=schema_def)
+    assert "$defs" not in openapi_schema
+    assert openapi_schema["components"]["schemas"]["SchemaName"]["properties"]["messages"]["items"]["anyOf"][0]["$ref"] == "#/components/schemas/UserMessage"
+    assert openapi_schema["components"]["schemas"]["SchemaName"]["properties"]["messages"]["items"]["anyOf"][1]["$ref"] == "#/components/schemas/AssistantMessage"
 
 def test_move_defs_to_components():
     """