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

@ -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}",