Skip to content

Commit 32c6019

Browse files
fix(_health_endpoints.py): protect /health/test_connection - only allow users who are allowed to create models, to call this endpoint
Closes LIT-989
1 parent fc18f4d commit 32c6019

File tree

3 files changed

+30
-3
lines changed

3 files changed

+30
-3
lines changed
File renamed without changes.

litellm/proxy/_experimental/out/onboarding.html

Lines changed: 0 additions & 1 deletion
This file was deleted.

litellm/proxy/health_endpoints/_health_endpoints.py

Lines changed: 30 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -869,6 +869,10 @@ async def test_model_connection(
869869
None,
870870
description="Parameters for litellm.completion, litellm.embedding for the health check",
871871
),
872+
model_info: Dict = fastapi.Body(
873+
None,
874+
description="Model info for the health check",
875+
),
872876
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
873877
):
874878
"""
@@ -897,7 +901,30 @@ async def test_model_connection(
897901
Returns:
898902
dict: A dictionary containing the health check result with either success information or error details.
899903
"""
904+
from litellm.proxy._types import CommonProxyErrors
905+
from litellm.proxy.management_endpoints.model_management_endpoints import (
906+
ModelManagementAuthChecks,
907+
)
908+
from litellm.proxy.proxy_server import premium_user, prisma_client
909+
from litellm.types.router import Deployment, LiteLLM_Params
910+
900911
try:
912+
if prisma_client is None:
913+
raise HTTPException(
914+
status_code=500,
915+
detail={"error": CommonProxyErrors.db_not_connected_error.value},
916+
)
917+
## Auth check
918+
await ModelManagementAuthChecks.can_user_make_model_call(
919+
model_params=Deployment(
920+
model_name="test_model",
921+
litellm_params=LiteLLM_Params(**litellm_params),
922+
model_info=model_info,
923+
),
924+
user_api_key_dict=user_api_key_dict,
925+
prisma_client=prisma_client,
926+
premium_user=premium_user,
927+
)
901928
# Include health_check_params if provided
902929
litellm_params = _update_litellm_params_for_health_check(
903930
model_info={},
@@ -925,11 +952,12 @@ async def test_model_connection(
925952
"result": cleaned_result,
926953
}
927954

955+
except HTTPException as e:
956+
raise e
928957
except Exception as e:
929-
verbose_proxy_logger.error(
958+
verbose_proxy_logger.debug(
930959
f"litellm.proxy.health_endpoints.test_model_connection(): Exception occurred - {str(e)}"
931960
)
932-
verbose_proxy_logger.debug(traceback.format_exc())
933961
raise HTTPException(
934962
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
935963
detail={"error": f"Failed to test connection: {str(e)}"},

0 commit comments

Comments
 (0)