fix(utils.py): map cohere timeout error

This commit is contained in:
Krrish Dholakia 2024-07-31 15:15:18 -07:00
parent 6202f9bbb0
commit bd68714f51
2 changed files with 9 additions and 0 deletions

View file

@ -719,6 +719,8 @@ def test_completion_cohere_command_r_plus_function_call():
force_single_step=True, force_single_step=True,
) )
print(second_response) print(second_response)
except litellm.Timeout:
pass
except Exception as e: except Exception as e:
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")

View file

@ -7328,6 +7328,13 @@ def exception_type(
model=model, model=model,
response=original_exception.response, response=original_exception.response,
) )
elif original_exception.status_code == 408:
exception_mapping_worked = True
raise Timeout(
message=f"CohereException - {original_exception.message}",
llm_provider="cohere",
model=model,
)
elif original_exception.status_code == 500: elif original_exception.status_code == 500:
exception_mapping_worked = True exception_mapping_worked = True
raise ServiceUnavailableError( raise ServiceUnavailableError(