phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

refactor(team_endpoints.py): refactor auth checks for team member endpoints to ui team admin to manage it

Browse source
This commit is contained in:
Krrish Dholakia 2024-08-20 16:57:18 -07:00
parent 19083a4d31
commit a61f3e7656
5 changed files with 72 additions and 39 deletions
Download patch file Download diff file Expand all files Collapse all files

46
litellm/proxy/management_endpoints/team_endpoints.py
View file

@ -30,7 +30,7 @@ from litellm.proxy._types import (
UpdateTeamRequest,
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.auth.user_api_key_auth import _is_user_proxy_admin, user_api_key_auth
from litellm.proxy.management_helpers.utils import (
add_new_member,
management_endpoint_wrapper,
@ -39,6 +39,16 @@ from litellm.proxy.management_helpers.utils import (
router = APIRouter()
def _is_user_team_admin(
user_api_key_dict: UserAPIKeyAuth, team_obj: LiteLLM_TeamTable
) -> bool:
for member in team_obj.members_with_roles:
if member.user_id is not None and member.user_id == user_api_key_dict.user_id:
return True
return False
#### TEAM MANAGEMENT ####
@router.post(
"/team/new",
@ -466,6 +476,23 @@ async def team_member_add(
complete_team_data = LiteLLM_TeamTable(**existing_team_row.model_dump())
## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=complete_team_data
)
):
raise HTTPException(
status_code=403,
detail={
"error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format(
"/team/member_add", complete_team_data.team_id
)
},
)
if isinstance(data.member, Member):
# add to team db
new_member = data.member
@ -570,6 +597,23 @@ async def team_member_delete(
)
existing_team_row = LiteLLM_TeamTable(**_existing_team_row.model_dump())
## CHECK IF USER IS PROXY ADMIN OR TEAM ADMIN
if (
user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value
and not _is_user_team_admin(
user_api_key_dict=user_api_key_dict, team_obj=existing_team_row
)
):
raise HTTPException(
status_code=403,
detail={
"error": "Call not allowed. User not proxy admin OR team admin. route={}, team_id={}".format(
"/team/member_delete", existing_team_row.team_id
)
},
)
## DELETE MEMBER FROM TEAM
new_team_members: List[Member] = []
for m in existing_team_row.members_with_roles: