scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-08 22:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge pull request #2552 from BerriAI/litellm_end_user_opt_out_db

Browse source 
fix(proxy_server.py): write blocked user list to a db table
This commit is contained in:
Krish Dholakia 2024-03-16 14:38:43 -07:00 • committed by GitHub
commit 3197ebeb7e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 210 additions and 42 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
enterprise/enterprise_hooks/blocked_user_list.py
View file

@ -9,8 +9,9 @@
from typing import Optional, Literal
import litellm
from litellm.proxy.utils import PrismaClient
from litellm.caching import DualCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy._types import UserAPIKeyAuth, LiteLLM_EndUserTable
from litellm.integrations.custom_logger import CustomLogger
from litellm._logging import verbose_proxy_logger
from fastapi import HTTPException
@ -19,13 +20,13 @@ import json, traceback
class _ENTERPRISE_BlockedUserList(CustomLogger):
# Class variables or attributes
def __init__(self):
blocked_user_list = litellm.blocked_user_list
def __init__(self, prisma_client: Optional[PrismaClient]):
self.prisma_client = prisma_client
blocked_user_list = litellm.blocked_user_list
if blocked_user_list is None:
raise Exception(
"`blocked_user_list` can either be a list or filepath. None set."
)
self.blocked_user_list = None
return
if isinstance(blocked_user_list, list):
self.blocked_user_list = blocked_user_list
@ -64,17 +65,56 @@ class _ENTERPRISE_BlockedUserList(CustomLogger):
"""
- check if user id part of call
- check if user id part of blocked list
- if blocked list is none or user not in blocked list
- check if end-user in cache
- check if end-user in db
"""
self.print_verbose(f"Inside Blocked User List Pre-Call Hook")
if "user_id" in data or "user" in data:
user = data.get("user_id", data.get("user", ""))
if user in self.blocked_user_list:
if (
self.blocked_user_list is not None
and user in self.blocked_user_list
):
raise HTTPException(
status_code=400,
detail={
"error": f"User blocked from making LLM API Calls. User={user}"
},
)
cache_key = f"litellm:end_user_id:{user}"
end_user_cache_obj: LiteLLM_EndUserTable = cache.get_cache(
key=cache_key
)
if end_user_cache_obj is None and self.prisma_client is not None:
# check db
end_user_obj = (
await self.prisma_client.db.litellm_endusertable.find_unique(
where={"user_id": user}
)
)
if end_user_obj is None: # user not in db - assume not blocked
end_user_obj = LiteLLM_EndUserTable(user_id=user, blocked=False)
cache.set_cache(key=cache_key, value=end_user_obj, ttl=60)
if end_user_obj is not None and end_user_obj.blocked == True:
raise HTTPException(
status_code=400,
detail={
"error": f"User blocked from making LLM API Calls. User={user}"
},
)
elif (
end_user_cache_obj is not None
and end_user_cache_obj.blocked == True
):
raise HTTPException(
status_code=400,
detail={
"error": f"User blocked from making LLM API Calls. User={user}"
},
)
except HTTPException as e:
raise e
except Exception as e:

16
litellm/proxy/_types.py
View file

@ -602,6 +602,22 @@ class LiteLLM_UserTable(LiteLLMBase):
protected_namespaces = ()
class LiteLLM_EndUserTable(LiteLLMBase):
user_id: str
blocked: bool
alias: Optional[str] = None
spend: float = 0.0
@root_validator(pre=True)
def set_model_info(cls, values):
if values.get("spend") is None:
values.update({"spend": 0.0})
return values
class Config:
protected_namespaces = ()
class LiteLLM_SpendLogs(LiteLLMBase):
request_id: str
api_key: str

71
litellm/proxy/proxy_server.py
View file

