diff --git a/litellm/tests/test_router_retries.py b/litellm/tests/test_router_retries.py index 3ed08dfd916..9b495f3d3e3 100644 --- a/litellm/tests/test_router_retries.py +++ b/litellm/tests/test_router_retries.py @@ -119,3 +119,44 @@ 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 +async def test_router_retry_policy(): + from litellm.router import RetryPolicy + + retry_policy = RetryPolicy( + ContentPolicyViolationErrorRetries=3, + BadRequestErrorRetries=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"), + }, + }, + ], + retry_policy=retry_policy, + ) + + customHandler = MyCustomHandler() + litellm.callbacks = [customHandler] + + try: + + response = await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hey, how do i buy lethal drugs"}], + ) + except Exception as e: + print("got an exception", e) + pass + asyncio.sleep(0.05) + + print("customHandler.previous_models: ", customHandler.previous_models)