From 55da6224b88508f765fd69b7539499b2e3dcf8b6 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Tue, 24 Dec 2024 18:03:14 -0700
Subject: [PATCH] enh: save status from the backend

---
 backend/open_webui/models/chats.py | 18 ++++++++++++++++++
 backend/open_webui/socket/main.py  |  9 +++++++++
 2 files changed, 27 insertions(+)

diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py
index 8b2766a1f..87bcdb8d3 100644
--- a/backend/open_webui/models/chats.py
+++ b/backend/open_webui/models/chats.py
@@ -235,6 +235,24 @@ class ChatTable:
         chat["history"] = history
         return self.update_chat_by_id(id, chat)
 
+    def add_message_status_to_chat_by_id_and_message_id(
+        self, id: str, message_id: str, status: dict
+    ) -> Optional[ChatModel]:
+        chat = self.get_chat_by_id(id)
+        if chat is None:
+            return None
+
+        chat = chat.chat
+        history = chat.get("history", {})
+
+        if message_id in history.get("messages", {}):
+            status_history = history["messages"][message_id].get("statusHistory", [])
+            status_history.append(status)
+            history["messages"][message_id]["statusHistory"] = status_history
+
+        chat["history"] = history
+        return self.update_chat_by_id(id, chat)
+
     def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
         with get_db() as db:
             # Get the existing chat to share
diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py
index 1412b9139..9c6885ad6 100644
--- a/backend/open_webui/socket/main.py
+++ b/backend/open_webui/socket/main.py
@@ -6,6 +6,8 @@ import time
 
 from open_webui.models.users import Users
 from open_webui.models.channels import Channels
+from open_webui.models.chats import Chats
+
 from open_webui.env import (
     ENABLE_WEBSOCKET_SUPPORT,
     WEBSOCKET_MANAGER,
@@ -259,6 +261,13 @@ def get_event_emitter(request_info):
                 to=session_id,
             )
 
+        if "type" in event_data and event_data["type"] == "status":
+            Chats.add_message_status_to_chat_by_id_and_message_id(
+                request_info["chat_id"],
+                request_info["message_id"],
+                event_data.get("data", {}),
+            )
+
     return __event_emitter__