fix(proxy_server.py): fix team member delete endpoint

fixes issue in for loop, removing member from team

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(

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