From 0e2fcde145e9ffb27812badbc314b1c6f2af74d9 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Tue, 9 Dec 2025 18:17:07 -0800 Subject: [PATCH] Change health check to read env varlike chat completions --- .../health_endpoints/_health_endpoints.py | 111 +++++++++++++++- .../health_endpoints/test_health_endpoints.py | 118 ++++++++++++++++++ 2 files changed, 223 insertions(+), 6 deletions(-) diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 79e9838d115..030843376bf 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -30,9 +30,48 @@ from litellm.proxy.health_check import ( perform_health_check, run_with_timeout, ) +from litellm.secret_managers.main import get_secret #### Health ENDPOINTS #### + +def _resolve_os_environ_variables(params: dict) -> dict: + """ + Resolve os.environ/ environment variables in litellm_params. + + This function recursively processes dictionary values that start with "os.environ/" + by replacing them with the actual environment variable values. + + Args: + params: Dictionary containing litellm_params that may have os.environ/ values + + Returns: + Dictionary with os.environ/ values resolved to actual environment variable values + """ + if not isinstance(params, dict): + return params + + resolved_params = {} + for key, value in params.items(): + if isinstance(value, str) and value.startswith("os.environ/"): + # Resolve the environment variable + resolved_value = get_secret(value) + resolved_params[key] = resolved_value + elif isinstance(value, dict): + # Recursively resolve nested dictionaries + resolved_params[key] = _resolve_os_environ_variables(value) + elif isinstance(value, list): + # Handle lists that might contain dictionaries with os.environ/ values + resolved_params[key] = [ + _resolve_os_environ_variables(item) if isinstance(item, dict) else item + for item in value + ] + else: + resolved_params[key] = value + + return resolved_params + + router = APIRouter() services = Union[ Literal[ @@ -1166,21 +1205,41 @@ async def test_model_connection( Example: ```bash + # If model is configured in proxy_config.yaml, you only need to specify the model name: curl -X POST 'http://localhost:4000/health/test_connection' \\ -H 'Authorization: Bearer sk-1234' \\ -H 'Content-Type: application/json' \\ -d '{ "litellm_params": { - "model": "gpt-4", - "custom_llm_provider": "azure_ai", - "litellm_credential_name": null, - "api_key": "6xxxxxxx", - "api_base": "https://litellm8397336933.openai.azure.com/openai/deployments/gpt-4o/chat/completions?api-version=2024-10-21", + "model": "gpt-4o" + }, + "mode": "chat" + }' + + # The endpoint will automatically use api_key, api_base, etc. from proxy_config.yaml + + # You can also override specific params or test with custom credentials: + curl -X POST 'http://localhost:4000/health/test_connection' \\ + -H 'Authorization: Bearer sk-1234' \\ + -H 'Content-Type: application/json' \\ + -d '{ + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "os.environ/AZURE_OPENAI_API_KEY", + "api_base": "os.environ/AZURE_OPENAI_ENDPOINT", + "api_version": "2024-10-21" }, "mode": "chat" }' ``` + Note: + - If the model is configured in proxy_config.yaml, credentials (api_key, api_base, etc.) + will be automatically loaded from the config (with resolved environment variables). + - You can override specific params by including them in the request. + - You can use `os.environ/VARIABLE_NAME` syntax to reference environment variables, + which will be resolved automatically (same as in proxy_config.yaml). + Returns: dict: A dictionary containing the health check result with either success information or error details. """ @@ -1188,7 +1247,7 @@ async def test_model_connection( from litellm.proxy.management_endpoints.model_management_endpoints import ( ModelManagementAuthChecks, ) - from litellm.proxy.proxy_server import premium_user, prisma_client + from litellm.proxy.proxy_server import llm_router, premium_user, prisma_client from litellm.types.router import Deployment, LiteLLM_Params try: @@ -1197,6 +1256,46 @@ async def test_model_connection( status_code=500, detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) + + # Get model name from litellm_params + request_litellm_params = litellm_params or {} + model_name = request_litellm_params.get("model") + + # Look up model configuration from router if model name is provided + # This gets the litellm_params from proxy config (with resolved env vars) + config_litellm_params = {} + if model_name and llm_router is not None: + try: + # First try to find by proxy model_name (e.g., "gpt-4o") + deployments = llm_router.get_model_list(model_name=model_name) + + # If not found, try to find by litellm model name (e.g., "azure/gpt-4o") + if not deployments or len(deployments) == 0: + all_deployments = llm_router.get_model_list(model_name=None) + if all_deployments: + for deployment in all_deployments: + if deployment.get("litellm_params", {}).get("model") == model_name: + deployments = [deployment] + break + + if deployments and len(deployments) > 0: + # Use the first deployment's litellm_params as base config + # These already have resolved environment variables from proxy config + config_litellm_params = deployments[0].get("litellm_params", {}).copy() + except Exception as e: + verbose_proxy_logger.debug( + f"Could not find model {model_name} in router: {e}. " + "Proceeding with request params only." + ) + + # Merge: config params (from proxy config) as base, request params override + # This allows users to override specific params while using config for credentials + merged_litellm_params = {**config_litellm_params, **request_litellm_params} + + # Resolve os.environ/ environment variables in any remaining request params + # This handles cases where user explicitly passes os.environ/ values to override config + litellm_params = _resolve_os_environ_variables(merged_litellm_params) + ## Auth check await ModelManagementAuthChecks.can_user_make_model_call( model_params=Deployment( diff --git a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py index 6939a19b7ef..594a92b0dc7 100644 --- a/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py +++ b/tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py @@ -17,6 +17,7 @@ from litellm.proxy.health_endpoints._health_endpoints import ( _db_health_readiness_check, db_health_cache, health_services_endpoint, + test_model_connection as health_test_model_connection, ) @@ -126,3 +127,120 @@ async def test_health_services_endpoint_sqs(status, error_message): assert result["message"] == error_message mock_instance.async_health_check.assert_awaited_once() + +@pytest.mark.asyncio +async def test_test_model_connection_loads_config_from_router(): + """ + Test that /health/test_connection automatically loads model configuration + (including resolved environment variables) from the router when model name is provided. + """ + # Mock request + mock_request = MagicMock() + + # Mock user_api_key_dict + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.user_id = "test-user" + mock_user_api_key_dict.token = "test-token" + + # Mock prisma_client + mock_prisma_client = MagicMock() + + # Mock router with model configuration + mock_router = MagicMock() + mock_deployment = { + "model_name": "gpt-4o", + "litellm_params": { + "model": "azure/gpt-4o", + "api_key": "resolved-api-key-from-env", + "api_base": "https://resolved-endpoint.openai.azure.com/", + "api_version": "2024-10-21", + }, + "model_info": {}, + } + mock_router.get_model_list.return_value = [mock_deployment] + + # Mock ModelManagementAuthChecks - patch at the source module since it's imported inside the function + mock_can_user_make_model_call = AsyncMock() + + # Mock litellm.ahealth_check + mock_health_check_result = { + "status": "healthy", + "response_time_ms": 100, + } + mock_ahealth_check = AsyncMock(return_value=mock_health_check_result) + + # Mock run_with_timeout + mock_run_with_timeout = AsyncMock(return_value=mock_health_check_result) + + # Mock _update_litellm_params_for_health_check + def mock_update_params(model_info, litellm_params): + # Just return params with messages added + params = litellm_params.copy() + params["messages"] = [{"role": "user", "content": "test"}] + return params + + # Mock _resolve_os_environ_variables + def mock_resolve_os_environ(params): + return params + + with patch( + "litellm.proxy.proxy_server.prisma_client", + mock_prisma_client, + ), patch( + "litellm.proxy.proxy_server.llm_router", + mock_router, + ), patch( + "litellm.proxy.proxy_server.premium_user", + False, + ), patch( + "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", + mock_can_user_make_model_call, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints.litellm.ahealth_check", + mock_ahealth_check, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints.run_with_timeout", + mock_run_with_timeout, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints._update_litellm_params_for_health_check", + mock_update_params, + ), patch( + "litellm.proxy.health_endpoints._health_endpoints._resolve_os_environ_variables", + mock_resolve_os_environ, + ): + # Call the endpoint with only model name (no credentials) + result = await health_test_model_connection( + request=mock_request, + mode="chat", + litellm_params={"model": "gpt-4o"}, + model_info={}, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Verify router.get_model_list was called with the model name + mock_router.get_model_list.assert_called_once_with(model_name="gpt-4o") + + # Verify that run_with_timeout was called (which wraps ahealth_check) + assert mock_run_with_timeout.called + + # Get the call args to verify merged params + call_args = mock_run_with_timeout.call_args + assert call_args is not None + + # The first arg should be the coroutine from ahealth_check + # We need to check what was passed to ahealth_check + ahealth_check_call_args = mock_ahealth_check.call_args + assert ahealth_check_call_args is not None + model_params = ahealth_check_call_args.kwargs.get("model_params", {}) + + # Verify that config params were loaded and merged + # Note: request params override config params, so model from request is used + assert model_params.get("api_key") == "resolved-api-key-from-env" + assert model_params.get("api_base") == "https://resolved-endpoint.openai.azure.com/" + assert model_params.get("api_version") == "2024-10-21" + assert model_params.get("model") == "gpt-4o" # Request param overrides config param + + # Verify result + assert result["status"] == "success" + assert "result" in result +