This commit is contained in:
Ishaan Jaff 2024-06-01 17:03:53 -07:00
parent 286d42a881
commit 3caa9dfec6

View file

@ -186,6 +186,51 @@ async def test_router_retry_policy(error_type):
assert customHandler.previous_models == 3
@pytest.mark.asyncio
@pytest.mark.skip(
reason="This is a local only test, use this to confirm if retry policy works"
)
async def test_router_retry_policy_on_429_errprs():
from litellm.router import RetryPolicy
retry_policy = RetryPolicy(
RateLimitErrorRetries=2,
)
router = Router(
model_list=[
{
"model_name": "gpt-3.5-turbo", # openai model name
"litellm_params": {
"model": "vertex_ai/gemini-1.5-pro-001",
},
},
],
retry_policy=retry_policy,
# set_verbose=True,
# debug_level="DEBUG",
allowed_fails=10,
)
customHandler = MyCustomHandler()
litellm.callbacks = [customHandler]
try:
# litellm.set_verbose = True
_one_message = [{"role": "user", "content": "Hello good morning"}]
messages = [_one_message] * 5
print("messages: ", messages)
responses = await router.abatch_completion(
models=["gpt-3.5-turbo"],
messages=messages,
)
print("responses: ", responses)
except Exception as e:
print("got an exception", e)
pass
asyncio.sleep(0.05)
print("customHandler.previous_models: ", customHandler.previous_models)
@pytest.mark.parametrize("model_group", ["gpt-3.5-turbo", "bad-model"])
@pytest.mark.asyncio
async def test_dynamic_router_retry_policy(model_group):