From e23c68b15a622f72c0799e68355546826c041c6e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 11 Mar 2024 12:14:13 -0700 Subject: [PATCH] (fix) failing usage based routing test --- litellm/router_strategy/lowest_tpm_rpm.py | 8 +++++--- litellm/tests/test_router_get_deployments.py | 4 ++-- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 9f0b9eb22..b2f9d6e4e 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -200,9 +200,11 @@ class LowestTPMLoggingHandler(CustomLogger): if item_tpm == 0: deployment = _deployment break - elif item_tpm + input_tokens > _deployment_tpm or ( - item in rpm_dict and rpm_dict[item] + 1 > _deployment_rpm - ): # if user passed in tpm / rpm in the model_list + elif item_tpm + input_tokens > _deployment_tpm: + continue + elif (rpm_dict is not None and item in rpm_dict) and ( + rpm_dict[item] + 1 > _deployment_rpm + ): continue elif item_tpm < lowest_tpm: lowest_tpm = item_tpm diff --git a/litellm/tests/test_router_get_deployments.py b/litellm/tests/test_router_get_deployments.py index 62630d7e7..7fc871743 100644 --- a/litellm/tests/test_router_get_deployments.py +++ b/litellm/tests/test_router_get_deployments.py @@ -429,11 +429,11 @@ def test_usage_based_routing(): mock_response="good morning", ) - # print(response) + # print("response", response) selection_counts[response["model"]] += 1 - print(selection_counts) + # print("selection counts", selection_counts) total_requests = sum(selection_counts.values())