scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

fix - use llm router for /health checks

Browse source
This commit is contained in:
Ishaan Jaff 2024-04-12 15:37:49 -07:00
parent 8fe8deebde
commit 54633e5cab
2 changed files with 15 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
litellm/proxy/health_check.py
View file

@ -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 = [

11
litellm/proxy/proxy_server.py
View file

@ -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 {