(test) router - test_model_group_aliases

This commit is contained in:
ishaan-jaff 2023-12-14 13:16:44 +05:30
parent 841e941ecd
commit b70bfbb06f

View file

@ -297,7 +297,7 @@ def test_model_group_aliases():
"litellm_params": { "litellm_params": {
"model": "gpt-3.5-turbo-0613", "model": "gpt-3.5-turbo-0613",
"api_key": os.getenv("OPENAI_API_KEY"), "api_key": os.getenv("OPENAI_API_KEY"),
"rpm": 6, "tpm": 1,
}, },
}, },
{ {
@ -307,14 +307,14 @@ def test_model_group_aliases():
"api_key": os.getenv("AZURE_API_KEY"), "api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"), "api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"), "api_version": os.getenv("AZURE_API_VERSION"),
"rpm": 1440, "tpm": 99,
}, },
}, },
{ {
"model_name": "claude-1", "model_name": "claude-1",
"litellm_params": { "litellm_params": {
"model": "bedrock/claude1.2", "model": "bedrock/claude1.2",
"rpm": 1440, "tpm": 1,
}, },
} }
] ]
@ -331,8 +331,23 @@ def test_model_group_aliases():
if selected_model_name != "gpt-3.5-turbo": if selected_model_name != "gpt-3.5-turbo":
pytest.fail(f"Selected model {selected_model_name} is not gpt-3.5-turbo") pytest.fail(f"Selected model {selected_model_name} is not gpt-3.5-turbo")
# test that
# call get_available_deployment 1k times, it should pick azure/chatgpt-v-2 about 90% of the time
selection_counts = defaultdict(int)
for _ in range(1000):
selected_model = router.get_available_deployment("gpt-3.5-turbo")
selected_model_id = selected_model["litellm_params"]["model"]
selected_model_name = litellm.utils.remove_model_id(selected_model_id)
selection_counts[selected_model_name] +=1
print(selection_counts)
total_requests = sum(selection_counts.values())
# Assert that 'azure/chatgpt-v-2' has about 90% of the total requests
assert selection_counts['azure/chatgpt-v-2'] / total_requests > 0.89, f"Assertion failed: 'azure/chatgpt-v-2' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
router.reset() router.reset()
except Exception as e: except Exception as e:
traceback.print_exc() traceback.print_exc()
pytest.fail(f"Error occurred: {e}") pytest.fail(f"Error occurred: {e}")
test_model_group_aliases() # test_model_group_aliases()