From b3493269b3fd2d4bb77d14058abd2f67beed6563 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 13 Mar 2024 12:00:23 -0700
Subject: [PATCH] fix(proxy_server.py): support checking openai user param

---
 docs/my-website/docs/proxy/enterprise.md      | 39 ++++++++++++++++++-
 .../enterprise_hooks/blocked_user_list.py     |  7 ++--
 2 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/docs/my-website/docs/proxy/enterprise.md b/docs/my-website/docs/proxy/enterprise.md
index 93786eff4..4f581846e 100644
--- a/docs/my-website/docs/proxy/enterprise.md
+++ b/docs/my-website/docs/proxy/enterprise.md
@@ -169,11 +169,43 @@ If any call is made to proxy with this user id, it'll be rejected - use this if
 ```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` 
+     blocked_user_list: ["user_id_1", "user_id_2", ...]  # can also be a .txt filepath e.g. `/relative/path/blocked_list.txt` 
 ```
 
 ### How to test
 
+<Tabs>
+
+
+<TabItem value="openai" label="OpenAI Python v1.0.0+">
+
+Set `user=<user_id>` to the user id of the user who might have opted out.
+
+```python
+import openai
+client = openai.OpenAI(
+    api_key="sk-1234",
+    base_url="http://0.0.0.0:4000"
+)
+
+# request sent to model set on litellm proxy, `litellm --model`
+response = client.chat.completions.create(
+    model="gpt-3.5-turbo",
+    messages = [
+        {
+            "role": "user",
+            "content": "this is a test request, write a short poem"
+        }
+    ],
+    user="user_id_1"
+)
+
+print(response)
+```
+</TabItem>
+
+<TabItem value="Curl" label="Curl Request">
+
 ```bash 
 curl --location 'http://0.0.0.0:4000/chat/completions' \
 --header 'Content-Type: application/json' \
@@ -185,11 +217,14 @@ curl --location 'http://0.0.0.0:4000/chat/completions' \
           "content": "what llm are you"
         }
       ],
-      "user_id": "user_id_1" # this is also an openai supported param 
+      "user": "user_id_1" # this is also an openai supported param 
     }
 '
 ```
 
+</TabItem>
+</Tabs>
+
 :::info 
 
 [Suggest a way to improve this](https://github.com/BerriAI/litellm/issues/new/choose)
diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py
index 26a1bd9f7..686fdf1de 100644
--- a/enterprise/enterprise_hooks/blocked_user_list.py
+++ b/enterprise/enterprise_hooks/blocked_user_list.py
@@ -66,12 +66,13 @@ class _ENTERPRISE_BlockedUserList(CustomLogger):
             - 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:
+            if "user_id" in data or "user" in data:
+                user = data.get("user_id", data.get("user", ""))
+                if user in self.blocked_user_list:
                     raise HTTPException(
                         status_code=400,
                         detail={
-                            "error": f"User blocked from making LLM API Calls. User={data['user_id']}"
+                            "error": f"User blocked from making LLM API Calls. User={user}"
                         },
                     )
         except HTTPException as e: