diff --git a/litellm/llms/chatgpt/image_generation/response_parsing.py b/litellm/llms/chatgpt/image_generation/response_parsing.py index df848d696ca..bd0c4d7d4a0 100644 --- a/litellm/llms/chatgpt/image_generation/response_parsing.py +++ b/litellm/llms/chatgpt/image_generation/response_parsing.py @@ -164,6 +164,33 @@ def parse_sse_payloads(body_text: str) -> List[dict]: return payloads +def extract_error_status_code(payload: dict, error_obj: Any) -> int: + possible_status_codes: List[Any] = [] + if isinstance(error_obj, dict): + possible_status_codes.extend( + [error_obj.get("status_code"), error_obj.get("status")] + ) + + response_payload = payload.get("response") + if isinstance(response_payload, dict): + possible_status_codes.extend( + [response_payload.get("status_code"), response_payload.get("status")] + ) + + possible_status_codes.extend([payload.get("status_code"), payload.get("status")]) + + for status_code in possible_status_codes: + if isinstance(status_code, bool): + continue + if isinstance(status_code, int) and 100 <= status_code <= 599: + return status_code + if isinstance(status_code, str) and status_code.isdigit(): + parsed_status_code = int(status_code) + if 100 <= parsed_status_code <= 599: + return parsed_status_code + return 400 + + def extract_images_from_payload(payload: dict) -> Tuple[List[str], List[str]]: event_type = payload.get("type") if event_type in ( @@ -171,7 +198,10 @@ def extract_images_from_payload(payload: dict) -> Tuple[List[str], List[str]]: ResponsesAPIStreamEvents.ERROR, ): error_obj = payload.get("error") or (payload.get("response") or {}).get("error") - raise OpenAIError(message=str(error_obj or payload), status_code=400) + raise OpenAIError( + message=str(error_obj or payload), + status_code=extract_error_status_code(payload, error_obj), + ) partial_images: List[str] = [] if event_type in ( diff --git a/tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py b/tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py index d8851f345eb..d29de526ebb 100644 --- a/tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py +++ b/tests/test_litellm/llms/chatgpt/image_generation/test_response_parsing.py @@ -129,6 +129,65 @@ def test_chatgpt_image_generation_raises_provider_error_event(): ) +def test_chatgpt_image_generation_preserves_response_error_status(): + config = ChatGPTImageGenerationConfig() + + with pytest.raises(OpenAIError, match="token expired") as exc_info: + config._extract_images_from_payload( + { + "type": "response.failed", + "response": {"error": {"message": "token expired", "status": 401}}, + } + ) + + assert exc_info.value.status_code == 401 + + +def test_chatgpt_image_generation_preserves_top_level_error_status_code(): + config = ChatGPTImageGenerationConfig() + + with pytest.raises(OpenAIError, match="rate limited") as exc_info: + config._extract_images_from_payload( + { + "type": "error", + "error": {"message": "rate limited", "status_code": 429}, + } + ) + + assert exc_info.value.status_code == 429 + + +def test_chatgpt_image_generation_parses_string_error_status(): + config = ChatGPTImageGenerationConfig() + + with pytest.raises(OpenAIError, match="server error") as exc_info: + config._extract_images_from_payload( + { + "type": "error", + "error": {"message": "server error", "status": "500"}, + } + ) + + assert exc_info.value.status_code == 500 + + +def test_chatgpt_image_generation_error_status_falls_back_to_400(): + config = ChatGPTImageGenerationConfig() + + with pytest.raises(OpenAIError, match="rate limit code") as exc_info: + config._extract_images_from_payload( + { + "type": "error", + "error": { + "message": "rate limit code", + "code": "rate_limit_exceeded", + }, + } + ) + + assert exc_info.value.status_code == 400 + + def test_chatgpt_image_generation_handles_invalid_json_payloads(): config = ChatGPTImageGenerationConfig() raw_response = httpx.Response(