diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 00256ed87..7005f88fc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -458,8 +458,16 @@ class TeamMemberDeleteRequest(LiteLLMBase): return values -class UpdateTeamRequest(TeamBase): +class UpdateTeamRequest(LiteLLMBase): team_id: str # required + team_alias: Optional[str] = None + organization_id: Optional[str] = None + metadata: Optional[dict] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + max_budget: Optional[float] = None + models: Optional[list] = None + blocked: Optional[bool] = None class DeleteTeamRequest(LiteLLMBase): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 06a51a384..f6a8f54f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6310,7 +6310,7 @@ async def update_team( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): """ - [RECOMMENDED] - use `/team/member_add` to add new team members instead + Use `/team/member_add` AND `/team/member/delete` to add/remove new team members You can now update team budget / rate limits via /team/update @@ -6360,60 +6360,6 @@ async def update_team( team_id=data.team_id, ) - ## ADD NEW USERS ## - existing_user_id_list = [] - ## Get new users - for user in existing_team_row.members_with_roles: - if user["user_id"] is not None: - existing_user_id_list.append(user["user_id"]) - - ## Update new user rows with team id (info used by /user/info to show all teams, user is a part of) - if data.members_with_roles is not None: - for user in data.members_with_roles: - if user.user_id not in existing_user_id_list: - await prisma_client.update_data( - user_id=user.user_id, - data={ - "user_id": user.user_id, - "teams": [team_row["team_id"]], - "models": team_row["data"].models, - }, - update_key_values_custom_query={ - "teams": { - "push": [team_row["team_id"]], - } - }, - table_name="user", - ) - - ## REMOVE DELETED USERS ## - ### Get list of deleted users (old list - new list) - deleted_user_id_list = [] - new_user_id_list = [] - ## Get old user list - if data.members_with_roles is not None: - for user in data.members_with_roles: - new_user_id_list.append(user.user_id) - ## Get diff - if existing_team_row.members_with_roles is not None: - for user in existing_team_row.members_with_roles: - if user["user_id"] not in new_user_id_list and len(new_user_id_list) > 0: - deleted_user_id_list.append(user["user_id"]) - - ## SET UPDATED LIST - if len(deleted_user_id_list) > 0: - # get the deleted users - existing_user_rows = await prisma_client.get_data( - user_id_list=deleted_user_id_list, table_name="user", query_type="find_all" - ) - for user in existing_user_rows: - if data.team_id in user["teams"]: - user["teams"].remove(data.team_id) - await prisma_client.update_data( - user_id=user["user_id"], - data=user, - update_key_values={"user_id": user["user_id"], "teams": user["teams"]}, - ) return team_row diff --git a/tests/test_team.py b/tests/test_team.py index 37ee69dd5..7bff7b36b 100644 --- a/tests/test_team.py +++ b/tests/test_team.py @@ -37,6 +37,29 @@ async def new_user( return await response.json() +async def add_member(session, i, team_id, user_id=None, user_email=None): + url = "http://0.0.0.0:4000/team/member_add" + headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} + data = {"team_id": team_id, "member": {"role": "user"}} + if user_email is not None: + data["member"]["user_email"] = user_email + elif user_id is not None: + data["member"]["user_id"] = user_id + + async with session.post(url, headers=headers, json=data) as response: + status = response.status + response_text = await response.text() + + print(f"ADD MEMBER Response {i} (Status code: {status}):") + print(response_text) + print() + + if status != 200: + raise Exception(f"Request {i} did not return a 200 status code: {status}") + + return await response.json() + + async def delete_member(session, i, team_id, user_id): url = "http://0.0.0.0:4000/team/member_delete" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} @@ -247,43 +270,75 @@ async def test_team_info(): await get_team_info(session=session, get_team=team_id, call_key="sk-1234") +""" +- Create team +- Add user (user exists in db) +- Update team +- Check if it works +""" + +""" +- Create team +- Add user (user doesn't exist in db) +- Update team +- Check if it works +""" + + @pytest.mark.asyncio -async def test_team_update(): +async def test_team_update_sc_2(): """ - - Create team with 1 admin, 1 user - - Create new user - - Replace existing user with new user in team + - Create team + - Add 1 user (doesn't exist in db) + - Change team alias + - Check if it works + - Assert team object unchanged besides team alias """ async with aiohttp.ClientSession() as session: ## Create admin admin_user = f"{uuid.uuid4()}" await new_user(session=session, i=0, user_id=admin_user) - ## Create normal user - normal_user = f"{uuid.uuid4()}" - await new_user(session=session, i=0, user_id=normal_user) ## Create team with 1 admin and 1 user member_list = [ {"role": "admin", "user_id": admin_user}, - {"role": "user", "user_id": normal_user}, ] team_data = await new_team(session=session, i=0, member_list=member_list) ## Create new normal user - new_normal_user = f"{uuid.uuid4()}" - await new_user(session=session, i=0, user_id=new_normal_user) - ## Update member list - member_list = [ - {"role": "admin", "user_id": admin_user}, - {"role": "user", "user_id": new_normal_user}, - ] - team_data = await update_team( + new_normal_user = f"krrish_{uuid.uuid4()}@berri.ai" + await add_member( session=session, i=0, - member_list=member_list, team_id=team_data["team_id"], - tpm_limit=100, + user_id=None, + user_email=new_normal_user, ) - assert team_data["data"]["tpm_limit"] == 100 + ## CHANGE TEAM ALIAS + + new_team_data = await update_team( + session=session, i=0, team_id=team_data["team_id"], team_alias="my-new-team" + ) + + assert new_team_data["data"]["team_alias"] == "my-new-team" + print(f"team_data: {team_data}") + ## assert rest of object is the same + for k, v in new_team_data["data"].items(): + if ( + k == "members_with_roles" + ): # assert 1 more member (role: "user", user_email: $user_email) + len(new_team_data["data"][k]) == len(team_data[k]) + 1 + elif ( + k == "created_at" + or k == "updated_at" + or k == "model_spend" + or k == "model_max_budget" + or k == "model_id" + or k == "litellm_organization_table" + or k == "litellm_model_table" + ): + pass + else: + assert new_team_data["data"][k] == team_data[k] @pytest.mark.asyncio