fix: fix linting errors

This commit is contained in:
Krrish Dholakia 2024-11-23 22:14:57 +05:30
parent fbf54900b1
commit cebfcfc2d7

View file

@ -179,77 +179,6 @@ 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()