From fff148a8353b7bdf38c58a6f8297ad17ebab7e5c Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 15 Dec 2023 07:07:25 +0530 Subject: [PATCH] (test) router_get_deployments - stop using remove_model_id --- litellm/proxy/health_check.py | 1 - litellm/tests/test_router_get_deployments.py | 12 ++++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index db07ae20a7f..f8e56c05996 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -69,7 +69,6 @@ async def _perform_health_check(model_list: list): for model in model_list: litellm_params = model["litellm_params"] model_info = model.get("model_info", {}) - litellm_params["model"] = litellm.utils.remove_model_id(litellm_params["model"]) litellm_params["messages"] = _get_random_llm_message() prepped_params.append(litellm_params) diff --git a/litellm/tests/test_router_get_deployments.py b/litellm/tests/test_router_get_deployments.py index d9433645886..a71fc3823d7 100644 --- a/litellm/tests/test_router_get_deployments.py +++ b/litellm/tests/test_router_get_deployments.py @@ -49,7 +49,7 @@ def test_weighted_selection_router(): 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) + selected_model_name = selected_model_id selection_counts[selected_model_name] +=1 print(selection_counts) @@ -101,7 +101,7 @@ def test_weighted_selection_router_tpm(): 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) + selected_model_name = selected_model_id selection_counts[selected_model_name] +=1 print(selection_counts) @@ -154,7 +154,7 @@ def test_weighted_selection_router_tpm_as_router_param(): 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) + selected_model_name = selected_model_id selection_counts[selected_model_name] +=1 print(selection_counts) @@ -210,7 +210,7 @@ def test_weighted_selection_router_rpm_as_router_param(): 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) + selected_model_name = selected_model_id selection_counts[selected_model_name] +=1 print(selection_counts) @@ -270,7 +270,7 @@ def test_weighted_selection_router_no_rpm_set(): for _ in range(1000): selected_model = router.get_available_deployment("claude-1") selected_model_id = selected_model["litellm_params"]["model"] - selected_model_name = litellm.utils.remove_model_id(selected_model_id) + selected_model_name = selected_model_id selection_counts[selected_model_name] +=1 print(selection_counts) @@ -337,7 +337,7 @@ def test_model_group_aliases(): 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) + selected_model_name = selected_model_id selection_counts[selected_model_name] +=1 print(selection_counts)