From 7f4da14d7273c4cc03c6c34047f4da9f354c41f4 Mon Sep 17 00:00:00 2001 From: kris3984 Date: Sat, 22 Aug 2026 19:48:10 +0530 Subject: [PATCH] fix(utils): handle Pydantic schema conversion and map validation errors to APIError --- litellm/main.py | 10 +- litellm/utils.py | 119 ++++++++++++--- .../test_litellm/test_pydantic_validation.py | 141 ++++++++++++++++++ 3 files changed, 243 insertions(+), 27 deletions(-) create mode 100644 tests/test_litellm/test_pydantic_validation.py diff --git a/litellm/main.py b/litellm/main.py index 7cfd322f3d0..b882f4e3517 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -152,6 +152,7 @@ from litellm.utils import ( get_secret, get_standard_openai_params, mock_completion_streaming_obj, + normalize_completion_response_format, pre_process_non_default_params, read_config_args, should_run_mock_completion, @@ -481,6 +482,8 @@ async def acompletion( - The `completion` function is called using `run_in_executor` to execute synchronously in the event loop. - If `stream` is True, the function returns an async generator that yields completion lines. """ + request_response_format: Final = normalize_completion_response_format(response_format, model=model) + fallbacks = kwargs.get("fallbacks", None) mock_timeout = kwargs.get("mock_timeout", None) @@ -572,7 +575,7 @@ async def acompletion( "frequency_penalty": frequency_penalty, "logit_bias": logit_bias, "user": user, - "response_format": response_format, + "response_format": request_response_format, "seed": seed, "tools": tools, "tool_choice": tool_choice, @@ -5018,6 +5021,7 @@ def completion( # model whose model_cost mode is "responses" but whose provider has no # Responses API config (get_provider_responses_api_config -> None). skip_responses_api_bridge: Final = kwargs.pop("_skip_responses_api_bridge", False) + request_response_format: Final = normalize_completion_response_format(response_format, model=model) skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) if not skip_mcp_handler and tools: @@ -5052,7 +5056,7 @@ def completion( frequency_penalty=frequency_penalty, logit_bias=logit_bias, user=user, - response_format=response_format, + response_format=request_response_format, seed=seed, tools=tools, tool_choice=tool_choice, @@ -5355,7 +5359,7 @@ def completion( # params to identify the model "model": model, "custom_llm_provider": custom_llm_provider, - "response_format": response_format, + "response_format": request_response_format, "seed": seed, "tools": tools, "tool_choice": tool_choice, diff --git a/litellm/utils.py b/litellm/utils.py index e5ce7157e77..d4eab782e64 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -42,9 +42,8 @@ import openai import tiktoken from httpx import Proxy from httpx._utils import get_environment_proxies -from openai.lib import _parsing, _pydantic from openai.types.chat.completion_create_params import ResponseFormat -from pydantic import BaseModel +from pydantic import BaseModel, ValidationError from tiktoken import Encoding from tokenizers import Tokenizer @@ -1253,6 +1252,93 @@ async def async_post_call_success_deployment_hook( return response +def process_response_format( + response_format: type[BaseModel] | dict | None, +) -> dict | None: + if response_format is None: + return None + if isinstance(response_format, dict): + return type_to_response_format_param(response_format) + if isinstance(response_format, type) and issubclass(response_format, BaseModel): + return type_to_response_format_param(response_format) + raise TypeError(f"Unsupported response_format type - {response_format}") + + +def normalize_completion_response_format( + response_format: type[BaseModel] | dict | None, + model: str, +) -> dict | type[BaseModel] | None: + try: + processed: Final = process_response_format(response_format) + except (ValidationError, json.JSONDecodeError) as e: + raise litellm.APIError( + status_code=400, + message=f"Invalid Pydantic response_format: {e}", + llm_provider="", + model=model, + ) from e + return processed if processed is not None else response_format + + +def _deserialize_pydantic_response_format( + response_format: type[BaseModel], + model_response: str, + model: str | None, +) -> None: + try: + model_validate_json = getattr(response_format, "model_validate_json", None) + if callable(model_validate_json): + model_validate_json(model_response) + return + parse_raw = getattr(response_format, "parse_raw", None) + if callable(parse_raw): + parse_raw(model_response) + return + json.loads(model_response) + except (ValidationError, json.JSONDecodeError) as e: + raise litellm.APIError( + status_code=500, + message=f"Structured output did not match the Pydantic response_format: {e}", + llm_provider="", + model=model or "", + ) from e + + +def _response_format_as_json_schema(response_format: object) -> dict | None: + if isinstance(response_format, type) and issubclass(response_format, BaseModel): + return process_response_format(response_format) + if isinstance(response_format, dict) and response_format.get("json_schema") is not None: + return response_format + return None + + +def _apply_response_format_validation( + response_format: object, + model_response: str, + model: str | None, +) -> None: + try: + if isinstance(response_format, type) and issubclass(response_format, BaseModel): + _deserialize_pydantic_response_format( + response_format=response_format, + model_response=model_response, + model=model, + ) + json_response_format: Final = _response_format_as_json_schema(response_format) + if json_response_format is not None: + litellm.litellm_core_utils.json_validation_rule.validate_schema( + schema=json_response_format["json_schema"]["schema"], + response=model_response, + ) + except (ValidationError, json.JSONDecodeError) as e: + raise litellm.APIError( + status_code=500, + message=f"Structured output did not match the Pydantic response_format: {e}", + llm_provider="", + model=model or "", + ) from e + + def post_call_processing( original_response, model, @@ -1294,26 +1380,11 @@ def post_call_processing( and "response_format" in optional_params and optional_params["response_format"] is not None ): - json_response_format: dict | None = None - if ( - isinstance( - optional_params["response_format"], - dict, - ) - and optional_params["response_format"].get("json_schema") is not None - ): - json_response_format = optional_params["response_format"] - elif _parsing._completions.is_basemodel_type( - optional_params["response_format"] - ): - json_response_format = type_to_response_format_param( - response_format=optional_params["response_format"] - ) - if json_response_format is not None: - litellm.litellm_core_utils.json_validation_rule.validate_schema( - schema=json_response_format["json_schema"]["schema"], - response=model_response, - ) + _apply_response_format_validation( + response_format=optional_params["response_format"], + model_response=model_response, + model=model, + ) except TypeError: pass if ( @@ -3817,8 +3888,8 @@ def pre_process_non_default_params( response_format=non_default_params["response_format"] ) else: - non_default_params["response_format"] = type_to_response_format_param( - response_format=non_default_params["response_format"] + non_default_params["response_format"] = process_response_format( + non_default_params["response_format"] ) if "tools" in non_default_params and isinstance( diff --git a/tests/test_litellm/test_pydantic_validation.py b/tests/test_litellm/test_pydantic_validation.py new file mode 100644 index 00000000000..e9bf7249676 --- /dev/null +++ b/tests/test_litellm/test_pydantic_validation.py @@ -0,0 +1,141 @@ +import json +from typing import Final +from unittest.mock import patch + +import pytest +from pydantic import BaseModel, ValidationError + +import litellm +from litellm.llms.base_llm.base_utils import _pydantic_model_json_schema, type_to_response_format_param +from litellm.types.utils import ModelResponse +from litellm.utils import Rules, post_call_processing, process_response_format + + +class MovieReview(BaseModel): + title: str + rating: int + + +def _mock_completion(): + pass + + +_mock_completion.__name__ = "completion" + + +def _make_response(content: str) -> ModelResponse: + response = ModelResponse() + response.choices[0].message.content = content + return response + + +def test_process_response_format_converts_pydantic_v2_basemodel(): + processed: Final = process_response_format(MovieReview) + + assert processed is not None + assert processed["type"] == "json_schema" + json_schema: Final = processed["json_schema"] + assert json_schema["name"] == "MovieReview" + assert json_schema["strict"] is True + schema: Final = json_schema["schema"] + assert schema["type"] == "object" + assert "title" in schema["properties"] + assert "rating" in schema["properties"] + assert schema["properties"]["title"]["type"] == "string" + assert schema["properties"]["rating"]["type"] == "integer" + + +def test_process_response_format_passthrough_none_and_dict(): + existing: Final = { + "type": "json_schema", + "json_schema": { + "name": "MovieReview", + "schema": {"type": "object", "properties": {"title": {"type": "string"}}}, + }, + } + assert process_response_format(None) is None + assert process_response_format(existing)["json_schema"]["name"] == "MovieReview" + + +def test_pydantic_v1_schema_fallback_when_model_json_schema_missing(): + class LegacyShape(BaseModel): + x: str + + def _v1_schema() -> dict: + return { + "title": "LegacyShape", + "type": "object", + "properties": {"x": {"title": "X", "type": "string"}}, + } + + with patch.object(LegacyShape, "model_json_schema", None): + with patch.object(LegacyShape, "schema", staticmethod(_v1_schema)): + schema: Final = _pydantic_model_json_schema(LegacyShape) + + assert schema["properties"]["x"]["type"] == "string" + assert schema["title"] == "LegacyShape" + + +def test_type_to_response_format_param_falls_back_when_strict_schema_fails(): + with patch( + "litellm.llms.base_llm.base_utils._pydantic.to_strict_json_schema", + side_effect=ValidationError.from_exception_data("MovieReview", []), + ): + processed: Final = type_to_response_format_param(MovieReview) + + assert processed is not None + assert processed["json_schema"]["schema"]["properties"]["title"]["type"] == "string" + + +def test_post_call_processing_raises_apierror_on_invalid_pydantic_json(): + with pytest.raises(litellm.APIError, match="Structured output"): + post_call_processing( + _make_response("not-json"), + "gpt-4o", + { + "response_format": MovieReview, + "enable_json_schema_validation": True, + }, + _mock_completion, + Rules(), + ) + + +def test_post_call_processing_raises_apierror_on_pydantic_validation_error(): + with pytest.raises(litellm.APIError, match="Structured output"): + post_call_processing( + _make_response(json.dumps({"title": "Inception", "rating": "nine"})), + "gpt-4o", + { + "response_format": MovieReview, + "enable_json_schema_validation": True, + }, + _mock_completion, + Rules(), + ) + + +def test_post_call_processing_accepts_valid_pydantic_response(): + post_call_processing( + _make_response(json.dumps({"title": "Inception", "rating": 9})), + "gpt-4o", + { + "response_format": MovieReview, + "enable_json_schema_validation": True, + }, + _mock_completion, + Rules(), + ) + + +def test_completion_converts_pydantic_response_format_with_mock_response(): + response: Final = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "review"}], + response_format=MovieReview, + mock_response=json.dumps({"title": "Inception", "rating": 9}), + ) + assert response.choices[0].message.content is not None + payload: Final = json.loads(response.choices[0].message.content) + assert payload["title"] == "Inception" + assert payload["rating"] == 9