From 536b9edec00547d5b84ef2e6ea0f929c054c1333 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Mon, 24 Aug 2026 18:35:04 -0400 Subject: [PATCH] refac --- backend/open_webui/models/chats.py | 25 ++++++++++++++++++++++ backend/open_webui/socket/main.py | 34 ++++++++++-------------------- 2 files changed, 36 insertions(+), 23 deletions(-) diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index d8edef363e..243803de1d 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -6,6 +6,7 @@ import logging import re import time import uuid +from typing import Literal # local imports from open_webui.env import ENABLE_ADMIN_CHAT_ACCESS @@ -1070,6 +1071,30 @@ class ChatTable: return chat.chat.get('history', {}).get('messages', {}).get(message_id, {}) + async def get_message_metadata_list( + self, + chat_id: str, + message_id: str, + metadata_key: Literal['files', 'sources', 'embeds'], + ) -> list[dict]: + """Read one list-valued message field without rebuilding the whole history.""" + async with get_async_db_context() as db: + # Read the column directly; some stored rows cannot be validated as full ChatMessageModel objects. + result = await db.execute( + select(getattr(ChatMessage, metadata_key)).where(ChatMessage.id == f'{chat_id}-{message_id}') + ) + metadata_row = result.first() + + if metadata_row is not None: + return metadata_row[0] or [] + + chat = await self.get_chat_by_id(chat_id) + if chat is None: + return [] + + message = chat.chat.get('history', {}).get('messages', {}).get(message_id, {}) + return message.get(metadata_key) or [] + async def upsert_message_to_chat_by_id_and_message_id( self, id: str, message_id: str, message: dict, *, touch: bool = True ) -> ChatModel | None: diff --git a/backend/open_webui/socket/main.py b/backend/open_webui/socket/main.py index 7739636887..7f06ecc2d0 100644 --- a/backend/open_webui/socket/main.py +++ b/backend/open_webui/socket/main.py @@ -1074,15 +1074,12 @@ async def get_event_emitter(request_info, update_db=True): embeds = event_payload.get('embeds', []) if not event_payload.get('replace', False): - message = await Chats.get_message_by_id_and_message_id( - request_info['chat_id'], - request_info['message_id'], - ) - embeds.extend(message.get('embeds', [])) + existing_embeds = await Chats.get_message_metadata_list(chat_id, message_id, 'embeds') + embeds.extend(existing_embeds) await Chats.upsert_message_to_chat_by_id_and_message_id( - request_info['chat_id'], - request_info['message_id'], + chat_id, + message_id, { 'embeds': embeds, }, @@ -1090,17 +1087,13 @@ async def get_event_emitter(request_info, update_db=True): ) elif event_type == 'files': - message = await Chats.get_message_by_id_and_message_id( - request_info['chat_id'], - request_info['message_id'], - ) - files = event_data.get('data', {}).get('files', []) - files.extend(message.get('files', [])) + existing_files = await Chats.get_message_metadata_list(chat_id, message_id, 'files') + files.extend(existing_files) await Chats.upsert_message_to_chat_by_id_and_message_id( - request_info['chat_id'], - request_info['message_id'], + chat_id, + message_id, { 'files': files, }, @@ -1110,17 +1103,12 @@ async def get_event_emitter(request_info, update_db=True): elif event_type in ('source', 'citation'): data = event_data.get('data', {}) if data.get('type') is None: - message = await Chats.get_message_by_id_and_message_id( - request_info['chat_id'], - request_info['message_id'], - ) - - sources = message.get('sources', []) + sources = await Chats.get_message_metadata_list(chat_id, message_id, 'sources') sources.append(data) await Chats.upsert_message_to_chat_by_id_and_message_id( - request_info['chat_id'], - request_info['message_id'], + chat_id, + message_id, { 'sources': sources, },