diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a6510c7f2..b1aac7791 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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": [, ...] - }' - ``` + (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": [, ...] + }' + ``` """ - 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): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 031db99d1..c11a387a8 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -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 diff --git a/schema.prisma b/schema.prisma index 031db99d1..c11a387a8 100644 --- a/schema.prisma +++ b/schema.prisma @@ -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