phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

refactor(test_router_caching.py): move tpm/rpm routing tests to separate file

Browse source
This commit is contained in:
Krrish Dholakia 2024-01-02 11:10:11 +05:30
parent 18ef244230
commit dff4c172d0
3 changed files with 18 additions and 45 deletions
Download patch file Download diff file Expand all files Collapse all files

44
litellm/tests/test_router_caching.py
View file

@ -209,47 +209,3 @@ async def test_acompletion_caching_on_router_caching_groups():
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
def test_usage_based_routing_completion():
litellm.set_verbose = True
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo-0301",
"api_key": os.getenv("OPENAI_API_KEY"),
"custom_llm_provider": "Custom-LLM",
},
"tpm": 10000,
"rpm": 5,
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo-0301",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"tpm": 10000,
"rpm": 5,
},
]
router = Router(
model_list=model_list, routing_strategy="usage-based-routing", set_verbose=False
)
max_requests = 5
while max_requests > 0:
try:
router.completion(
model="gpt-3.5-turbo",
messages=[{"content": "write a one sentence poem.", "role": "user"}],
)
except ValueError as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
finally:
max_requests -= 1
router.reset()
test_usage_based_routing_completion()