fix(proxy_server.py): write blocked user list to a db table

lets this persist across workers
This commit is contained in:
Krrish Dholakia 2024-03-16 12:26:29 -07:00
parent 2c2f322d5a
commit ef0002f31c
3 changed files with 42 additions and 25 deletions

View file

@ -5094,39 +5094,36 @@ async def user_get_requests():
@router.post(
"/user/block",
tags=["user management"],
tags=["End User Management"],
dependencies=[Depends(user_api_key_auth)],
)
async def block_user(data: BlockUsers):
"""
[BETA] Reject calls with this user id
[BETA] Reject calls with this end-user id
```
curl -X POST "http://0.0.0.0:8000/user/block"
-H "Authorization: Bearer sk-1234"
-D '{
"user_ids": [<user_id>, ...]
}'
```
(any /chat/completion call with this user={end-user-id} param, will be rejected.)
```
curl -X POST "http://0.0.0.0:8000/user/block"
-H "Authorization: Bearer sk-1234"
-D '{
"user_ids": [<user_id>, ...]
}'
```
"""
from enterprise.enterprise_hooks.blocked_user_list import (
_ENTERPRISE_BlockedUserList,
)
if not any(isinstance(x, _ENTERPRISE_BlockedUserList) for x in litellm.callbacks):
blocked_user_list = _ENTERPRISE_BlockedUserList()
litellm.callbacks.append(blocked_user_list) # type: ignore
if litellm.blocked_user_list is None:
litellm.blocked_user_list = data.user_ids
elif isinstance(litellm.blocked_user_list, list):
litellm.blocked_user_list = litellm.blocked_user_list + data.user_ids
if prisma_client is not None:
for id in data.user_ids:
await prisma_client.db.litellm_endusertable.upsert(
where={"id": id},
data={
"create": {"id": id, "blocked": True},
"update": {"blocked": True},
},
)
else:
raise HTTPException(
status_code=500,
detail={
"error": "`blocked_user_list` must be a list or not set. Filepaths can't be updated."
},
detail={"error": "Postgres DB Not connected"},
)
return {"blocked_users": litellm.blocked_user_list}
@ -5134,7 +5131,7 @@ async def block_user(data: BlockUsers):
@router.post(
"/user/unblock",
tags=["user management"],
tags=["End User Management"],
dependencies=[Depends(user_api_key_auth)],
)
async def unblock_user(data: BlockUsers):

View file

@ -24,6 +24,7 @@ model LiteLLM_BudgetTable {
updated_by String
organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget
keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
}
model LiteLLM_OrganizationTable {
@ -127,6 +128,15 @@ model LiteLLM_VerificationToken {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
}
model LiteLLM_EndUserTable {
id String @id
alias String? // admin-facing alias
spend Float @default(0.0)
budget_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
blocked Boolean @default(false)
}
// store proxy config.yaml
model LiteLLM_Config {
param_name String @id

View file

@ -24,6 +24,7 @@ model LiteLLM_BudgetTable {
updated_by String
organization LiteLLM_OrganizationTable[] // multiple orgs can have the same budget
keys LiteLLM_VerificationToken[] // multiple keys can have the same budget
end_users LiteLLM_EndUserTable[] // multiple end-users can have the same budget
}
model LiteLLM_OrganizationTable {
@ -127,6 +128,15 @@ model LiteLLM_VerificationToken {
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
}
model LiteLLM_EndUserTable {
id String @id
alias String? // admin-facing alias
spend Float @default(0.0)
budget_id String?
litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id])
blocked Boolean @default(false)
}
// store proxy config.yaml
model LiteLLM_Config {
param_name String @id