fix(proxy_server.py): use consistent naming schema - move to /team/block

This commit is contained in:
Krrish Dholakia 2024-03-26 16:59:36 -07:00
parent 4f7ba902d8
commit 77472b80eb
5 changed files with 42 additions and 12 deletions

View file

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

View file

@ -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 (

View file

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

View file

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

View file

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