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

Change health check to read env varlike chat completions

Browse source
This commit is contained in:
yuneng-jiang 2025-12-09 18:17:07 -08:00
parent 9a7831dd33
commit 0e2fcde145
2 changed files with 223 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
litellm/proxy/health_endpoints/_health_endpoints.py
View file

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

118
tests/test_litellm/proxy/health_endpoints/test_health_endpoints.py
View file

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