diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e64f43eed8d..23bca14a694 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -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( diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index ec645c0b156..170ece20488 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -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={