Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/common_utils/custom_openapi_spec.py: 79%
155 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1from collections.abc import Mapping, Sequence
2from types import MappingProxyType
3from typing import Final, TypeAlias, Union, cast
5from pydantic import TypeAdapter
7from litellm._logging import verbose_proxy_logger
9JsonValue: TypeAlias = Union["JsonObject", "JsonArray", str, int, float, bool, None]
10JsonObject: TypeAlias = dict[str, JsonValue]
11JsonArray: TypeAlias = list[JsonValue]
14class CustomOpenAPISpec:
15 """
16 Handler for customizing OpenAPI specifications with Pydantic models
17 for documentation purposes without runtime validation.
18 """
20 CHAT_COMPLETION_PATHS = [
21 "/v1/chat/completions",
22 "/chat/completions",
23 "/engines/{model}/chat/completions",
24 "/openai/deployments/{model}/chat/completions",
25 ]
27 EMBEDDING_PATHS = [
28 "/v1/embeddings",
29 "/embeddings",
30 "/engines/{model}/embeddings",
31 "/openai/deployments/{model}/embeddings",
32 ]
34 RESPONSES_API_PATHS = ["/v1/responses", "/responses", "/openai/v1/responses"]
36 @staticmethod
37 def _as_object(node: JsonValue) -> JsonObject:
38 return node if isinstance(node, dict) else {}
40 @staticmethod
41 def _as_array(node: JsonValue) -> JsonArray:
42 return node if isinstance(node, list) else []
44 @staticmethod
45 def _components_schemas(openapi_schema: JsonObject) -> JsonObject:
46 components: Final = CustomOpenAPISpec._as_object(openapi_schema.setdefault("components", {}))
47 return CustomOpenAPISpec._as_object(components.setdefault("schemas", {}))
49 @staticmethod
50 def get_pydantic_schema(model_class: type) -> JsonObject | None:
51 """
52 Get JSON schema for a request or response model class, including TypedDicts.
54 Args:
55 model_class: Pydantic model class or TypedDict
57 Returns:
58 JSON schema dict or None if failed
59 """
60 try:
61 return cast(JsonObject, TypeAdapter(model_class).json_schema()) # cast-ok: pydantic returns dict[str, Any]
62 except Exception as e:
63 # FastAPI 0.120+ may fail schema generation for certain types (e.g., openai.Timeout)
64 # Log the error and return None to skip schema generation for this model
65 verbose_proxy_logger.debug("Failed to generate schema for %s: %s", model_class, e)
66 return None
68 @staticmethod
69 def add_schema_to_components(openapi_schema: JsonObject, schema_name: str, schema_def: JsonObject) -> None:
70 """
71 Add a schema definition to the OpenAPI components/schemas section.
73 Args:
74 openapi_schema: The OpenAPI schema dict to modify
75 schema_name: Name for the schema component
76 schema_def: The schema definition
77 """
78 # Ensure components/schemas structure exists
79 _ = CustomOpenAPISpec._components_schemas(openapi_schema)
81 defs: Final[Mapping[str, JsonValue]] = (
82 CustomOpenAPISpec._as_object(schema_def["$defs"]) if "$defs" in schema_def else MappingProxyType({})
83 )
84 renames: Final = CustomOpenAPISpec._move_defs_to_components(openapi_schema, defs, schema_name)
85 schemas: Final = CustomOpenAPISpec._components_schemas(openapi_schema)
86 schemas[schema_name] = CustomOpenAPISpec._rewrite_defs_refs(schema_def, renames)
88 @staticmethod
89 def _expanded_request_field(field_name: str, field_def: JsonValue) -> JsonValue:
90 expanded: Final = CustomOpenAPISpec._rewrite_defs_refs(
91 CustomOpenAPISpec._expand_field_definition(CustomOpenAPISpec._as_object(field_def)),
92 MappingProxyType({}),
93 )
94 if field_name != "messages":
95 return expanded
96 return {
97 **CustomOpenAPISpec._as_object(expanded),
98 "example": [{"role": "user", "content": "Hello, how are you?"}],
99 }
101 @staticmethod
102 def add_request_body_to_paths(openapi_schema: JsonObject, paths: Sequence[str], schema_ref: str) -> None:
103 """
104 Add request body with expanded form fields for better Swagger UI display.
105 This keeps the request body but expands it to show individual fields in the UI.
107 Args:
108 openapi_schema: The OpenAPI schema dict to modify
109 paths: List of paths to update
110 schema_ref: Reference to the schema component (e.g., "#/components/schemas/ModelName")
111 """
112 for path in paths:
113 path_item = CustomOpenAPISpec._as_object(
114 CustomOpenAPISpec._as_object(openapi_schema.get("paths")).get(path)
115 )
116 if "post" not in path_item: 116 ↛ 117line 116 didn't jump to line 117 because the condition on line 116 was never true
117 continue
119 post_operation = CustomOpenAPISpec._as_object(path_item["post"])
121 # Get the actual schema to extract ALL field definitions
122 schema_name = schema_ref.split("/")[-1] # Extract "ProxyChatCompletionRequest" from the ref
123 components = CustomOpenAPISpec._as_object(openapi_schema.get("components"))
124 actual_schema = CustomOpenAPISpec._as_object(
125 CustomOpenAPISpec._as_object(components.get("schemas")).get(schema_name)
126 )
127 schema_properties = CustomOpenAPISpec._as_object(actual_schema.get("properties"))
128 required_fields = actual_schema.get("required", [])
130 # Create an expanded inline schema instead of just a $ref
131 # This makes Swagger UI show all individual fields in the request body editor
132 expanded_schema: JsonObject = {
133 "type": "object",
134 "required": required_fields,
135 "properties": {
136 field_name: CustomOpenAPISpec._expanded_request_field(field_name, field_def)
137 for field_name, field_def in schema_properties.items()
138 },
139 }
141 # Set the request body with the expanded schema
142 post_operation["requestBody"] = {
143 "required": True,
144 "content": {"application/json": {"schema": expanded_schema}},
145 }
147 # Keep any existing parameters (like path parameters) but remove conflicting query params
148 if "parameters" in post_operation:
149 # Only keep path parameters, remove query params that conflict with request body
150 post_operation["parameters"] = [
151 param
152 for param in CustomOpenAPISpec._as_array(post_operation["parameters"])
153 if CustomOpenAPISpec._as_object(param).get("in") == "path"
154 ]
156 @staticmethod
157 def _move_defs_to_components(
158 openapi_schema: JsonObject, defs: Mapping[str, JsonValue], namespace: str
159 ) -> Mapping[str, str]:
160 """
161 Move $defs from Pydantic v2 schema to OpenAPI components/schemas.
162 This makes the definitions resolvable in Swagger/OpenAPI viewers.
164 Args:
165 openapi_schema: The OpenAPI schema dict to modify
166 defs: The $defs dictionary from Pydantic schema
167 namespace: Prefix used to rename defs that would overwrite an existing component
169 Returns:
170 Map of original def names to renamed component names for collision cases
171 """
172 schemas: Final = CustomOpenAPISpec._components_schemas(openapi_schema)
173 renames: Final = CustomOpenAPISpec._fixed_renames(schemas, defs, namespace, MappingProxyType({}))
174 for def_name, def_schema in defs.items():
175 if def_name in schemas and def_name not in renames:
176 continue
177 schemas[renames.get(def_name, def_name)] = CustomOpenAPISpec._rewrite_defs_refs(def_schema, renames)
178 return renames
180 @staticmethod
181 def _def_collisions(
182 schemas: JsonObject, defs: Mapping[str, JsonValue], namespace: str, renames: Mapping[str, str]
183 ) -> Mapping[str, str]:
184 return MappingProxyType(
185 {
186 name: f"{namespace}_{name}"
187 for name, d in defs.items()
188 if name in schemas
189 and not CustomOpenAPISpec._same_shape(schemas[name], CustomOpenAPISpec._rewrite_defs_refs(d, renames))
190 }
191 )
193 @staticmethod
194 def _same_shape(existing: JsonValue, incoming: JsonValue) -> bool:
195 if existing == incoming: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 return True
197 existing_obj: Final = CustomOpenAPISpec._as_object(existing)
198 incoming_obj: Final = CustomOpenAPISpec._as_object(incoming)
199 existing_props: Final = CustomOpenAPISpec._as_object(existing_obj.get("properties"))
200 incoming_props: Final = CustomOpenAPISpec._as_object(incoming_obj.get("properties"))
201 if not existing_props or not incoming_props: 201 ↛ 202line 201 didn't jump to line 202 because the condition on line 201 was never true
202 return False
203 return existing_props.keys() == incoming_props.keys() and frozenset(
204 x for x in CustomOpenAPISpec._as_array(existing_obj.get("required")) if isinstance(x, str)
205 ) == frozenset(x for x in CustomOpenAPISpec._as_array(incoming_obj.get("required")) if isinstance(x, str))
207 @staticmethod
208 def _fixed_renames(
209 schemas: JsonObject, defs: Mapping[str, JsonValue], namespace: str, renames: Mapping[str, str]
210 ) -> Mapping[str, str]:
211 next_renames: Final = MappingProxyType(
212 {**renames, **CustomOpenAPISpec._def_collisions(schemas, defs, namespace, renames)}
213 )
214 if next_renames == renames:
215 return renames
216 return CustomOpenAPISpec._fixed_renames(schemas, defs, namespace, next_renames)
218 @staticmethod
219 def _rewritten_defs_entry(key: str, value: JsonValue, renames: Mapping[str, str]) -> JsonValue:
220 if key == "$ref" and isinstance(value, str) and value.startswith("#/$defs/"):
221 # Rewrite the reference to use components/schemas
222 def_name: Final = value.replace("#/$defs/", "")
223 return f"#/components/schemas/{renames.get(def_name, def_name)}"
224 # Recursively process nested structures
225 return CustomOpenAPISpec._rewrite_defs_refs(value, renames)
227 @staticmethod
228 def _rewrite_defs_refs(schema: JsonValue, renames: Mapping[str, str]) -> JsonValue:
229 """
230 Recursively rewrite $ref values from #/$defs/... to #/components/schemas/...
231 This converts Pydantic v2 references to OpenAPI-compatible references.
233 Args:
234 schema: Schema object to process (can be dict, list, or primitive)
236 Returns:
237 Schema with rewritten references
238 """
239 if isinstance(schema, dict):
240 return {
241 key: CustomOpenAPISpec._rewritten_defs_entry(key, value, renames)
242 for key, value in schema.items()
243 if key != "$defs"
244 }
245 if isinstance(schema, list):
246 return [CustomOpenAPISpec._rewrite_defs_refs(item, renames) for item in schema]
247 return schema
249 @staticmethod
250 def _extract_field_schema(field_def: JsonObject) -> JsonValue:
251 """
252 Extract a simple schema from a Pydantic field definition for parameter display.
254 Args:
255 field_def: Pydantic field definition
257 Returns:
258 Simplified schema for OpenAPI parameter
259 """
260 # Handle simple types
261 if "type" in field_def:
262 return {"type": field_def["type"]}
264 # Handle anyOf (Optional fields in Pydantic v2)
265 if "anyOf" in field_def:
266 any_of: Final = CustomOpenAPISpec._as_array(field_def["anyOf"])
267 # Find the non-null type
268 for option in any_of:
269 if CustomOpenAPISpec._as_object(option).get("type") != "null":
270 return option
271 # Fallback to string if all else fails
272 return {"type": "string"}
274 # Default fallback
275 return {"type": "string"}
277 @staticmethod
278 def _expand_field_definition(field_def: JsonObject) -> JsonObject:
279 """
280 Expand a Pydantic field definition for inline use in OpenAPI schema.
281 This creates a full field definition that Swagger UI can render as individual form fields.
283 Args:
284 field_def: Pydantic field definition
286 Returns:
287 Expanded field definition for OpenAPI schema
288 """
289 # Return the field definition as-is since Pydantic already provides proper schemas
290 return field_def.copy()
292 @staticmethod
293 def add_request_schema(
294 openapi_schema: JsonObject,
295 model_class: type,
296 schema_name: str,
297 paths: Sequence[str],
298 operation_name: str,
299 ) -> JsonObject:
300 """
301 Generic method to add a request schema to OpenAPI specification.
303 Args:
304 openapi_schema: The OpenAPI schema dict to modify
305 model_class: The Pydantic model class to get schema from
306 schema_name: Name for the schema component
307 paths: List of paths to add the request body to
308 operation_name: Name of the operation for logging (e.g., "chat completion", "embedding")
310 Returns:
311 Modified OpenAPI schema
312 """
313 try:
314 # Get the schema for the model class
315 request_schema: Final = CustomOpenAPISpec.get_pydantic_schema(model_class)
317 # Only proceed if we successfully got the schema
318 if request_schema is not None: 318 ↛ 329line 318 didn't jump to line 329 because the condition on line 318 was always true
319 # Add schema to components
320 CustomOpenAPISpec.add_schema_to_components(openapi_schema, schema_name, request_schema)
322 # Add request body to specified endpoints
323 CustomOpenAPISpec.add_request_body_to_paths(
324 openapi_schema, paths, f"#/components/schemas/{schema_name}"
325 )
327 verbose_proxy_logger.debug("Successfully added %s schema to OpenAPI spec", schema_name)
328 else:
329 verbose_proxy_logger.debug("Could not get schema for %s", schema_name)
331 except Exception as e:
332 # If schema addition fails, continue without it
333 verbose_proxy_logger.debug("Failed to add %s request schema: %s", operation_name, e)
335 return openapi_schema
337 @staticmethod
338 def add_chat_completion_request_schema(
339 openapi_schema: JsonObject,
340 ) -> JsonObject:
341 """
342 Add ProxyChatCompletionRequest schema to chat completion endpoints for documentation.
343 This shows the request body in Swagger without runtime validation.
345 Args:
346 openapi_schema: The OpenAPI schema dict to modify
348 Returns:
349 Modified OpenAPI schema
350 """
351 try:
352 from litellm.proxy._types import ProxyChatCompletionRequest
354 return CustomOpenAPISpec.add_request_schema(
355 openapi_schema=openapi_schema,
356 model_class=ProxyChatCompletionRequest,
357 schema_name="ProxyChatCompletionRequest",
358 paths=CustomOpenAPISpec.CHAT_COMPLETION_PATHS,
359 operation_name="chat completion",
360 )
361 except ImportError as e:
362 verbose_proxy_logger.debug("Failed to import ProxyChatCompletionRequest: %s", e)
363 return openapi_schema
365 @staticmethod
366 def add_embedding_request_schema(openapi_schema: JsonObject) -> JsonObject:
367 """
368 Add EmbeddingRequest schema to embedding endpoints for documentation.
369 This shows the request body in Swagger without runtime validation.
371 Args:
372 openapi_schema: The OpenAPI schema dict to modify
374 Returns:
375 Modified OpenAPI schema
376 """
377 try:
378 from litellm.types.embedding import EmbeddingRequest
380 return CustomOpenAPISpec.add_request_schema(
381 openapi_schema=openapi_schema,
382 model_class=EmbeddingRequest,
383 schema_name="EmbeddingRequest",
384 paths=CustomOpenAPISpec.EMBEDDING_PATHS,
385 operation_name="embedding",
386 )
387 except ImportError as e:
388 verbose_proxy_logger.debug("Failed to import EmbeddingRequest: %s", e)
389 return openapi_schema
391 @staticmethod
392 def add_responses_api_request_schema(
393 openapi_schema: JsonObject,
394 ) -> JsonObject:
395 """
396 Add ResponsesAPIRequestParams schema to responses API endpoints for documentation.
397 This shows the request body in Swagger without runtime validation.
399 Args:
400 openapi_schema: The OpenAPI schema dict to modify
402 Returns:
403 Modified OpenAPI schema
404 """
405 try:
406 from litellm.types.llms.openai import ResponsesAPIRequestParams
408 return CustomOpenAPISpec.add_request_schema(
409 openapi_schema=openapi_schema,
410 model_class=ResponsesAPIRequestParams,
411 schema_name="ResponsesAPIRequestParams",
412 paths=CustomOpenAPISpec.RESPONSES_API_PATHS,
413 operation_name="responses API",
414 )
415 except ImportError as e:
416 verbose_proxy_logger.debug("Failed to import ResponsesAPIRequestParams: %s", e)
417 return openapi_schema
419 @staticmethod
420 def add_llm_api_request_schema_body(
421 openapi_schema: JsonObject,
422 ) -> JsonObject:
423 """
424 Add LLM API request schema bodies to OpenAPI specification for documentation.
426 Args:
427 openapi_schema: The base OpenAPI schema
429 Returns:
430 OpenAPI schema with added request body schemas
431 """
432 # Add chat completion request schema
433 with_chat_completions: Final = CustomOpenAPISpec.add_chat_completion_request_schema(openapi_schema)
435 # Add embedding request schema
436 with_embeddings: Final = CustomOpenAPISpec.add_embedding_request_schema(with_chat_completions)
438 # Add responses API request schema
439 return CustomOpenAPISpec.add_responses_api_request_schema(with_embeddings)