diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py
index 87bcdb8d3..18f802afe 100644
--- a/backend/open_webui/models/chats.py
+++ b/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]:
diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py
index c0bb3c662..f3e9a033e 100644
--- a/backend/open_webui/socket/main.py
+++ b/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__
 
 
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index bb6c56a3b..c476e8de8 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/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: