diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index 7cf55e80e0c..2029b5d4d9d 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -8,6 +8,7 @@ from litellm.exceptions import ( AuthenticationError, BadRequestError, ContentPolicyViolationError, + InternalServerError, RateLimitError, Timeout, ) @@ -26,6 +27,7 @@ def get_num_retries_from_retry_policy( TimeoutErrorRetries: Optional[int] = None RateLimitErrorRetries: Optional[int] = None ContentPolicyViolationErrorRetries: Optional[int] = None + InternalServerErrorRetries: Optional[int] = None """ # if we can find the exception then in the retry policy -> return the number of retries @@ -50,6 +52,8 @@ def get_num_retries_from_retry_policy( return retry_policy.ContentPolicyViolationErrorRetries if isinstance(exception, BadRequestError) and retry_policy.BadRequestErrorRetries is not None: return retry_policy.BadRequestErrorRetries + if isinstance(exception, InternalServerError) and retry_policy.InternalServerErrorRetries is not None: + return retry_policy.InternalServerErrorRetries def reset_retry_policy() -> RetryPolicy: diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index f81578dbd99..62ccfa44f66 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -1354,6 +1354,11 @@ def test_track_deployment_metrics(model_list): "ContentPolicyViolationError", 7, ), + ( + litellm.exceptions.InternalServerError, + "InternalServerError", + 5, + ), ], ) def test_get_num_retries_from_retry_policy(