From 90f76d24ecda8e6d31727697a74d1a10dc0f3387 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Wed, 19 Nov 2025 02:16:09 -0500
Subject: [PATCH] refac

---
 backend/open_webui/models/chats.py     | 21 +++++++++++++++++
 backend/open_webui/utils/middleware.py | 31 +++++++++++++-------------
 2 files changed, 37 insertions(+), 15 deletions(-)

diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py
index c559932bcd..f1607f0707 100644
--- a/backend/open_webui/models/chats.py
+++ b/backend/open_webui/models/chats.py
@@ -297,6 +297,27 @@ class ChatTable:
         chat["history"] = history
         return self.update_chat_by_id(id, chat)
 
+    def add_message_files_by_id_and_message_id(
+        self, id: str, message_id: str, files: list[dict]
+    ) -> list[dict]:
+        chat = self.get_chat_by_id(id)
+        if chat is None:
+            return None
+
+        chat = chat.chat
+        history = chat.get("history", {})
+
+        message_files = []
+
+        if message_id in history.get("messages", {}):
+            message_files = history["messages"][message_id].get("files", [])
+            message_files = message_files + files
+            history["messages"][message_id]["files"] = message_files
+
+        chat["history"] = history
+        self.update_chat_by_id(id, chat)
+        return message_files
+
     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/utils/middleware.py b/backend/open_webui/utils/middleware.py
index a3e34e746b..2af14e4233 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -742,11 +742,11 @@ def get_last_images(message_list):
     return images
 
 
-def extract_and_process_urls(delta_images, request, metadata, user) -> list[str]:
+def get_image_urls(delta_images, request, metadata, user) -> list[str]:
     if not isinstance(delta_images, list):
         return []
 
-    processed_urls = []
+    image_urls = []
     for img in delta_images:
         if not isinstance(img, dict) or img.get("type") != "image_url":
             continue
@@ -758,9 +758,9 @@ def extract_and_process_urls(delta_images, request, metadata, user) -> list[str]
         if url.startswith("data:image/png;base64"):
             url = get_image_url_from_base64(request, url, metadata, user)
 
-        processed_urls.append(url)
+        image_urls.append(url)
 
-    return processed_urls
+    return image_urls
 
 
 async def chat_image_generation_handler(
@@ -2602,22 +2602,23 @@ async def process_chat_response(
                                                             "arguments"
                                                         ] += delta_arguments
 
-                                    processed_image_urls = extract_and_process_urls(
+                                    image_urls = get_image_urls(
                                         delta.get("images", []), request, metadata, user
                                     )
-                                    if processed_image_urls:
+                                    if image_urls:
+                                        message_files = Chats.add_message_files_by_id_and_message_id(
+                                            metadata["chat_id"],
+                                            metadata["message_id"],
+                                            [
+                                                {"type": "image", "url": url}
+                                                for url in image_urls
+                                            ],
+                                        )
+
                                         await event_emitter(
                                             {
                                                 "type": "files",
-                                                "data": {
-                                                    "files": [
-                                                        {
-                                                            "type": "image",
-                                                            "url": url,
-                                                        }
-                                                        for url in processed_image_urls
-                                                    ]
-                                                },
+                                                "data": {"files": message_files},
                                             }
                                         )