diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4fd1bf3b0..c8921dfa8 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -396,6 +396,7 @@ class TeamBase(LiteLLMBase): rpm_limit: Optional[int] = None max_budget: Optional[float] = None models: list = [] + disabled: bool = False class NewTeamRequest(TeamBase): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index b8f7c6e3f..ed2603b34 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -30,12 +30,17 @@ def common_checks( """ Common checks across jwt + key-based auth. - 1. If user can call model - 2. If user is in budget - 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget + 1. If team is disabled + 2. If team can call model + 3. If team is in budget + 4. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget """ _model = request_body.get("model", None) - # 1. If user can call model + if team_object.disabled == True: + raise Exception( + f"Team={team_object.team_id} is disabled. Update via `/team/update`." + ) + # 2. If user can call model if ( _model is not None and len(team_object.models) > 0 @@ -44,7 +49,7 @@ def common_checks( raise Exception( f"Team={team_object.team_id} not allowed to call model={_model}. Allowed team models = {team_object.models}" ) - # 2. If team is in budget + # 3. If team is in budget if ( team_object.max_budget is not None and team_object.spend is not None @@ -53,7 +58,7 @@ def common_checks( raise Exception( f"Team={team_object.team_id} over budget. Spend={team_object.spend}, Budget={team_object.max_budget}" ) - # 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget + # 4. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget if end_user_object is not None and end_user_object.litellm_budget_table is not None: end_user_budget = end_user_object.litellm_budget_table.max_budget if end_user_budget is not None and end_user_object.spend > end_user_budget: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 656abe474..d510b5d7a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -6151,6 +6151,26 @@ async def team_info( ) +@router.post( + "/team/disable", tags=["team management"], dependencies=[Depends(user_api_key_auth)] +) +async def disable_team( + data: DeleteTeamRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Sets + """ + global prisma_client + + if prisma_client is None: + raise Exception("No DB Connected.") + + await prisma_client.db.litellm_teamtable.update_many( + where={"team_id": {"in": data.team_ids}}, data={"disabled": True} + ) + + #### ORGANIZATION MANAGEMENT #### diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index cccad973a..c89b4e27f 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -70,6 +70,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + disabled Boolean @default(false) budget_duration String? budget_reset_at DateTime? created_at DateTime @default(now()) @map("created_at") diff --git a/schema.prisma b/schema.prisma index cccad973a..99ed89380 100644 --- a/schema.prisma +++ b/schema.prisma @@ -55,6 +55,7 @@ model LiteLLM_ModelTable { team LiteLLM_TeamTable? } + // Assign prod keys to groups, not individuals model LiteLLM_TeamTable { team_id String @id @default(uuid()) @@ -72,6 +73,7 @@ model LiteLLM_TeamTable { rpm_limit BigInt? budget_duration String? budget_reset_at DateTime? + disabled Boolean @default(false) created_at DateTime @default(now()) @map("created_at") updated_at DateTime @default(now()) @updatedAt @map("updated_at") model_spend Json @default("{}")