refac

backend/open_webui/models/chats.py

@@ -212,6 +212,15 @@ class ChatTable:
 
         return chat.chat.get("history", {}).get("messages", {}) or {}
 
+    def get_message_by_id_and_message_id(
+        self, id: str, message_id: str
+    ) -> Optional[dict]:
+        chat = self.get_chat_by_id(id)
+        if chat is None:
+            return None
+
+        return chat.chat.get("history", {}).get("messages", {}).get(message_id, {})
+
     def upsert_message_to_chat_by_id_and_message_id(
         self, id: str, message_id: str, message: dict
     ) -> Optional[ChatModel]:

backend/open_webui/socket/main.py

@@ -292,6 +292,34 @@ def get_event_emitter(request_info):
                 event_data.get("data", {}),
             )
 
+        if "type" in event_data and event_data["type"] == "message":
+            message = Chats.get_message_by_id_and_message_id(
+                request_info["chat_id"],
+                request_info["message_id"],
+            )
+
+            content = message.get("content", "")
+            content += event_data.get("data", {}).get("content", "")
+
+            Chats.upsert_message_to_chat_by_id_and_message_id(
+                request_info["chat_id"],
+                request_info["message_id"],
+                {
+                    "content": content,
+                },
+            )
+
+        if "type" in event_data and event_data["type"] == "replace":
+            content = event_data.get("data", {}).get("content", "")
+
+            Chats.upsert_message_to_chat_by_id_and_message_id(
+                request_info["chat_id"],
+                request_info["message_id"],
+                {
+                    "content": content,
+                },
+            )
+
     return __event_emitter__
 
 

backend/open_webui/utils/middleware.py

@@ -929,9 +929,10 @@ async def process_chat_response(
 
         # Handle as a background task
         async def post_response_handler(response, events):
-
-            assistant_message = get_last_assistant_message(form_data["messages"])
-            content = assistant_message if assistant_message else ""
+            message = Chats.get_message_by_id_and_message_id(
+                metadata["chat_id"], metadata["message_id"]
+            )
+            content = message.get("content", "") if message else ""
 
             try:
                 for event in events: