test(test_completion.py): handle gemini timeout error

This commit is contained in:
Krrish Dholakia 2024-03-06 19:05:39 -08:00
parent 7d824225a5
commit 7e3734d037
2 changed files with 6 additions and 1 deletions

View file

@ -2139,6 +2139,8 @@ async def test_acompletion_gemini():
response = await litellm.acompletion(model=model_name, messages=messages) response = await litellm.acompletion(model=model_name, messages=messages)
# Add any assertions here to check the response # Add any assertions here to check the response
print(f"response: {response}") print(f"response: {response}")
except litellm.Timeout as e:
pass
except litellm.APIError as e: except litellm.APIError as e:
pass pass
except Exception as e: except Exception as e:

View file

@ -7105,7 +7105,10 @@ def exception_type(
llm_provider="palm", llm_provider="palm",
response=original_exception.response, 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 exception_mapping_worked = True
raise Timeout( raise Timeout(
message=f"PalmException - {original_exception.message}", message=f"PalmException - {original_exception.message}",