diff --git a/litellm/router_utils/get_retry_from_policy.py b/litellm/router_utils/get_retry_from_policy.py index ec326ebb50d..5ec7175c3a9 100644 --- a/litellm/router_utils/get_retry_from_policy.py +++ b/litellm/router_utils/get_retry_from_policy.py @@ -10,6 +10,7 @@ from litellm.exceptions import ( AuthenticationError, BadRequestError, ContentPolicyViolationError, + InternalServerError, RateLimitError, Timeout, ) @@ -28,6 +29,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 @@ -55,6 +57,11 @@ def get_num_retries_from_retry_policy( and retry_policy.RateLimitErrorRetries is not None ): return retry_policy.RateLimitErrorRetries + if ( + isinstance(exception, InternalServerError) + and retry_policy.InternalServerErrorRetries is not None + ): + return retry_policy.InternalServerErrorRetries if ( isinstance(exception, ContentPolicyViolationError) and retry_policy.ContentPolicyViolationErrorRetries is not None diff --git a/tests/litellm/router_utils/test_get_retry_from_policy.py b/tests/litellm/router_utils/test_get_retry_from_policy.py new file mode 100644 index 00000000000..9c8473333cf --- /dev/null +++ b/tests/litellm/router_utils/test_get_retry_from_policy.py @@ -0,0 +1,39 @@ +"""Unit tests for litellm.router_utils.get_retry_from_policy.""" + +import litellm +from litellm.router_utils.get_retry_from_policy import ( + get_num_retries_from_retry_policy, +) +from litellm.types.router import RetryPolicy + + +def test_internal_server_error_retries_is_honored(): + """Regression: `InternalServerErrorRetries` must be returned for + `InternalServerError` exceptions. Previously the dispatcher had no + branch for this field and silently returned `None`, causing the + caller to fall back to `num_retries`.""" + retry_policy = RetryPolicy(InternalServerErrorRetries=5) + exc = litellm.exceptions.InternalServerError( + message="test", llm_provider="openai", model="gpt-3.5-turbo" + ) + + num_retries = get_num_retries_from_retry_policy( + exception=exc, retry_policy=retry_policy + ) + + assert num_retries == 5 + + +def test_internal_server_error_retries_unset_returns_none(): + """When the field is not set, the dispatcher should return `None` + so the caller falls back to `num_retries`.""" + retry_policy = RetryPolicy() + exc = litellm.exceptions.InternalServerError( + message="test", llm_provider="openai", model="gpt-3.5-turbo" + ) + + num_retries = get_num_retries_from_retry_policy( + exception=exc, retry_policy=retry_policy + ) + + assert num_retries is None diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index 34a19f5ce79..68890e6ee4a 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -957,6 +957,7 @@ def test_track_deployment_metrics(model_list): "ContentPolicyViolationError", 7, ), + (litellm.exceptions.InternalServerError, "InternalServerError", 5), ], ) def test_get_num_retries_from_retry_policy(