feat(proxy_server.py): enable new /team/disable endpoint

reject all requests from this team id, without deleting it.

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):

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:

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 ####
 
 

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")

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("{}")