feat(team_endpoints.py): support using default team member budget id, if set

allows all team members to use the same budget id
This commit is contained in:
Krrish Dholakia 2025-06-20 15:37:51 -07:00
parent acef5324b1
commit 90fa3f61a2
2 changed files with 17 additions and 7 deletions

View file

@ -896,8 +896,6 @@ async def team_member_add(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
[BETA]
Add new members (either via user_email or user_id) to a team
If user doesn't exist, new user row will also be added to User Table
@ -991,6 +989,11 @@ async def team_member_add(
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=(
complete_team_data.metadata.get("team_member_budget_id")
if complete_team_data.metadata is not None
else None
),
)
except Exception as e:
raise HTTPException(
@ -1016,6 +1019,11 @@ async def team_member_add(
user_api_key_dict=user_api_key_dict,
litellm_proxy_admin_name=litellm_proxy_admin_name,
team_id=data.team_id,
default_team_budget_id=(
complete_team_data.metadata.get("team_member_budget_id")
if complete_team_data.metadata is not None
else None
),
)
except Exception as e:
raise HTTPException(

View file

@ -60,6 +60,7 @@ async def add_new_member(
team_id: str,
user_api_key_dict: UserAPIKeyAuth,
litellm_proxy_admin_name: str,
default_team_budget_id: Optional[str] = None,
) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]:
"""
Add a new member to a team
@ -119,11 +120,8 @@ async def add_new_member(
)
# Check if trying to set a budget for team member
if (
max_budget_in_team is not None
and returned_user is not None
and returned_user.user_id is not None
):
if max_budget_in_team is not None:
# create a new budget item for this member
response = await prisma_client.db.litellm_budgettable.create(
data={
@ -134,6 +132,10 @@ async def add_new_member(
)
_budget_id = response.budget_id
else:
_budget_id = default_team_budget_id
if returned_user is not None and returned_user.user_id is not None:
_returned_team_membership = (
await prisma_client.db.litellm_teammembership.create(
data={