From 94c5d9f82fe02e9658b93be5f07ca0bb4173f0cf Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 26 Mar 2024 16:59:36 -0700
Subject: [PATCH] fix(proxy_server.py): use consistent naming schema - move to
 `/team/block`

---
 litellm/proxy/_types.py           |  8 ++++++-
 litellm/proxy/auth/auth_checks.py |  6 +++---
 litellm/proxy/proxy_server.py     | 36 +++++++++++++++++++++++++------
 litellm/proxy/schema.prisma       |  2 +-
 schema.prisma                     |  2 +-
 5 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py
index c8921dfa8c..2cd979b4b8 100644
--- a/litellm/proxy/_types.py
+++ b/litellm/proxy/_types.py
@@ -82,6 +82,8 @@ class LiteLLMRoutes(enum.Enum):
         "/team/update",
         "/team/delete",
         "/team/info",
+        "/team/block",
+        "/team/unblock",
         # model
         "/model/new",
         "/model/update",
@@ -396,7 +398,7 @@ class TeamBase(LiteLLMBase):
     rpm_limit: Optional[int] = None
     max_budget: Optional[float] = None
     models: list = []
-    disabled: bool = False
+    blocked: bool = False
 
 
 class NewTeamRequest(TeamBase):
@@ -437,6 +439,10 @@ class DeleteTeamRequest(LiteLLMBase):
     team_ids: List[str]  # required
 
 
+class BlockTeamRequest(LiteLLMBase):
+    team_id: str  # required
+
+
 class LiteLLM_TeamTable(TeamBase):
     spend: Optional[float] = None
     max_parallel_requests: Optional[int] = None
diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py
index ed2603b343..37ec2065f6 100644
--- a/litellm/proxy/auth/auth_checks.py
+++ b/litellm/proxy/auth/auth_checks.py
@@ -30,15 +30,15 @@ def common_checks(
     """
     Common checks across jwt + key-based auth.
 
-    1. If team is disabled
+    1. If team is blocked
     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)
-    if team_object.disabled == True:
+    if team_object.blocked == True:
         raise Exception(
-            f"Team={team_object.team_id} is disabled. Update via `/team/update`."
+            f"Team={team_object.team_id} is blocked. Update via `/team/unblock` if your admin."
         )
     # 2. If user can call model
     if (
diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index d510b5d7a5..2f1e61d9bb 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -6152,24 +6152,48 @@ async def team_info(
 
 
 @router.post(
-    "/team/disable", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
+    "/team/block", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
 )
-async def disable_team(
-    data: DeleteTeamRequest,
+async def block_team(
+    data: BlockTeamRequest,
     user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
 ):
     """
-    Sets
+    Blocks all calls from keys with this team id.
     """
     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}
+    record = await prisma_client.db.litellm_teamtable.update(
+        where={"team_id": data.team_id}, data={"blocked": True}
     )
 
+    return record
+
+
+@router.post(
+    "/team/unblock", tags=["team management"], dependencies=[Depends(user_api_key_auth)]
+)
+async def unblock_team(
+    data: BlockTeamRequest,
+    user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
+):
+    """
+    Blocks all calls from keys with this team id.
+    """
+    global prisma_client
+
+    if prisma_client is None:
+        raise Exception("No DB Connected.")
+
+    record = await prisma_client.db.litellm_teamtable.update(
+        where={"team_id": data.team_id}, data={"blocked": False}
+    )
+
+    return record
+
 
 #### ORGANIZATION MANAGEMENT ####
 
diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma
index c89b4e27fd..c662455447 100644
--- a/litellm/proxy/schema.prisma
+++ b/litellm/proxy/schema.prisma
@@ -70,7 +70,7 @@ model LiteLLM_TeamTable {
     max_parallel_requests Int?
     tpm_limit     BigInt?
     rpm_limit     BigInt?
-    disabled Boolean @default(false) 
+    blocked 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 99ed89380d..5988487768 100644
--- a/schema.prisma
+++ b/schema.prisma
@@ -73,7 +73,7 @@ model LiteLLM_TeamTable {
     rpm_limit     BigInt?
     budget_duration String? 
     budget_reset_at DateTime?
-    disabled Boolean @default(false) 
+    blocked Boolean @default(false)
     created_at    DateTime               @default(now()) @map("created_at")
     updated_at    DateTime               @default(now()) @updatedAt @map("updated_at")
     model_spend      Json @default("{}")