mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 19:24:27 +00:00
add test
This commit is contained in:
parent
286d42a881
commit
3caa9dfec6
1 changed files with 45 additions and 0 deletions
|
@ -186,6 +186,51 @@ async def test_router_retry_policy(error_type):
|
||||||
assert customHandler.previous_models == 3
|
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.parametrize("model_group", ["gpt-3.5-turbo", "bad-model"])
|
||||||
@pytest.mark.asyncio
|
@pytest.mark.asyncio
|
||||||
async def test_dynamic_router_retry_policy(model_group):
|
async def test_dynamic_router_retry_policy(model_group):
|
||||||
|
|
Loading…
Add table
Add a link
Reference in a new issue