diff --git a/litellm/tests/test_image_generation.py b/litellm/tests/test_image_generation.py index 1d3d53f70a1..49ec18f24cd 100644 --- a/litellm/tests/test_image_generation.py +++ b/litellm/tests/test_image_generation.py @@ -195,6 +195,8 @@ async def test_aimage_generation_vertex_ai(): assert isinstance(d, litellm.ImageObject) print("data in response.data", d) assert d.b64_json is not None + except litellm.ServiceUnavailableError as e: + pass except litellm.RateLimitError as e: pass except litellm.ContentPolicyViolationError: diff --git a/litellm/utils.py b/litellm/utils.py index dbaa3d0939d..5e85419dcd0 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -9080,7 +9080,7 @@ def exception_type( ): exception_mapping_worked = True raise RateLimitError( - message=f"VertexAIException RateLimitError - {error_str}", + message=f"litellm.RateLimitError: VertexAIException - {error_str}", model=model, llm_provider="vertex_ai", litellm_debug_info=extra_information, @@ -9092,7 +9092,14 @@ def exception_type( ), ), ) - + elif "500 Internal Server Error" in error_str: + exception_mapping_worked = True + raise ServiceUnavailableError( + message=f"litellm.ServiceUnavailableError: VertexAIException - {error_str}", + model=model, + llm_provider="vertex_ai", + litellm_debug_info=extra_information, + ) if hasattr(original_exception, "status_code"): if original_exception.status_code == 400: exception_mapping_worked = True