diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 2f2f6b2dd..7ff209094 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -12,12 +12,3 @@ model_list: vertex_ai_project: "adroit-crow-413218" vertex_ai_location: "us-east5" -litellm_settings: - success_callback: ["langfuse"] - callbacks: ["prometheus"] - key_generation_settings: - team_key_generation: - allowed_team_member_roles: ["admin"] - required_params: ["tags"] - personal_key_generation: # maps to 'Default Team' on UI - allowed_user_roles: ["proxy_admin"] \ No newline at end of file diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 511e5a940..7a2aa9de5 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -179,6 +179,77 @@ def key_generation_check( ) + +def _is_team_key(data: GenerateKeyRequest): + return data.team_id is not None + + +def _team_key_generation_check(user_api_key_dict: UserAPIKeyAuth): + if ( + litellm.key_generation_settings is None + or litellm.key_generation_settings.get("team_key_generation") is None + ): + return True + + if user_api_key_dict.team_member is None: + raise HTTPException( + status_code=400, + detail=f"User not assigned to team. Got team_member={user_api_key_dict.team_member}", + ) + + team_member_role = user_api_key_dict.team_member.role + if ( + team_member_role + not in litellm.key_generation_settings["team_key_generation"][ # type: ignore + "allowed_team_member_roles" + ] + ): + raise HTTPException( + status_code=400, + detail=f"Team member role {team_member_role} not in allowed_team_member_roles={litellm.key_generation_settings['team_key_generation']['allowed_team_member_roles']}", # type: ignore + ) + return True + + +def _personal_key_generation_check(user_api_key_dict: UserAPIKeyAuth): + + if ( + litellm.key_generation_settings is None + or litellm.key_generation_settings.get("personal_key_generation") is None + ): + return True + + if ( + user_api_key_dict.user_role + not in litellm.key_generation_settings["personal_key_generation"][ # type: ignore + "allowed_user_roles" + ] + ): + raise HTTPException( + status_code=400, + detail=f"Personal key creation has been restricted by admin. Allowed roles={litellm.key_generation_settings['personal_key_generation']['allowed_user_roles']}. Your role={user_api_key_dict.user_role}", # type: ignore + ) + return True + + +def key_generation_check( + user_api_key_dict: UserAPIKeyAuth, data: GenerateKeyRequest +) -> bool: + """ + Check if admin has restricted key creation to certain roles for teams or individuals + """ + if litellm.key_generation_settings is None: + return True + + ## check if key is for team or individual + is_team_key = _is_team_key(data=data) + + if is_team_key: + return _team_key_generation_check(user_api_key_dict) + else: + return _personal_key_generation_check(user_api_key_dict=user_api_key_dict) + + router = APIRouter()