From 6efca03a8f1613a1494b63fa16378f7f1c68ef85 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Fri, 13 Dec 2024 22:51:43 -0800
Subject: [PATCH] refac

---
 backend/open_webui/functions.py        | 1 +
 backend/open_webui/socket/main.py      | 5 +++++
 backend/open_webui/utils/chat.py       | 2 ++
 backend/open_webui/utils/middleware.py | 1 +
 4 files changed, 9 insertions(+)

diff --git a/backend/open_webui/functions.py b/backend/open_webui/functions.py
index d424d4663..9c241432a 100644
--- a/backend/open_webui/functions.py
+++ b/backend/open_webui/functions.py
@@ -230,6 +230,7 @@ async def generate_function_chat_completion(
             "role": user.role,
         },
         "__metadata__": metadata,
+        "__request__": request,
     }
     extra_params["__tools__"] = get_tools(
         request,
diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py
index ba5eeb6ae..c0f45c9a0 100644
--- a/backend/open_webui/socket/main.py
+++ b/backend/open_webui/socket/main.py
@@ -173,6 +173,11 @@ async def user_count(sid):
     await sio.emit("user-count", {"count": len(USER_POOL.items())})
 
 
+@sio.on("chat")
+async def chat(sid, data):
+    print("chat", sid, SESSION_POOL[sid], data)
+
+
 @sio.event
 async def disconnect(sid):
     if sid in SESSION_POOL:
diff --git a/backend/open_webui/utils/chat.py b/backend/open_webui/utils/chat.py
index 676a4a203..56904d1d8 100644
--- a/backend/open_webui/utils/chat.py
+++ b/backend/open_webui/utils/chat.py
@@ -237,6 +237,7 @@ async def chat_completed(request: Request, form_data: dict, user: Any):
                 "__id__": filter_id,
                 "__event_emitter__": __event_emitter__,
                 "__event_call__": __event_call__,
+                "__request__": request,
             }
 
             # Add extra params in contained in function signature
@@ -334,6 +335,7 @@ async def chat_action(request: Request, action_id: str, form_data: dict, user: A
                 "__id__": sub_action_id if sub_action_id is not None else action_id,
                 "__event_emitter__": __event_emitter__,
                 "__event_call__": __event_call__,
+                "__request__": request,
             }
 
             # Add extra params in contained in function signature
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 9c5186e00..1d2bc2b99 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -382,6 +382,7 @@ async def process_chat_payload(request, form_data, user, model):
             "role": user.role,
         },
         "__metadata__": metadata,
+        "__request__": request,
     }
 
     # Initialize events to store additional event to be sent to the client