diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6ceb7df48..8866781d6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10283,36 +10283,39 @@ async def team_member_delete( detail={"error": "Either user_id or user_email needs to be passed in"}, ) - existing_team_row = await prisma_client.get_data( # type: ignore - team_id=data.team_id, table_name="team", query_type="find_unique" + _existing_team_row = await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": data.team_id} ) + if _existing_team_row is None: + raise HTTPException( + status_code=400, + detail={"error": "Team id={} does not exist in db".format(data.team_id)}, + ) + existing_team_row = LiteLLM_TeamTable(**_existing_team_row.model_dump()) + ## DELETE MEMBER FROM TEAM - new_team_members = [] + new_team_members: List[dict] = [] for m in existing_team_row.members_with_roles: if ( data.user_id is not None - and m["user_id"] is not None - and data.user_id == m["user_id"] + and m.user_id is not None + and data.user_id == m.user_id ): continue elif ( data.user_email is not None - and m["user_email"] is not None - and data.user_email == m["user_email"] + and m.user_email is not None + and data.user_email == m.user_email ): continue - new_team_members.append(m) - existing_team_row.members_with_roles = new_team_members - complete_team_data = LiteLLM_TeamTable( - **existing_team_row.model_dump(), - ) + new_team_members.append(m.model_dump()) - team_row = await prisma_client.update_data( - update_key_values=complete_team_data.json(exclude_none=True), - data=complete_team_data.json(exclude_none=True), - table_name="team", - team_id=data.team_id, + _ = await prisma_client.db.litellm_teamtable.update( + where={ + "team_id": data.team_id, + }, + data={"members_with_roles": json.dumps(new_team_members)}, # type: ignore ) ## DELETE TEAM ID from USER ROW, IF EXISTS ## @@ -10322,36 +10325,26 @@ async def team_member_delete( key_val["user_id"] = data.user_id elif data.user_email is not None: key_val["user_email"] = data.user_email - existing_user_rows = await prisma_client.get_data( - key_val=key_val, - table_name="user", - query_type="find_all", + existing_user_rows = await prisma_client.db.litellm_usertable.find_many( + where=key_val # type: ignore ) - user_data = { # type: ignore - "teams": [], - "models": team_row["data"].models, - } + if existing_user_rows is not None and ( isinstance(existing_user_rows, list) and len(existing_user_rows) > 0 ): for existing_user in existing_user_rows: team_list = [] - if hasattr(existing_user, "teams"): + if data.team_id in existing_user.teams: team_list = existing_user.teams team_list.remove(data.team_id) - user_data["user_id"] = existing_user.user_id - await prisma_client.update_data( - user_id=existing_user.user_id, - data=user_data, - update_key_values_custom_query={ - "teams": { - "set": [team_row["team_id"]], - } + await prisma_client.db.litellm_usertable.update( + where={ + "user_id": existing_user.user_id, }, - table_name="user", + data={"teams": {"set": team_list}}, ) - return team_row["data"] + return existing_team_row @router.post( diff --git a/tests/test_team.py b/tests/test_team.py index 17491b396..3f7ed71b5 100644 --- a/tests/test_team.py +++ b/tests/test_team.py @@ -15,6 +15,7 @@ async def new_user( budget_duration=None, models=["azure-models"], team_id=None, + user_email=None, ): url = "http://0.0.0.0:4000/user/new" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} @@ -24,6 +25,7 @@ async def new_user( "duration": None, "max_budget": budget, "budget_duration": budget_duration, + "user_email": user_email, } if user_id is not None: @@ -74,10 +76,14 @@ async def add_member( return await response.json() -async def delete_member(session, i, team_id, user_id): +async def delete_member(session, i, team_id, user_id=None, user_email=None): url = "http://0.0.0.0:4000/team/member_delete" headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"} - data = {"team_id": team_id, "user_id": user_id} + data = {"team_id": team_id} + if user_id is not None: + data["user_id"] = user_id + elif user_email is not None: + data["user_email"] = user_email async with session.post(url, headers=headers, json=data) as response: status = response.status @@ -415,8 +421,9 @@ async def test_team_delete(): await delete_team(session=session, i=0, team_id=team_data["team_id"]) +@pytest.mark.parametrize("dimension", ["user_id", "user_email"]) @pytest.mark.asyncio -async def test_member_delete(): +async def test_member_delete(dimension): """ - Create team - Add member @@ -431,30 +438,58 @@ async def test_member_delete(): await new_user(session=session, i=0, user_id=admin_user) ## Create normal user normal_user = f"{uuid.uuid4()}" + normal_user_email = "{}@berri.ai".format(normal_user) print(f"normal_user: {normal_user}") - await new_user(session=session, i=0, user_id=normal_user) + await new_user( + session=session, i=0, user_id=normal_user, user_email=normal_user_email + ) ## Create team with 1 admin and 1 user member_list = [ {"role": "admin", "user_id": admin_user}, - {"role": "user", "user_id": normal_user}, ] + if dimension == "user_id": + member_list.append({"role": "user", "user_id": normal_user}) + elif dimension == "user_email": + member_list.append({"role": "user", "user_email": normal_user_email}) team_data = await new_team(session=session, i=0, member_list=member_list) - print(f"team_data: {team_data}") - member_id_list = [] + + user_in_team = False for member in team_data["members_with_roles"]: - member_id_list.append(member["user_id"]) + if dimension == "user_id" and member["user_id"] == normal_user: + user_in_team = True + elif ( + dimension == "user_email" and member["user_email"] == normal_user_email + ): + user_in_team = True - assert normal_user in member_id_list - # Delete member - updated_team_data = await delete_member( - session=session, i=0, team_id=team_data["team_id"], user_id=normal_user + assert ( + user_in_team is True + ), "User not in team. Team list={}, User details - id={}, email={}. Dimension={}".format( + team_data["members_with_roles"], normal_user, normal_user_email, dimension ) + # Delete member + if dimension == "user_id": + updated_team_data = await delete_member( + session=session, i=0, team_id=team_data["team_id"], user_id=normal_user + ) + elif dimension == "user_email": + updated_team_data = await delete_member( + session=session, + i=0, + team_id=team_data["team_id"], + user_email=normal_user_email, + ) print(f"updated_team_data: {updated_team_data}") - member_id_list = [] - for member in updated_team_data["members_with_roles"]: - member_id_list.append(member["user_id"]) + user_in_team = False + for member in team_data["members_with_roles"]: + if dimension == "user_id" and member["user_id"] == normal_user: + user_in_team = True + elif ( + dimension == "user_email" and member["user_email"] == normal_user_email + ): + user_in_team = True - assert normal_user not in member_id_list + assert user_in_team is True @pytest.mark.asyncio