From 164898a213a7e1c6e80f4553a813a75bd2808ae3 Mon Sep 17 00:00:00 2001
From: CLARKBENHAM <clark.benham@gmail.com>
Date: Tue, 2 Apr 2024 19:10:40 -0700
Subject: [PATCH] fix lowest latency tests

---
 litellm/router_strategy/lowest_latency.py    | 25 ++++++----------
 litellm/tests/test_lowest_latency_routing.py | 30 ++++++--------------
 2 files changed, 16 insertions(+), 39 deletions(-)

diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py
index 57b56e87f88..493b25eda39 100644
--- a/litellm/router_strategy/lowest_latency.py
+++ b/litellm/router_strategy/lowest_latency.py
@@ -35,9 +35,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
     logged_success: int = 0
     logged_failure: int = 0
 
-    def __init__(
-        self, router_cache: DualCache, model_list: list, routing_args: dict = {}
-    ):
+    def __init__(self, router_cache: DualCache, model_list: list, routing_args: dict = {}):
         self.router_cache = router_cache
         self.model_list = model_list
         self.routing_args = RoutingArgs(**routing_args)
@@ -50,9 +48,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
             else:
-                model_group = kwargs["litellm_params"]["metadata"].get(
-                    "model_group", None
-                )
+                model_group = kwargs["litellm_params"]["metadata"].get("model_group", None)
 
                 id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
                 if model_group is None or id is None:
@@ -65,8 +61,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 {
                     {model_group}_map: {
                         id: {
-                            "latency": [..]  
-                            f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}      
+                            "latency": [..]
+                            f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
                         }
                     }
                 }
@@ -132,9 +128,7 @@ class LowestLatencyLoggingHandler(CustomLogger):
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
             else:
-                model_group = kwargs["litellm_params"]["metadata"].get(
-                    "model_group", None
-                )
+                model_group = kwargs["litellm_params"]["metadata"].get("model_group", None)
 
                 id = kwargs["litellm_params"].get("model_info", {}).get("id", None)
                 if model_group is None or id is None:
@@ -147,8 +141,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 {
                     {model_group}_map: {
                         id: {
-                            "latency": [..]  
-                            f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}      
+                            "latency": [..]
+                            f"{date:hour:minute}" : {"tpm": 34, "rpm": 3}
                         }
                     }
                 }
@@ -287,13 +281,10 @@ class LowestLatencyLoggingHandler(CustomLogger):
                 deployment = _deployment
                 break
             elif (
-                item_tpm + input_tokens > _deployment_tpm
-                or item_rpm + 1 > _deployment_rpm
+                item_tpm + input_tokens > _deployment_tpm or item_rpm + 1 > _deployment_rpm
             ):  # if user passed in tpm / rpm in the model_list
                 continue
             elif item_latency < lowest_latency:
                 lowest_latency = item_latency
                 deployment = _deployment
-        if deployment is None:
-            deployment = random.choice(healthy_deployments)
         return deployment
diff --git a/litellm/tests/test_lowest_latency_routing.py b/litellm/tests/test_lowest_latency_routing.py
index b07132f9e52..13627b71363 100644
--- a/litellm/tests/test_lowest_latency_routing.py
+++ b/litellm/tests/test_lowest_latency_routing.py
@@ -230,19 +230,11 @@ def test_get_available_endpoints_tpm_rpm_check_async(ans_rpm):
     d2 = [(lowest_latency_logger, "5678", 50, 0.01)] * non_ans_rpm
     asyncio.run(_gather_deploy([*d1, *d2]))
     ## CHECK WHAT'S SELECTED ##
-    print(dir(lowest_latency_logger))
-    print(
-        "availible",
-        lowest_latency_logger.get_available_deployments(
-            model_group=model_group, healthy_deployments=model_list
-        ),
-    )
-    assert (
-        lowest_latency_logger.get_available_deployments(
-            model_group=model_group, healthy_deployments=model_list
-        )["model_info"]["id"]
-        == ans
+    d_ans = lowest_latency_logger.get_available_deployments(
+        model_group=model_group, healthy_deployments=model_list
     )
+    print(d_ans)
+    assert (d_ans and d_ans["model_info"]["id"]) == ans
 
 
 # test_get_available_endpoints_tpm_rpm_check_async()
@@ -325,17 +317,11 @@ def test_get_available_endpoints_tpm_rpm_check(ans_rpm):
         )
 
     ## CHECK WHAT'S SELECTED ##
-    print(
-        lowest_latency_logger.get_available_deployments(
-            model_group=model_group, healthy_deployments=model_list
-        )
-    )
-    assert (
-        lowest_latency_logger.get_available_deployments(
-            model_group=model_group, healthy_deployments=model_list
-        )["model_info"]["id"]
-        == ans
+    d_ans = lowest_latency_logger.get_available_deployments(
+        model_group=model_group, healthy_deployments=model_list
     )
+    print(d_ans)
+    assert (d_ans and d_ans["model_info"]["id"]) == ans
 
 
 def test_router_get_available_deployments():