From 2e49614230c08a5af188470d3f987f5b32ca92c9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 10 Jun 2024 21:36:48 -0700 Subject: [PATCH] fix(utils.py): exception map vertex ai 500 internal server errors --- litellm/tests/test_image_generation.py | 2 ++ litellm/utils.py | 11 +++++++++-- 2 files changed, 11 insertions(+), 2 deletions(-) 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