From 4f7ba902d84631185fa3f10f1f28c4ea258cad4e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 26 Mar 2024 15:40:50 -0700 Subject: [PATCH 1/3] feat(proxy_server.py): enable new `/team/disable` endpoint reject all requests from this team id, without deleting it. --- litellm/proxy/_types.py | 1 + litellm/proxy/auth/auth_checks.py | 17 +++++++++++------ litellm/proxy/proxy_server.py | 20 ++++++++++++++++++++ litellm/proxy/schema.prisma | 1 + schema.prisma | 2 ++ 5 files changed, 35 insertions(+), 6 deletions(-) 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("{}") From 77472b80eb1a760648c232899c2581b9f3cb2ed6 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 26 Mar 2024 16:59:36 -0700 Subject: [PATCH 2/3] 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 c8921dfa8..2cd979b4b 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 ed2603b34..37ec2065f 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 d510b5d7a..2f1e61d9b 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 c89b4e27f..c66245544 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 99ed89380..598848776 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("{}") From 5b66cb3864b9c5a5c2b7c72315bc2e86280f4d4d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 26 Mar 2024 18:06:49 -0700 Subject: [PATCH 3/3] test(test_exceptions.py): handle api instability --- litellm/tests/test_exceptions.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/tests/test_exceptions.py b/litellm/tests/test_exceptions.py index 311bbfa57..13ad42b50 100644 --- a/litellm/tests/test_exceptions.py +++ b/litellm/tests/test_exceptions.py @@ -85,6 +85,8 @@ def test_context_window_with_fallbacks(model): ) except litellm.ServiceUnavailableError as e: pass + except litellm.APIConnectionError as e: + pass # for model in litellm.models_by_provider["bedrock"]: