Merge pull request #2144 from BerriAI/litellm_blocked_user_list

feat(proxy_server.py): add support for blocked user lists (enterprise-only)

docs/my-website/docs/proxy/enterprise.md

@@ -1,7 +1,7 @@
 import Tabs from '@theme/Tabs';
 import TabItem from '@theme/TabItem';
 
-# ✨ Enterprise Features - Content Moderation
+# ✨ Enterprise Features - Content Moderation, Blocked Users
 
 Features here are behind a commercial license in our `/enterprise` folder. [**See Code**](https://github.com/BerriAI/litellm/tree/main/enterprise)
 
@@ -15,6 +15,7 @@ Features:
 - [ ] Content Moderation with LlamaGuard 
 - [ ] Content Moderation with Google Text Moderations 
 - [ ] Content Moderation with LLM Guard
+- [ ] Reject calls from Blocked User list 
 - [ ] Tracking Spend for Custom Tags
  
 ## Content Moderation with LlamaGuard 
@@ -132,6 +133,39 @@ Here are the category specific values:
 
 
 
+## Enable Blocked User Lists 
+If any call is made to proxy with this user id, it'll be rejected - use this if you want to let users opt-out of ai features 
+
+```yaml
+litellm_settings: 
+     callbacks: ["blocked_user_check"] 
+     blocked_user_id_list: ["user_id_1", "user_id_2", ...]  # can also be a .txt filepath e.g. `/relative/path/blocked_list.txt` 
+```
+
+### How to test
+
+```bash 
+curl --location 'http://0.0.0.0:8000/chat/completions' \
+--header 'Content-Type: application/json' \
+--data ' {
+      "model": "gpt-3.5-turbo",
+      "messages": [
+        {
+          "role": "user",
+          "content": "what llm are you"
+        }
+      ],
+      "user_id": "user_id_1" # this is also an openai supported param 
+    }
+'
+```
+
+:::info 
+
+[Suggest a way to improve this](https://github.com/BerriAI/litellm/issues/new/choose)
+
+:::
+
 ## Tracking Spend for Custom Tags
 
 Requirements: 

enterprise/enterprise_hooks/blocked_user_list.py

@@ -0,0 +1,80 @@
+# +------------------------------+
+#
+#        Blocked User List
+#
+# +------------------------------+
+#  Thank you users! We ❤️ you! - Krrish & Ishaan
+## This accepts a list of user id's for whom calls will be rejected
+
+
+from typing import Optional, Literal
+import litellm
+from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.integrations.custom_logger import CustomLogger
+from litellm._logging import verbose_proxy_logger
+from fastapi import HTTPException
+import json, traceback
+
+
+class _ENTERPRISE_BlockedUserList(CustomLogger):
+    # Class variables or attributes
+    def __init__(self):
+        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."
+            )
+
+        if isinstance(blocked_user_list, list):
+            self.blocked_user_list = blocked_user_list
+
+        if isinstance(blocked_user_list, str):  # assume it's a filepath
+            try:
+                with open(blocked_user_list, "r") as file:
+                    data = file.read()
+                    self.blocked_user_list = data.split("\n")
+            except FileNotFoundError:
+                raise Exception(
+                    f"File not found. blocked_user_list={blocked_user_list}"
+                )
+            except Exception as e:
+                raise Exception(
+                    f"An error occurred: {str(e)}, blocked_user_list={blocked_user_list}"
+                )
+
+    def print_verbose(self, print_statement, level: Literal["INFO", "DEBUG"] = "DEBUG"):
+        if level == "INFO":
+            verbose_proxy_logger.info(print_statement)
+        elif level == "DEBUG":
+            verbose_proxy_logger.debug(print_statement)
+
+        if litellm.set_verbose is True:
+            print(print_statement)  # noqa
+
+    async def async_pre_call_hook(
+        self,
+        user_api_key_dict: UserAPIKeyAuth,
+        cache: DualCache,
+        data: dict,
+        call_type: str,
+    ):
+        try:
+            """
+            - check if user id part of call
+            - check if user id part of blocked list
+            """
+            self.print_verbose(f"Inside Blocked User List Pre-Call Hook")
+            if "user_id" in data:
+                if data["user_id"] in self.blocked_user_list:
+                    raise HTTPException(
+                        status_code=400,
+                        detail={
+                            "error": f"User blocked from making LLM API Calls. User={data['user_id']}"
+                        },
+                    )
+        except HTTPException as e:
+            raise e
+        except Exception as e:
+            traceback.print_exc()

litellm/__init__.py

@@ -60,6 +60,7 @@ llamaguard_model_name: Optional[str] = None
 presidio_ad_hoc_recognizers: Optional[str] = None
 google_moderation_confidence_threshold: Optional[float] = None
 llamaguard_unsafe_content_categories: Optional[str] = None
+blocked_user_list: Optional[Union[str, List]] = None
 ##################
 logging: bool = True
 caching: bool = (

litellm/proxy/proxy_server.py

@@ -1479,6 +1479,16 @@ class ProxyConfig:
 
                                 llm_guard_moderation_obj = _ENTERPRISE_LLMGuard()
                                 imported_list.append(llm_guard_moderation_obj)
+                            elif (
+                                isinstance(callback, str)
+                                and callback == "blocked_user_check"
+                            ):
+                                from litellm.proxy.enterprise.enterprise_hooks.blocked_user_list import (
+                                    _ENTERPRISE_BlockedUserList,
+                                )
+
+                                blocked_user_list = _ENTERPRISE_BlockedUserList()
+                                imported_list.append(blocked_user_list)
                             else:
                                 imported_list.append(
                                     get_instance_fn(

litellm/tests/test_blocked_user_list.py

@@ -0,0 +1,63 @@
+# What is this?
+## This tests the blocked user pre call hook for the proxy server
+
+
+import sys, os, asyncio, time, random
+from datetime import datetime
+import traceback
+from dotenv import load_dotenv
+
+load_dotenv()
+import os
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import pytest
+import litellm
+from litellm.proxy.enterprise.enterprise_hooks.blocked_user_list import (
+    _ENTERPRISE_BlockedUserList,
+)
+from litellm import Router, mock_completion
+from litellm.proxy.utils import ProxyLogging
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.caching import DualCache
+
+
+@pytest.mark.asyncio
+async def test_block_user_check():
+    """
+    - 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
+    """
+    litellm.blocked_user_list = ["user_id_1"]
+
+    blocked_user_obj = _ENTERPRISE_BlockedUserList()
+
+    _api_key = "sk-12345"
+    user_api_key_dict = UserAPIKeyAuth(api_key=_api_key)
+    local_cache = DualCache()
+
+    ## Case 1: blocked user id passed
+    try:
+        await blocked_user_obj.async_pre_call_hook(
+            user_api_key_dict=user_api_key_dict,
+            cache=local_cache,
+            call_type="completion",
+            data={"user_id": "user_id_1"},
+        )
+        pytest.fail(f"Expected call to fail")
+    except Exception as e:
+        pass
+
+    ## Case 2: normal user id passed
+    try:
+        await blocked_user_obj.async_pre_call_hook(
+            user_api_key_dict=user_api_key_dict,
+            cache=local_cache,
+            call_type="completion",
+            data={"user_id": "user_id_2"},
+        )
+    except Exception as e:
+        pytest.fail(f"An error occurred - {str(e)}")