From ef0002f31c90a107115d0cf403689f7530b381f5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 16 Mar 2024 12:26:29 -0700 Subject: [PATCH 1/5] fix(proxy_server.py): write blocked user list to a db table lets this persist across workers --- litellm/proxy/proxy_server.py | 47 ++++++++++++++++------------------- litellm/proxy/schema.prisma | 10 ++++++++ schema.prisma | 10 ++++++++ 3 files changed, 42 insertions(+), 25 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index a6510c7f220..b1aac779167 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 031db99d13c..c11a387a8b5 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 031db99d13c..c11a387a8b5 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 From dd151869a3ace1f6e422f137b90889431daeb424 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 16 Mar 2024 13:03:52 -0700 Subject: [PATCH 2/5] fix(blocked_user_list.py): check if end user blocked in db --- .../enterprise_hooks/blocked_user_list.py | 54 ++++++++++++++++--- litellm/proxy/_types.py | 16 ++++++ litellm/proxy/proxy_server.py | 37 +++++++------ litellm/proxy/schema.prisma | 2 +- 4 files changed, 86 insertions(+), 23 deletions(-) diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py index 686fdf1de2a..cbc14d2c2b8 100644 --- a/enterprise/enterprise_hooks/blocked_user_list.py +++ b/enterprise/enterprise_hooks/blocked_user_list.py @@ -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: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8a7efa1a168..a8c0c3d270e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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 diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b1aac779167..91fc8295d79 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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) @@ -5111,22 +5113,27 @@ async def block_user(data: BlockUsers): }' ``` """ - 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}, - }, + 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}, + data={ + "create": {"user_id": id, "blocked": True}, + "update": {"blocked": True}, + }, + ) + records.append(record) + else: + raise HTTPException( + status_code=500, + detail={"error": "Postgres DB Not connected"}, ) - else: - raise HTTPException( - status_code=500, - detail={"error": "Postgres DB Not connected"}, - ) - return {"blocked_users": litellm.blocked_user_list} + return {"blocked_users": records} + except Exception as e: + raise HTTPException(status_code=500, detail={"error": str(e)}) @router.post( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index c11a387a8b5..6dd89bd8534 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -129,7 +129,7 @@ model LiteLLM_VerificationToken { } model LiteLLM_EndUserTable { - id String @id + user_id String @id alias String? // admin-facing alias spend Float @default(0.0) budget_id String? From 844d6828c86630452b0b85137c835a87e54ed7ee Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 16 Mar 2024 13:53:10 -0700 Subject: [PATCH 3/5] fix(proxy_server.py): rename beta end user blocking endpoints separating end user management from user management --- litellm/proxy/proxy_server.py | 5 +- litellm/tests/test_blocked_user_list.py | 91 ++++++++++++++++++++++++- 2 files changed, 92 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 91fc8295d79..4eb21ea0709 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5095,7 +5095,7 @@ async def user_get_requests(): @router.post( - "/user/block", + "/end_user/block", tags=["End User Management"], dependencies=[Depends(user_api_key_auth)], ) @@ -5133,11 +5133,12 @@ async def block_user(data: BlockUsers): 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", + "/end_user/unblock", tags=["End User Management"], dependencies=[Depends(user_api_key_auth)], ) diff --git a/litellm/tests/test_blocked_user_list.py b/litellm/tests/test_blocked_user_list.py index b40d8296c30..d3f9f6a1a32 100644 --- a/litellm/tests/test_blocked_user_list.py +++ b/litellm/tests/test_blocked_user_list.py @@ -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 From 3de766c071d2f8e512c7835b73676633059f8534 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 16 Mar 2024 14:02:36 -0700 Subject: [PATCH 4/5] fix(proxy_server.py): fix linting errors --- litellm/proxy/proxy_server.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4eb21ea0709..8bcb45f812b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5118,9 +5118,9 @@ async def block_user(data: BlockUsers): if prisma_client is not None: for id in data.user_ids: record = await prisma_client.db.litellm_endusertable.upsert( - where={"user_id": id}, + where={"user_id": id}, # type: ignore data={ - "create": {"user_id": id, "blocked": True}, + "create": {"user_id": id, "blocked": True}, # type: ignore "update": {"blocked": True}, }, ) From 26e783ab59e011977939b195442e8671e97197a8 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 16 Mar 2024 14:15:01 -0700 Subject: [PATCH 5/5] build: fix field name --- schema.prisma | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/schema.prisma b/schema.prisma index c11a387a8b5..6dd89bd8534 100644 --- a/schema.prisma +++ b/schema.prisma @@ -129,7 +129,7 @@ model LiteLLM_VerificationToken { } model LiteLLM_EndUserTable { - id String @id + user_id String @id alias String? // admin-facing alias spend Float @default(0.0) budget_id String?