mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-26 19:24:27 +00:00
fix(proxy_server.py): use consistent naming schema - move to /team/block
This commit is contained in:
parent
ce44645000
commit
94c5d9f82f
5 changed files with 42 additions and 12 deletions
|
@ -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
|
||||
|
|
|
@ -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 (
|
||||
|
|
|
@ -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 ####
|
||||
|
||||
|
|
|
@ -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")
|
||||
|
|
|
@ -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("{}")
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue