From 50461eb22c604f2117b514e5783b772b977c720b Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 22 May 2024 17:16:02 -0700 Subject: [PATCH] feat - create budgets when team/member_add --- litellm/proxy/_types.py | 1 + litellm/proxy/proxy_server.py | 28 ++++++++++++++++++++++------ schema.prisma | 13 ++++++++++++- 3 files changed, 35 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index ce38f4f9e..e834d98b0 100644 --- a/litellm/proxy/_types.py +++ b/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): diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fbcff31c2..1bbb050f4 100644 --- a/litellm/proxy/proxy_server.py +++ b/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": "krrish247652@berri.ai"}}' - -D '{ - "team_id": "45e3e396-ee08-4a61-a88e-16b3ce7e0849", - "member": {"role": "user", "user_id": "krrish247652@berri.ai"} - }' ``` """ 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 diff --git a/schema.prisma b/schema.prisma index 7ffeab9ba..0fceaa12a 100644 --- a/schema.prisma +++ b/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 -} \ No newline at end of file +} + +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]) +}