From cd86161f3364e49116182e894318a6b6a199f0e2 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Tue, 24 Dec 2024 19:34:56 -0700
Subject: [PATCH] refac

---
 backend/open_webui/utils/middleware.py | 32 ++++++++++++++++++--------
 src/lib/components/chat/Chat.svelte    | 12 +++-------
 2 files changed, 25 insertions(+), 19 deletions(-)

diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 9e55596da..e13730b39 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -742,8 +742,6 @@ async def process_chat_payload(request, form_data, metadata, user, model):
             }
         )
 
-    print(f"form_data, events")
-
     return form_data, events
 
 
@@ -810,22 +808,36 @@ async def process_chat_response(request, response, user, events, metadata, tasks
 
                     try:
                         data = json.loads(data)
-                        value = (
-                            data.get("choices", [])[0].get("delta", {}).get("content")
-                        )
 
-                        if value:
-                            content = f"{content}{value}"
-
-                            # Save message in the database
+                        if "selected_model_id" in data:
                             Chats.upsert_message_to_chat_by_id_and_message_id(
                                 metadata["chat_id"],
                                 metadata["message_id"],
                                 {
-                                    "content": content,
+                                    "selectedModelId": data["selected_model_id"],
                                 },
                             )
 
+                        else:
+
+                            value = (
+                                data.get("choices", [])[0]
+                                .get("delta", {})
+                                .get("content")
+                            )
+
+                            if value:
+                                content = f"{content}{value}"
+
+                                # Save message in the database
+                                Chats.upsert_message_to_chat_by_id_and_message_id(
+                                    metadata["chat_id"],
+                                    metadata["message_id"],
+                                    {
+                                        "content": content,
+                                    },
+                                )
+
                     except Exception as e:
                         done = "data: [DONE]" in line
                         title = Chats.get_chat_title_by_id(metadata["chat_id"])
diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte
index 9ede63d9f..c3fbefab8 100644
--- a/src/lib/components/chat/Chat.svelte
+++ b/src/lib/components/chat/Chat.svelte
@@ -1052,7 +1052,7 @@
 	};
 
 	const chatCompletionEventHandler = async (data, message, chatId) => {
-		const { id, done, choices, sources, selectedModelId, error, usage } = data;
+		const { id, done, choices, sources, selected_model_id, error, usage } = data;
 
 		if (error) {
 			await handleOpenAIError(error, message);
@@ -1060,12 +1060,6 @@
 
 		if (sources) {
 			message.sources = sources;
-			// Only remove status if it was initially set
-			if (model?.info?.meta?.knowledge ?? false) {
-				message.statusHistory = message.statusHistory.filter(
-					(status) => status.action !== 'knowledge_search'
-				);
-			}
 		}
 
 		if (choices) {
@@ -1103,8 +1097,8 @@
 			}
 		}
 
-		if (selectedModelId) {
-			message.selectedModelId = selectedModelId;
+		if (selected_model_id) {
+			message.selectedModelId = selected_model_id;
 			message.arena = true;
 		}