@@ -869,6 +869,10 @@ async def test_model_connection(
869
869
None ,
870
870
description = "Parameters for litellm.completion, litellm.embedding for the health check" ,
871
871
),
872
+ model_info : Dict = fastapi .Body (
873
+ None ,
874
+ description = "Model info for the health check" ,
875
+ ),
872
876
user_api_key_dict : UserAPIKeyAuth = Depends (user_api_key_auth ),
873
877
):
874
878
"""
@@ -897,7 +901,30 @@ async def test_model_connection(
897
901
Returns:
898
902
dict: A dictionary containing the health check result with either success information or error details.
899
903
"""
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
+
900
911
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
+ )
901
928
# Include health_check_params if provided
902
929
litellm_params = _update_litellm_params_for_health_check (
903
930
model_info = {},
@@ -925,11 +952,12 @@ async def test_model_connection(
925
952
"result" : cleaned_result ,
926
953
}
927
954
955
+ except HTTPException as e :
956
+ raise e
928
957
except Exception as e :
929
- verbose_proxy_logger .error (
958
+ verbose_proxy_logger .debug (
930
959
f"litellm.proxy.health_endpoints.test_model_connection(): Exception occurred - { str (e )} "
931
960
)
932
- verbose_proxy_logger .debug (traceback .format_exc ())
933
961
raise HTTPException (
934
962
status_code = status .HTTP_500_INTERNAL_SERVER_ERROR ,
935
963
detail = {"error" : f"Failed to test connection: { str (e )} " },
0 commit comments