@ -1774,7 +1774,9 @@ class ProxyConfig:
_ENTERPRISE_BlockedUserList,
)
blocked_user_list = _ENTERPRISE_BlockedUserList()
blocked_user_list = _ENTERPRISE_BlockedUserList(
prisma_client=prisma_client
)
imported_list.append(blocked_user_list)
elif (
isinstance(callback, str)
@ -5093,48 +5095,51 @@ async def user_get_requests():
@router.post(
"/user/block",
tags=["user management"],
"/end_user/block",
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,
)
try:
records = []
if prisma_client is not None:
for id in data.user_ids:
record = await prisma_client.db.litellm_endusertable.upsert(
where={"user_id": id}, # type: ignore
data={
"create": {"user_id": id, "blocked": True}, # type: ignore
"update": {"blocked": True},
},
)
records.append(record)
else:
raise HTTPException(
status_code=500,
detail={"error": "Postgres DB Not connected"},
)
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
else:
raise HTTPException(
status_code=500,
detail={
"error": "`blocked_user_list` must be a list or not set. Filepaths can't be updated."
},
)
return {"blocked_users": litellm.blocked_user_list}
return {"blocked_users": records}
except Exception as e:
verbose_proxy_logger.error(f"An error occurred - {str(e)}")
raise HTTPException(status_code=500, detail={"error": str(e)})
@router.post(
"/user/unblock",
tags=["user management"],
"/end_user/unblock",
tags=["End User Management"],
dependencies=[Depends(user_api_key_auth)],
)
async def unblock_user(data: BlockUsers):

10
litellm/proxy/schema.prisma
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 {
user_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

91
litellm/tests/test_blocked_user_list.py
View file

@ -6,6 +6,7 @@ import sys, os, asyncio, time, random
from datetime import datetime
import traceback
from dotenv import load_dotenv
from fastapi import Request
load_dotenv()
import os
@ -22,18 +23,87 @@ from litellm import Router, mock_completion
from litellm.proxy.utils import ProxyLogging
from litellm.proxy._types import UserAPIKeyAuth
from litellm.caching import DualCache
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
import pytest, logging, asyncio
import litellm, asyncio
from litellm.proxy.proxy_server import (
new_user,
generate_key_fn,
user_api_key_auth,
user_update,
delete_key_fn,
info_key_fn,
update_key_fn,
generate_key_fn,
generate_key_helper_fn,
spend_user_fn,
spend_key_fn,
view_spend_logs,
user_info,
block_user,
)
from litellm.proxy.utils import PrismaClient, ProxyLogging, hash_token
from litellm._logging import verbose_proxy_logger
verbose_proxy_logger.setLevel(level=logging.DEBUG)
from litellm.proxy._types import (
NewUserRequest,
GenerateKeyRequest,
DynamoDBArgs,
KeyRequest,
UpdateKeyRequest,
GenerateKeyRequest,
BlockUsers,
)
from litellm.proxy.utils import DBClient
from starlette.datastructures import URL
from litellm.caching import DualCache
proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache())
@pytest.fixture
def prisma_client():
from litellm.proxy.proxy_cli import append_query_params
### add connection pool + pool timeout args
params = {"connection_limit": 100, "pool_timeout": 60}
database_url = os.getenv("DATABASE_URL")
modified_url = append_query_params(database_url, params)
os.environ["DATABASE_URL"] = modified_url
# Assuming DBClient is a class that needs to be instantiated
prisma_client = PrismaClient(
database_url=os.environ["DATABASE_URL"], proxy_logging_obj=proxy_logging_obj
)
# Reset litellm.proxy.proxy_server.prisma_client to None
litellm.proxy.proxy_server.custom_db_client = None
litellm.proxy.proxy_server.litellm_proxy_budget_name = (
f"litellm-proxy-budget-{time.time()}"
)
litellm.proxy.proxy_server.user_custom_key_generate = None
return prisma_client
@pytest.mark.asyncio
async def test_block_user_check():
async def test_block_user_check(prisma_client):
"""
- Set a blocked user as a litellm module value
- Test to see if a call with that user id is made, an error is raised
- Test to see if a call without that user is passes
"""
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
litellm.blocked_user_list = ["user_id_1"]
blocked_user_obj = _ENTERPRISE_BlockedUserList()
blocked_user_obj = _ENTERPRISE_BlockedUserList(
prisma_client=litellm.proxy.proxy_server.prisma_client
)
_api_key = "sk-12345"
user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
@ -61,3 +131,20 @@ async def test_block_user_check():
)
except Exception as e:
pytest.fail(f"An error occurred - {str(e)}")
@pytest.mark.asyncio
async def test_block_user_db_check(prisma_client):
"""
- Block end user via "/user/block"
- Check returned value
"""
setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client)
setattr(litellm.proxy.proxy_server, "master_key", "sk-1234")
await litellm.proxy.proxy_server.prisma_client.connect()
_block_users = BlockUsers(user_ids=["user_id_1"])
result = await block_user(data=_block_users)
result = result["blocked_users"]
assert len(result) == 1
assert result[0].user_id == "user_id_1"
assert result[0].blocked == True

10
schema.prisma
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 {
user_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