diff --git a/docs/my-website/docs/routing.md b/docs/my-website/docs/routing.md index 0aa7901c058..4fa166548f6 100644 --- a/docs/my-website/docs/routing.md +++ b/docs/my-website/docs/routing.md @@ -616,6 +616,57 @@ response = router.completion(model="gpt-3.5-turbo", messages=messages) print(f"response: {response}") ``` +#### Retries based on Error Type + +Use `RetryPolicy` if you want to set a `num_retries` based on the Exception receieved + +Example: +- 4 retries for `ContentPolicyViolationError` +- 0 retries for `RateLimitErrors` + +Example Usage + +```python +from litellm.router import RetryPolicy +retry_policy = RetryPolicy( + ContentPolicyViolationErrorRetries=3, # run 3 retries for ContentPolicyViolationErrors + AuthenticationErrorRetries=0, # run 0 retries for AuthenticationErrorRetries + BadRequestErrorRetries=1, + TimeoutErrorRetries=2, + RateLimitErrorRetries=3, +) + +router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": os.getenv("AZURE_API_KEY"), + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + }, + { + "model_name": "bad-model", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": "bad-key", + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + }, + ], + retry_policy=retry_policy, +) + +response = await router.acompletion( + model=model, + messages=messages, +) +``` + + ### Fallbacks If a call fails after num_retries, fall back to another model group. diff --git a/litellm/main.py b/litellm/main.py index 59d98580cf1..d19463f532f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -620,6 +620,7 @@ def completion( "model_list", "num_retries", "context_window_fallback_dict", + "retry_policy", "roles", "final_prompt_value", "bos_token", @@ -2687,6 +2688,7 @@ def embedding( "model_list", "num_retries", "context_window_fallback_dict", + "retry_policy", "roles", "final_prompt_value", "bos_token", @@ -3556,6 +3558,7 @@ def image_generation( "model_list", "num_retries", "context_window_fallback_dict", + "retry_policy", "roles", "final_prompt_value", "bos_token", diff --git a/litellm/router.py b/litellm/router.py index d64deecec10..3b1c1d1022a 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -42,6 +42,7 @@ from litellm.types.router import ( RouterErrors, updateDeployment, updateLiteLLMParams, + RetryPolicy, ) from litellm.integrations.custom_logger import CustomLogger @@ -82,6 +83,9 @@ class Router: model_group_alias: Optional[dict] = {}, enable_pre_call_checks: bool = False, retry_after: int = 0, # min time to wait before retrying a failed request + retry_policy: Optional[ + RetryPolicy + ] = None, # set custom retries for different exceptions allowed_fails: Optional[ int ] = None, # Number of times a deployment can failbefore being added to cooldown @@ -303,6 +307,7 @@ class Router: f"Intialized router with Routing strategy: {self.routing_strategy}\n\nRouting fallbacks: {self.fallbacks}\n\nRouting context window fallbacks: {self.context_window_fallbacks}\n\nRouter Redis Caching={self.cache.redis_cache}" ) # noqa self.routing_strategy_args = routing_strategy_args + self.retry_policy: Optional[RetryPolicy] = retry_policy def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: dict): if routing_strategy == "least-busy": @@ -1504,6 +1509,15 @@ class Router: ) await asyncio.sleep(_timeout) ## LOGGING + if self.retry_policy is not None or kwargs.get("retry_policy") is not None: + # get num_retries from retry policy + _retry_policy_retries = self.get_num_retries_from_retry_policy( + exception=original_exception, + dynamic_retry_policy=kwargs.get("retry_policy"), + ) + if _retry_policy_retries is not None: + num_retries = _retry_policy_retries + if num_retries > 0: kwargs = self.log_retry(kwargs=kwargs, e=original_exception) @@ -3254,6 +3268,48 @@ class Router: except Exception as e: verbose_router_logger.error(f"Error in _track_deployment_metrics: {str(e)}") + def get_num_retries_from_retry_policy( + self, exception: Exception, dynamic_retry_policy: Optional[RetryPolicy] = None + ): + """ + BadRequestErrorRetries: Optional[int] = None + AuthenticationErrorRetries: Optional[int] = None + TimeoutErrorRetries: Optional[int] = None + RateLimitErrorRetries: Optional[int] = None + ContentPolicyViolationErrorRetries: Optional[int] = None + """ + # if we can find the exception then in the retry policy -> return the number of retries + retry_policy = self.retry_policy + if dynamic_retry_policy is not None: + retry_policy = dynamic_retry_policy + if retry_policy is None: + return None + if ( + isinstance(exception, litellm.BadRequestError) + and retry_policy.BadRequestErrorRetries is not None + ): + return retry_policy.BadRequestErrorRetries + if ( + isinstance(exception, litellm.AuthenticationError) + and retry_policy.AuthenticationErrorRetries is not None + ): + return retry_policy.AuthenticationErrorRetries + if ( + isinstance(exception, litellm.Timeout) + and retry_policy.TimeoutErrorRetries is not None + ): + return retry_policy.TimeoutErrorRetries + if ( + isinstance(exception, litellm.RateLimitError) + and retry_policy.RateLimitErrorRetries is not None + ): + return retry_policy.RateLimitErrorRetries + if ( + isinstance(exception, litellm.ContentPolicyViolationError) + and retry_policy.ContentPolicyViolationErrorRetries is not None + ): + return retry_policy.ContentPolicyViolationErrorRetries + def flush_cache(self): litellm.cache = None self.cache.flush_cache() @@ -3264,4 +3320,5 @@ class Router: litellm.__async_success_callback = [] litellm.failure_callback = [] litellm._async_failure_callback = [] + self.retry_policy = None self.flush_cache() diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py index 3ed08dfd916..8828e286e75 100644 --- a/litellm/tests/test_router_retries.py +++ b/litellm/tests/test_router_retries.py @@ -119,3 +119,123 @@ async def test_router_retries_errors(sync_mode, error_type): assert customHandler.previous_models == 0 # 0 retries else: assert customHandler.previous_models == 2 # 2 retries + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "error_type", ["AuthenticationErrorRetries", "ContentPolicyViolationErrorRetries"] +) +async def test_router_retry_policy(error_type): + from litellm.router import RetryPolicy + + retry_policy = RetryPolicy( + ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0 + ) + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": os.getenv("AZURE_API_KEY"), + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + }, + { + "model_name": "bad-model", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": "bad-key", + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + }, + ], + retry_policy=retry_policy, + ) + + customHandler = MyCustomHandler() + litellm.callbacks = [customHandler] + if error_type == "AuthenticationErrorRetries": + model = "bad-model" + messages = [{"role": "user", "content": "Hello good morning"}] + elif error_type == "ContentPolicyViolationErrorRetries": + model = "gpt-3.5-turbo" + messages = [{"role": "user", "content": "where do i buy lethal drugs from"}] + + try: + litellm.set_verbose = True + response = await router.acompletion( + model=model, + messages=messages, + ) + except Exception as e: + print("got an exception", e) + pass + asyncio.sleep(0.05) + + print("customHandler.previous_models: ", customHandler.previous_models) + + if error_type == "AuthenticationErrorRetries": + assert customHandler.previous_models == 0 + elif error_type == "ContentPolicyViolationErrorRetries": + assert customHandler.previous_models == 3 + + +@pytest.mark.parametrize("model_group", ["gpt-3.5-turbo", "bad-model"]) +@pytest.mark.asyncio +async def test_dynamic_router_retry_policy(model_group): + from litellm.router import RetryPolicy + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": os.getenv("AZURE_API_KEY"), + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + }, + { + "model_name": "bad-model", # openai model name + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/chatgpt-v-2", + "api_key": "bad-key", + "api_version": os.getenv("AZURE_API_VERSION"), + "api_base": os.getenv("AZURE_API_BASE"), + }, + }, + ] + ) + + customHandler = MyCustomHandler() + litellm.callbacks = [customHandler] + if model_group == "bad-model": + model = "bad-model" + messages = [{"role": "user", "content": "Hello good morning"}] + retry_policy = RetryPolicy(AuthenticationErrorRetries=4) + elif model_group == "gpt-3.5-turbo": + model = "gpt-3.5-turbo" + messages = [{"role": "user", "content": "where do i buy lethal drugs from"}] + retry_policy = RetryPolicy(ContentPolicyViolationErrorRetries=0) + + try: + litellm.set_verbose = True + response = await router.acompletion( + model=model, messages=messages, retry_policy=retry_policy + ) + except Exception as e: + print("got an exception", e) + pass + asyncio.sleep(0.05) + + print("customHandler.previous_models: ", customHandler.previous_models) + + if model_group == "bad-model": + assert customHandler.previous_models == 4 + elif model_group == "gpt-3.5-turbo": + assert customHandler.previous_models == 0 diff --git a/litellm/types/router.py b/litellm/types/router.py index 068a99b0059..1f5fb7103f9 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -266,3 +266,18 @@ class RouterErrors(enum.Enum): user_defined_ratelimit_error = "Deployment over user-defined ratelimit." no_deployments_available = "No deployments available for selected model" + + +class RetryPolicy(BaseModel): + """ + Use this to set a custom number of retries per exception type + If RateLimitErrorRetries = 3, then 3 retries will be made for RateLimitError + Mapping of Exception type to number of retries + https://docs.litellm.ai/docs/exception_mapping + """ + + BadRequestErrorRetries: Optional[int] = None + AuthenticationErrorRetries: Optional[int] = None + TimeoutErrorRetries: Optional[int] = None + RateLimitErrorRetries: Optional[int] = None + ContentPolicyViolationErrorRetries: Optional[int] = None