feat - create budgets when team/member_add

litellm/proxy/_types.py

@@ -563,6 +563,7 @@ class GlobalEndUsersSpend(LiteLLMBase):
 class TeamMemberAddRequest(LiteLLMBase):
     team_id: str
     member: Member
+    max_budget_in_team: Optional[float] = None  # Users max budget within the team
 
 
 class TeamMemberDeleteRequest(LiteLLMBase):

litellm/proxy/proxy_server.py

@@ -7605,16 +7605,12 @@ async def team_member_add(
     If user doesn't exist, new user row will also be added to User Table
 
     ```
-    curl -X POST 'http://0.0.0.0:8000/team/update' \
 
+    curl -X POST 'http://0.0.0.0:4000/team/member_add' \
     -H 'Authorization: Bearer sk-1234' \
-
     -H 'Content-Type: application/json' \
+    -d '{"team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849", "member": {"role": "user", "user_id": "[email protected]"}}'
 
-    -D '{
-        "team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849",
-        "member": {"role": "user", "user_id": "[email protected]"}
-    }'
     ```
     """
     if prisma_client is None:
@@ -7685,6 +7681,26 @@ async def team_member_add(
 
             await prisma_client.insert_data(data=user_data, table_name="user")
 
+    # Check if trying to set a budget for team member
+    if data.max_budget_in_team is not None and new_member.user_id is not None:
+        # create a new budget item for this member
+        response = await prisma_client.db.litellm_budgettable.create(
+            data={
+                "max_budget": data.max_budget_in_team,
+                "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+                "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name,
+            }
+        )
+
+        _budget_id = response.budget_id
+        await prisma_client.db.litellm_teammembership.create(
+            data={
+                "team_id": data.team_id,
+                "user_id": new_member.user_id,
+                "budget_id": _budget_id,
+            }
+        )
+
     return team_row
 
 

schema.prisma

@@ -25,6 +25,7 @@ model LiteLLM_BudgetTable {
   organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget
   keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
   end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
+  team_membership LiteLLM_TeamMembership[] // budgets of Users within a Team 
 }
 
 // Models on proxy
@@ -208,4 +209,14 @@ model LiteLLM_UserNotifications {
   models              String[]
   justification       String
   status              String // approved, disapproved, pending
-}
+}
+
+model LiteLLM_TeamMembership {
+  // Use this table to track the Internal User's Spend within a Team + Set Budgets, rpm limits for the user within the team
+  user_id    String
+  team_id    String
+  spend      Float    @default(0.0)
+  budget_id String?
+  litellm_budget_table LiteLLM_BudgetTable?   @relation(fields: [budget_id], references: [budget_id])
+  @@id([user_id, team_id])
+}