From 3f0b3ea90e6dc0a9e8a8bde302412d45c4d5f4eb Mon Sep 17 00:00:00 2001
From: Duc Nguyen <duc.nguyen@wealthsimple.com>
Date: Wed, 30 Oct 2024 11:02:56 -0400
Subject: [PATCH] feat: Add option to forward user info as headers to OpenAI
 API

---
 backend/open_webui/apps/audio/main.py  |  7 +++++++
 backend/open_webui/apps/images/main.py |  7 +++++++
 backend/open_webui/apps/openai/main.py | 21 +++++++++++++++++++++
 backend/open_webui/config.py           |  7 +++++++
 4 files changed, 42 insertions(+)

diff --git a/backend/open_webui/apps/audio/main.py b/backend/open_webui/apps/audio/main.py
index b138f82de..2b06e382f 100644
--- a/backend/open_webui/apps/audio/main.py
+++ b/backend/open_webui/apps/audio/main.py
@@ -23,6 +23,7 @@ from open_webui.config import (
     AUDIO_TTS_VOICE,
     AUDIO_TTS_AZURE_SPEECH_REGION,
     AUDIO_TTS_AZURE_SPEECH_OUTPUT_FORMAT,
+    OPENAI_FORWARD_USER_INFO_AS_HEADERS,
     CACHE_DIR,
     CORS_ALLOW_ORIGIN,
     WHISPER_MODEL,
@@ -248,6 +249,12 @@ async def speech(request: Request, user=Depends(get_verified_user)):
         headers["Authorization"] = f"Bearer {app.state.config.TTS_OPENAI_API_KEY}"
         headers["Content-Type"] = "application/json"
 
+        if OPENAI_FORWARD_USER_INFO_AS_HEADERS:
+            headers["X-OpenWebUI-User-Name"] = user.name
+            headers["X-OpenWebUI-User-Id"] = user.id
+            headers["X-OpenWebUI-User-Email"] = user.email
+            headers["X-OpenWebUI-User-Role"] = user.role
+
         try:
             body = body.decode("utf-8")
             body = json.loads(body)
diff --git a/backend/open_webui/apps/images/main.py b/backend/open_webui/apps/images/main.py
index d3cbd4d91..bb85b809e 100644
--- a/backend/open_webui/apps/images/main.py
+++ b/backend/open_webui/apps/images/main.py
@@ -25,6 +25,7 @@ from open_webui.config import (
     COMFYUI_WORKFLOW,
     COMFYUI_WORKFLOW_NODES,
     CORS_ALLOW_ORIGIN,
+    OPENAI_FORWARD_USER_INFO_AS_HEADERS,
     ENABLE_IMAGE_GENERATION,
     IMAGE_GENERATION_ENGINE,
     IMAGE_GENERATION_MODEL,
@@ -456,6 +457,12 @@ async def image_generations(
             headers["Authorization"] = f"Bearer {app.state.config.OPENAI_API_KEY}"
             headers["Content-Type"] = "application/json"
 
+            if OPENAI_FORWARD_USER_INFO_AS_HEADERS:
+                headers["X-OpenWebUI-User-Name"] = user.name
+                headers["X-OpenWebUI-User-Id"] = user.id
+                headers["X-OpenWebUI-User-Email"] = user.email
+                headers["X-OpenWebUI-User-Role"] = user.role
+
             data = {
                 "model": (
                     app.state.config.MODEL
diff --git a/backend/open_webui/apps/openai/main.py b/backend/open_webui/apps/openai/main.py
index 3647977ca..756f5d937 100644
--- a/backend/open_webui/apps/openai/main.py
+++ b/backend/open_webui/apps/openai/main.py
@@ -16,6 +16,7 @@ from open_webui.config import (
     MODEL_FILTER_LIST,
     OPENAI_API_BASE_URLS,
     OPENAI_API_KEYS,
+    OPENAI_FORWARD_USER_INFO_AS_HEADERS,
     AppConfig,
 )
 from open_webui.env import (
@@ -140,6 +141,11 @@ async def speech(request: Request, user=Depends(get_verified_user)):
         if "openrouter.ai" in app.state.config.OPENAI_API_BASE_URLS[idx]:
             headers["HTTP-Referer"] = "https://openwebui.com/"
             headers["X-Title"] = "Open WebUI"
+        if OPENAI_FORWARD_USER_INFO_AS_HEADERS:
+            headers["X-OpenWebUI-User-Name"] = user.name
+            headers["X-OpenWebUI-User-Id"] = user.id
+            headers["X-OpenWebUI-User-Email"] = user.email
+            headers["X-OpenWebUI-User-Role"] = user.role
         r = None
         try:
             r = requests.post(
@@ -325,6 +331,11 @@ async def get_models(url_idx: Optional[int] = None, user=Depends(get_verified_us
         headers = {}
         headers["Authorization"] = f"Bearer {key}"
         headers["Content-Type"] = "application/json"
+        if OPENAI_FORWARD_USER_INFO_AS_HEADERS:
+            headers["X-OpenWebUI-User-Name"] = user.name
+            headers["X-OpenWebUI-User-Id"] = user.id
+            headers["X-OpenWebUI-User-Email"] = user.email
+            headers["X-OpenWebUI-User-Role"] = user.role
 
         r = None
 
@@ -437,6 +448,11 @@ async def generate_chat_completion(
     if "openrouter.ai" in app.state.config.OPENAI_API_BASE_URLS[idx]:
         headers["HTTP-Referer"] = "https://openwebui.com/"
         headers["X-Title"] = "Open WebUI"
+    if OPENAI_FORWARD_USER_INFO_AS_HEADERS:
+        headers["X-OpenWebUI-User-Name"] = user.name
+        headers["X-OpenWebUI-User-Id"] = user.id
+        headers["X-OpenWebUI-User-Email"] = user.email
+        headers["X-OpenWebUI-User-Role"] = user.role
 
     r = None
     session = None
@@ -505,6 +521,11 @@ async def proxy(path: str, request: Request, user=Depends(get_verified_user)):
     headers = {}
     headers["Authorization"] = f"Bearer {key}"
     headers["Content-Type"] = "application/json"
+    if OPENAI_FORWARD_USER_INFO_AS_HEADERS:
+        headers["X-OpenWebUI-User-Name"] = user.name
+        headers["X-OpenWebUI-User-Id"] = user.id
+        headers["X-OpenWebUI-User-Email"] = user.email
+        headers["X-OpenWebUI-User-Role"] = user.role
 
     r = None
     session = None
diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py
index 9d1bd72d8..30d8b5463 100644
--- a/backend/open_webui/config.py
+++ b/backend/open_webui/config.py
@@ -422,6 +422,7 @@ OAUTH_ADMIN_ROLES = PersistentConfig(
 )
 
 
+
 def load_oauth_providers():
     OAUTH_PROVIDERS.clear()
     if GOOGLE_CLIENT_ID.value and GOOGLE_CLIENT_SECRET.value:
@@ -657,6 +658,12 @@ except Exception:
 
 OPENAI_API_BASE_URL = "https://api.openai.com/v1"
 
+OPENAI_FORWARD_USER_INFO_AS_HEADERS = PersistentConfig(
+    "OPENAI_FORWARD_USER_INFO_AS_HEADERS",
+    "oauth.OPENAI_FORWARD_USER_INFO_AS_HEADERS",
+    os.environ.get("OPENAI_FORWARD_USER_INFO_AS_HEADERS", "False").lower() == "true",
+)
+
 ####################################
 # WEBUI
 ####################################