From b92700cc19d863ce101de2a2bbe67b81900cb7fc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 14 Nov 2024 23:47:14 +0530 Subject: [PATCH] test: handle gemini overloaded model error --- litellm/litellm_core_utils/exception_mapping_utils.py | 9 ++++++--- tests/local_testing/test_stream_chunk_builder.py | 4 ++-- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index a4a30fc31..ca1de75be 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -1124,10 +1124,13 @@ def exception_type( # type: ignore # noqa: PLR0915 ), ), ) - elif "500 Internal Server Error" in error_str: + elif ( + "500 Internal Server Error" in error_str + or "The model is overloaded." in error_str + ): exception_mapping_worked = True - raise ServiceUnavailableError( - message=f"litellm.ServiceUnavailableError: VertexAIException - {error_str}", + raise litellm.InternalServerError( + message=f"litellm.InternalServerError: VertexAIException - {error_str}", model=model, llm_provider="vertex_ai", litellm_debug_info=extra_information, diff --git a/tests/local_testing/test_stream_chunk_builder.py b/tests/local_testing/test_stream_chunk_builder.py index a65d440a7..2548abdb7 100644 --- a/tests/local_testing/test_stream_chunk_builder.py +++ b/tests/local_testing/test_stream_chunk_builder.py @@ -205,8 +205,8 @@ def test_stream_chunk_builder_litellm_usage_chunks(): complete_response=True, stream_options={"include_usage": True}, ) - except litellm.ServiceUnavailableError as e: - pytest.skip(f"ServiceUnavailableError - {str(e)}") + except litellm.InternalServerError as e: + pytest.skip(f"Skipping test due to internal server error - {str(e)}") usage: litellm.Usage = response.usage