From bf3c807047ed8641d188bdc9cc64b6f4947cbf50 Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Tue, 22 Jul 2025 11:38:47 +0400
Subject: [PATCH] refac

---
 backend/open_webui/main.py             | 9 ++++++++-
 backend/open_webui/retrieval/utils.py  | 4 +++-
 backend/open_webui/utils/middleware.py | 8 --------
 3 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py
index 595d551d75..f89db29d74 100644
--- a/backend/open_webui/main.py
+++ b/backend/open_webui/main.py
@@ -1404,7 +1404,6 @@ async def chat_completion(
         form_data, metadata, events = await process_chat_payload(
             request, form_data, user, metadata, model
         )
-
     except Exception as e:
         log.debug(f"Error processing chat payload: {e}")
         if metadata.get("chat_id") and metadata.get("message_id"):
@@ -1424,6 +1423,14 @@ async def chat_completion(
 
     try:
         response = await chat_completion_handler(request, form_data, user)
+        if metadata.get("chat_id") and metadata.get("message_id"):
+            Chats.upsert_message_to_chat_by_id_and_message_id(
+                metadata["chat_id"],
+                metadata["message_id"],
+                {
+                    "model": model_id,
+                },
+            )
 
         return await process_chat_response(
             request, response, form_data, user, metadata, model, events, tasks
diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py
index 9158f8536e..7edc3c35ac 100644
--- a/backend/open_webui/retrieval/utils.py
+++ b/backend/open_webui/retrieval/utils.py
@@ -508,7 +508,9 @@ def get_sources_from_items(
             # Note Attached
             note = Notes.get_note_by_id(item.get("id"))
 
-            if user.role == "admin" or has_access(user.id, "read", note.access_control):
+            if note and (
+                user.role == "admin" or has_access(user.id, "read", note.access_control)
+            ):
                 # User has access to the note
                 query_result = {
                     "documents": [[note.data.get("content", {}).get("md", "")]],
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 7000d37863..b00a30c667 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -1381,14 +1381,6 @@ async def process_chat_response(
         task_id = str(uuid4())  # Create a unique task ID.
         model_id = form_data.get("model", "")
 
-        Chats.upsert_message_to_chat_by_id_and_message_id(
-            metadata["chat_id"],
-            metadata["message_id"],
-            {
-                "model": model_id,
-            },
-        )
-
         def split_content_and_whitespace(content):
             content_stripped = content.rstrip()
             original_whitespace = (