From 7e3734d037dbde3538d207082d5e03a8211dbea7 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Wed, 6 Mar 2024 19:05:39 -0800 Subject: [PATCH] test(test_completion.py): handle gemini timeout error --- litellm/tests/test_completion.py | 2 ++ litellm/utils.py | 5 ++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index f5e145769..8e695f3f7 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -2139,6 +2139,8 @@ async def test_acompletion_gemini(): response = await litellm.acompletion(model=model_name, messages=messages) # Add any assertions here to check the response print(f"response: {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 c2b1a730f..d09f330a1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7105,7 +7105,10 @@ def exception_type( llm_provider="palm", response=original_exception.response, ) - if "504 Deadline expired before operation could complete." in error_str: + if ( + "504 Deadline expired before operation could complete." in error_str + or "504 Deadline Exceeded" in error_str + ): exception_mapping_worked = True raise Timeout( message=f"PalmException - {original_exception.message}",