phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(team_endpoints.py): return team member budgets in /team/info call

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/5390
This commit is contained in:
Krrish Dholakia 2024-08-28 19:14:01 -07:00
parent f0fb8bdf45
commit 18689c25e9
5 changed files with 53 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -499,28 +499,10 @@ async def team_member_add(
},
)
if isinstance(data.member, Member):
# add to team db
new_member = data.member
complete_team_data.members_with_roles.append(new_member)
elif isinstance(data.member, List):
# add to team db
new_members = data.member
complete_team_data.members_with_roles.extend(new_members)
# ADD MEMBER TO TEAM
_db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles]
updated_team = await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore
)
updated_users: List[LiteLLM_UserTable] = []
updated_team_memberships: List[LiteLLM_TeamMembership] = []
## VALIDATE IF NEW MEMBER ##
if isinstance(data.member, Member):
try:
updated_user, updated_tm = await add_new_member(
@ -571,6 +553,41 @@ async def team_member_add(
await asyncio.gather(*tasks)
## ADD TO TEAM ##
if isinstance(data.member, Member):
# add to team db
new_member = data.member
# get user id
if new_member.user_id is None and new_member.user_email is not None:
for user in updated_users:
if (
user.user_email is not None
and user.user_email == new_member.user_email
):
new_member.user_id = user.user_id
complete_team_data.members_with_roles.append(new_member)
elif isinstance(data.member, List):
# add to team db
new_members = data.member
for nm in new_members:
if nm.user_id is None and nm.user_email is not None:
for user in updated_users:
if user.user_email is not None and user.user_email == nm.user_email:
nm.user_id = user.user_id
complete_team_data.members_with_roles.extend(new_members)
# ADD MEMBER TO TEAM
_db_team_members = [m.model_dump() for m in complete_team_data.members_with_roles]
updated_team = await prisma_client.db.litellm_teamtable.update(
where={"team_id": data.team_id},
data={"members_with_roles": json.dumps(_db_team_members)}, # type: ignore
)
# Check if updated_team is None
if updated_team is None:
raise HTTPException(
@ -897,7 +914,17 @@ async def team_info(
# if using pydantic v1
key = key.dict()
key.pop("token", None)
return {"team_id": team_id, "team_info": team_info, "keys": keys}
## GET ALL MEMBERSHIPS ##
team_memberships = await prisma_client.db.litellm_teammembership.find_many(
where={"team_id": team_id}
)
return {
"team_id": team_id,
"team_info": team_info,
"keys": keys,
"team_memberships": team_memberships,
}
except Exception as e:
verbose_proxy_logger.error(