From f8bec3a86c14ff4fda7a31508a64c6953b998428 Mon Sep 17 00:00:00 2001 From: Florian Greinacher Date: Thu, 18 Jul 2024 17:17:04 +0200 Subject: [PATCH] feat(proxy): support hiding health check details --- docs/my-website/docs/proxy/health.md | 14 ++++++++++- litellm/proxy/health_check.py | 25 +++++++++++-------- .../health_endpoints/_health_endpoints.py | 5 ++-- litellm/proxy/proxy_server.py | 11 +++++--- 4 files changed, 39 insertions(+), 16 deletions(-) diff --git a/docs/my-website/docs/proxy/health.md b/docs/my-website/docs/proxy/health.md index 1e2d4945bc4..6d383fc4161 100644 --- a/docs/my-website/docs/proxy/health.md +++ b/docs/my-website/docs/proxy/health.md @@ -124,6 +124,18 @@ model_list: mode: audio_transcription ``` +### Hide details + +The health check response contains details like endpoint URLs, error messages, +and other LiteLLM params. While this is useful for debugging, it can be +problematic when exposing the proxy server to a broad audience. + +You can hide these details by setting the `health_check_details` setting to `False`. + +```yaml +general_settings: + health_check_details: False +``` ## `/health/readiness` @@ -218,4 +230,4 @@ curl -X POST 'http://localhost:4000/chat/completions' \ ], } ' -``` \ No newline at end of file +``` diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index a20ec06e588..aa6205c7c3f 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -14,6 +14,7 @@ logger = logging.getLogger(__name__) ILLEGAL_DISPLAY_PARAMS = ["messages", "api_key", "prompt", "input"] +MINIMAL_DISPLAY_PARAMS = ["model"] def _get_random_llm_message(): """ @@ -24,14 +25,18 @@ def _get_random_llm_message(): return [{"role": "user", "content": random.choice(messages)}] -def _clean_litellm_params(litellm_params: dict): +def _clean_endpoint_data(endpoint_data: dict, details: bool): """ - Clean the litellm params for display to users. + Clean the endpoint data for display to users. """ - return {k: v for k, v in litellm_params.items() if k not in ILLEGAL_DISPLAY_PARAMS} + return ( + {k: v for k, v in endpoint_data.items() if k not in ILLEGAL_DISPLAY_PARAMS} + if details + else {k: v for k, v in endpoint_data.items() if k in MINIMAL_DISPLAY_PARAMS} + ) -async def _perform_health_check(model_list: list): +async def _perform_health_check(model_list: list, details: bool): """ Perform a health check for each model in the list. """ @@ -56,20 +61,20 @@ async def _perform_health_check(model_list: list): unhealthy_endpoints = [] for is_healthy, model in zip(results, model_list): - cleaned_litellm_params = _clean_litellm_params(model["litellm_params"]) + litellm_params = model["litellm_params"] if isinstance(is_healthy, dict) and "error" not in is_healthy: - healthy_endpoints.append({**cleaned_litellm_params, **is_healthy}) + healthy_endpoints.append(_clean_endpoint_data({**litellm_params, **is_healthy}, details)) elif isinstance(is_healthy, dict): - unhealthy_endpoints.append({**cleaned_litellm_params, **is_healthy}) + unhealthy_endpoints.append(_clean_endpoint_data({**litellm_params, **is_healthy}, details)) else: - unhealthy_endpoints.append(cleaned_litellm_params) + unhealthy_endpoints.append(_clean_endpoint_data(litellm_params, details)) return healthy_endpoints, unhealthy_endpoints async def perform_health_check( - model_list: list, model: Optional[str] = None, cli_model: Optional[str] = None + model_list: list, model: Optional[str] = None, cli_model: Optional[str] = None, details: Optional[bool] = True ): """ Perform a health check on the system. @@ -93,6 +98,6 @@ async def perform_health_check( _new_model_list = [x for x in model_list if x["model_name"] == model] model_list = _new_model_list - healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list) + healthy_endpoints, unhealthy_endpoints = await _perform_health_check(model_list, details) return healthy_endpoints, unhealthy_endpoints diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index e5ba03aacda..494d9aa0991 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -287,6 +287,7 @@ async def health_endpoint( llm_model_list, use_background_health_checks, user_model, + health_check_details ) try: @@ -294,7 +295,7 @@ async def health_endpoint( # 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 + model_list=[], cli_model=user_model, details=health_check_details ) return { "healthy_endpoints": healthy_endpoints, @@ -316,7 +317,7 @@ async def health_endpoint( return health_check_results else: healthy_endpoints, unhealthy_endpoints = await perform_health_check( - _llm_model_list, model + _llm_model_list, model, details=health_check_details ) return { diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9dc735d4619..5fe9289f49d 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -416,6 +416,7 @@ user_custom_key_generate = None use_background_health_checks = None use_queue = False health_check_interval = None +health_check_details = None health_check_results = {} queue: List = [] litellm_proxy_budget_name = "litellm-proxy-budget" @@ -1204,14 +1205,14 @@ async def _run_background_health_check(): Update health_check_results, based on this. """ - global health_check_results, llm_model_list, health_check_interval + global health_check_results, llm_model_list, health_check_interval, health_check_details # make 1 deep copy of llm_model_list -> use this for all background health checks _llm_model_list = copy.deepcopy(llm_model_list) while True: healthy_endpoints, unhealthy_endpoints = await perform_health_check( - model_list=_llm_model_list + model_list=_llm_model_list, details=health_check_details ) # Update the global variable with the health check results @@ -1363,7 +1364,7 @@ class ProxyConfig: """ Load config values into proxy global state """ - global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger + global master_key, user_config_file_path, otel_logging, user_custom_auth, user_custom_auth_path, user_custom_key_generate, use_background_health_checks, health_check_interval, use_queue, custom_db_client, proxy_budget_rescheduler_max_time, proxy_budget_rescheduler_min_time, ui_access_mode, litellm_master_key_hash, proxy_batch_write_at, disable_spend_logs, prompt_injection_detection_obj, redis_usage_cache, store_model_in_db, premium_user, open_telemetry_logger, health_check_details # Load existing config config = await self.get_config(config_file_path=config_file_path) @@ -1733,6 +1734,9 @@ class ProxyConfig: "background_health_checks", False ) health_check_interval = general_settings.get("health_check_interval", 300) + health_check_details = general_settings.get( + "health_check_details", True + ) ## check if user has set a premium feature in general_settings if ( @@ -9418,6 +9422,7 @@ def cleanup_router_config_variables(): user_custom_key_generate = None use_background_health_checks = None health_check_interval = None + health_check_details = None prisma_client = None custom_db_client = None