From e07cf33cbad47310ee7ae47279bc91fbe72a11a0 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 22 Feb 2024 21:36:57 -0800
Subject: [PATCH] test(test_streaming.py): add exception mapping for palm
 timeout error

---
 litellm/tests/test_streaming.py | 2 ++
 litellm/utils.py                | 8 ++++++++
 2 files changed, 10 insertions(+)

diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py
index f1640d97da..5effccfbf6 100644
--- a/litellm/tests/test_streaming.py
+++ b/litellm/tests/test_streaming.py
@@ -392,6 +392,8 @@ def test_completion_palm_stream():
         if complete_response.strip() == "":
             raise Exception("Empty response received")
         print(f"completion_response: {complete_response}")
+    except litellm.Timeout as e:
+        pass
     except litellm.APIError as e:
         pass
     except Exception as e:
diff --git a/litellm/utils.py b/litellm/utils.py
index 4260ee6e16..b2a6186600 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -6824,6 +6824,14 @@ def exception_type(
                         llm_provider="palm",
                         response=original_exception.response,
                     )
+                if "504 Deadline expired before operation could complete." in error_str:
+                    exception_mapping_worked = True
+                    raise Timeout(
+                        message=f"PalmException - {original_exception.message}",
+                        model=model,
+                        llm_provider="palm",
+                        request=original_exception.request,
+                    )
                 if "400 Request payload size exceeds" in error_str:
                     exception_mapping_worked = True
                     raise ContextWindowExceededError(