perf: collapse chat write helpers into single db session

Author: Algorithm5838

backend/open_webui/models/chats.py

@@ -2,7 +2,7 @@
 import json
 import time
 import uuid
-from typing import Optional
+from typing import Optional, Callable
 
 from sqlalchemy.orm import Session
 from open_webui.internal.db import Base, JSONField, get_db, get_db_context
@@ -462,62 +462,89 @@ def get_message_by_id_and_message_id(self, id: str, message_id: str) -> Optional
         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
+        self,
+        id: str,
+        message_id: str,
+        message: dict,
+        merge_fn: Optional[Callable[[dict], dict]] = None,
     ) -> Optional[ChatModel]:
-        chat = self.get_chat_by_id(id)
-        if chat is None:
-            return None
+        try:
+            with get_db_context() as db:
+                chat_item = db.get(Chat, id)
+                if chat_item is None:
+                    return None
 
-        # Sanitize message content for null characters before upserting
-        if isinstance(message.get('content'), str):
-            message['content'] = sanitize_text_for_db(message['content'])
+                user_id = chat_item.user_id
+                chat = chat_item.chat
+                history = chat.get('history', {})
 
-        user_id = chat.user_id
-        chat = chat.chat
-        history = chat.get('history', {})
+                if merge_fn is not None:
+                    existing_message = history.get('messages', {}).get(message_id, {})
+                    message = merge_fn(existing_message)
 
-        if message_id in history.get('messages', {}):
-            history['messages'][message_id] = {
-                **history['messages'][message_id],
-                **message,
-            }
-        else:
-            history['messages'][message_id] = message
+                # Sanitize message content for null characters before upserting
+                if isinstance(message.get('content'), str):
+                    message['content'] = sanitize_text_for_db(message['content'])
 
-        history['currentId'] = message_id
+                if message_id in history.get('messages', {}):
+                    history['messages'][message_id] = {
+                        **history['messages'][message_id],
+                        **message,
+                    }
+                else:
+                    history['messages'][message_id] = message
 
-        chat['history'] = history
+                history['currentId'] = message_id
+                chat['history'] = history
 
-        # Dual-write to chat_message table
-        try:
-            ChatMessages.upsert_message(
-                message_id=message_id,
-                chat_id=id,
-                user_id=user_id,
-                data=history['messages'][message_id],
-            )
-        except Exception as e:
-            log.warning(f'Failed to write to chat_message table: {e}')
+                chat_item.chat = self._clean_null_bytes(chat)
+                chat_item.title = self._clean_null_bytes(chat['title']) if 'title' in chat else 'New Chat'
+                chat_item.updated_at = int(time.time())
+                db.commit()
+                db.refresh(chat_item)
+
+                # Dual-write to chat_message table
+                try:
+                    ChatMessages.upsert_message(
+                        message_id=message_id,
+                        chat_id=id,
+                        user_id=user_id,
+                        data=history['messages'][message_id],
+                    )
+                except Exception as e:
+                    log.warning(f'Failed to write to chat_message table: {e}')
 
-        return self.update_chat_by_id(id, chat)
+                return ChatModel.model_validate(chat_item)
+        except Exception:
+            return None
 
     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
+        try:
+            with get_db_context() as db:
+                chat_item = db.get(Chat, id)
+                if chat_item is None:
+                    return None
+
+                chat = chat_item.chat
+                history = chat.get('history', {})
 
-        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
 
-        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
 
-        chat['history'] = history
-        return self.update_chat_by_id(id, chat)
+                chat_item.chat = self._clean_null_bytes(chat)
+                chat_item.updated_at = int(time.time())
+                db.commit()
+                db.refresh(chat_item)
+
+                return ChatModel.model_validate(chat_item)
+        except Exception:
+            return None
 
     def add_message_files_by_id_and_message_id(self, id: str, message_id: str, files: list[dict]) -> list[dict]:
         with get_db_context() as db: