From 694488d2ce1d285d7bab842a74208cb37cb49fc8 Mon Sep 17 00:00:00 2001
From: jinno <klvx01@gmail.com>
Date: Sat, 16 Nov 2024 23:22:47 +0900
Subject: [PATCH] fix: duplicate exception_type for gemini

---
 .../exception_mapping_utils.py                | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py
index 3fb276611a..3c747d79ce 100644
--- a/litellm/litellm_core_utils/exception_mapping_utils.py
+++ b/litellm/litellm_core_utils/exception_mapping_utils.py
@@ -1018,7 +1018,6 @@ def exception_type(  # type: ignore  # noqa: PLR0915
             elif (
                 custom_llm_provider == "vertex_ai"
                 or custom_llm_provider == "vertex_ai_beta"
-                or custom_llm_provider == "gemini"
             ):
                 if (
                     "Vertex AI API has not been used in project" in error_str
@@ -1216,7 +1215,7 @@ def exception_type(  # type: ignore  # noqa: PLR0915
                     raise BadRequestError(
                         message="GeminiException - Invalid api key",
                         model=model,
-                        llm_provider="palm",
+                        llm_provider=custom_llm_provider,
                         response=getattr(original_exception, "response", None),
                     )
                 if (
@@ -1227,14 +1226,14 @@ def exception_type(  # type: ignore  # noqa: PLR0915
                     raise Timeout(
                         message=f"GeminiException - {original_exception.message}",
                         model=model,
-                        llm_provider="palm",
+                        llm_provider=custom_llm_provider,
                     )
                 if "400 Request payload size exceeds" in error_str:
                     exception_mapping_worked = True
                     raise ContextWindowExceededError(
                         message=f"GeminiException - {error_str}",
                         model=model,
-                        llm_provider="palm",
+                        llm_provider=custom_llm_provider,
                         response=getattr(original_exception, "response", None),
                     )
                 if (
@@ -1245,7 +1244,7 @@ def exception_type(  # type: ignore  # noqa: PLR0915
                     raise APIError(
                         status_code=getattr(original_exception, "status_code", 500),
                         message=f"GeminiException - {original_exception.message}",
-                        llm_provider="palm",
+                        llm_provider=custom_llm_provider,
                         model=model,
                         request=httpx.Response(
                             status_code=429,
@@ -1261,7 +1260,15 @@ def exception_type(  # type: ignore  # noqa: PLR0915
                         raise BadRequestError(
                             message=f"GeminiException - {error_str}",
                             model=model,
-                            llm_provider="palm",
+                            llm_provider=custom_llm_provider,
+                            response=getattr(original_exception, "response", None),
+                        )
+                    if original_exception.status_code == 503:
+                        exception_mapping_worked = True
+                        raise BadRequestError(
+                            message=f"GeminiException - {error_str}",
+                            model=model,
+                            llm_provider=custom_llm_provider,
                             response=getattr(original_exception, "response", None),
                         )
                 # Dailed: Error occurred: 400 Request payload size exceeds the limit: 20000 bytes