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 a6510c7f220..8bcb45f812b 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) @@ -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": [, ...] - }' - ``` + (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, - ) + 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): diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 031db99d13c..6dd89bd8534 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 { + 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 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 diff --git a/schema.prisma b/schema.prisma index 031db99d13c..6dd89bd8534 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 { + 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