From 54633e5cab2f875aa9e2382fe8c5c0f974909845 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Fri, 12 Apr 2024 15:37:49 -0700
Subject: [PATCH] fix - use llm router for /health checks

---
 litellm/proxy/health_check.py |  9 ++++++++-
 litellm/proxy/proxy_server.py | 11 +++++++----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py
index b05bd4b6ab5..6815ef4cb8e 100644
--- a/litellm/proxy/health_check.py
+++ b/litellm/proxy/health_check.py
@@ -7,6 +7,7 @@ from typing import Optional
 import litellm
 import logging
 from litellm._logging import print_verbose
+from litellm import Router
 
 
 logger = logging.getLogger(__name__)
@@ -69,7 +70,9 @@ async def _perform_health_check(model_list: list):
 
 
 async def perform_health_check(
-    model_list: list, model: Optional[str] = None, cli_model: Optional[str] = None
+    litellm_router_instance: Optional[Router] = None,
+    model: Optional[str] = None,
+    cli_model: Optional[str] = None,
 ):
     """
     Perform a health check on the system.
@@ -77,6 +80,10 @@ async def perform_health_check(
     Returns:
         (bool): True if the health check passes, False otherwise.
     """
+    model_list = None
+    if litellm_router_instance is not None:
+        model_list = litellm_router_instance.get_model_list()
+
     if not model_list:
         if cli_model:
             model_list = [
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 17ddb2f0535..73b17328a0e 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -8329,13 +8329,13 @@ async def health_endpoint(
     ```
     else, the health checks will be run on models when /health is called.
     """
-    global health_check_results, use_background_health_checks, user_model
+    global health_check_results, use_background_health_checks, user_model, llm_router
     try:
-        if llm_model_list is None:
+        if llm_router is None:
             # if no router set, check if user set a model using litellm --model ollama/llama2
             if user_model is not None:
                 healthy_endpoints, unhealthy_endpoints = await perform_health_check(
-                    model_list=[], cli_model=user_model
+                    litellm_router_instance=None, cli_model=user_model
                 )
                 return {
                     "healthy_endpoints": healthy_endpoints,
@@ -8356,8 +8356,11 @@ async def health_endpoint(
         if use_background_health_checks:
             return health_check_results
         else:
+            verbose_proxy_logger.debug(
+                "about to make a health check - llm_model_list= %s", llm_model_list
+            )
             healthy_endpoints, unhealthy_endpoints = await perform_health_check(
-                llm_model_list, model
+                llm_router, model
             )
 
             return {