From cf5a3e1645470fec6eed1c04dff76dea710a55fc Mon Sep 17 00:00:00 2001 From: Classic298 <27028174+Classic298@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:32:04 +0200 Subject: [PATCH] perf: write a chat's messages in one transaction instead of one per message Saving a chat rewrote its message rows one at a time. Each message took its own session out of the pool and committed on its own, and the save endpoint hands over the entire merged history rather than only what changed, so a two hundred message chat cost two hundred sessions and two hundred commits on every save. The messages now go through a single select and a single commit. The field mapping for the insert and the update branch moved into two small helpers, so the batch and the single-message path cannot drift apart. Measured on a two hundred message chat with one message edited: 201 queries and 200 transactions before, 2 queries and 1 transaction after, ~149 ms against ~6 ms. Re-saving an unchanged history now costs one select and no writes at all. One behaviour change worth stating: a message the database cannot store used to be skipped on its own, and now costs the rest of that same save. This table is a rebuildable fast path, so the reader falls back to the history on the chat row and re-triggers the backfill, and the next save reconciles everything still present. A per-message retry was tried and dropped, because a commit that lands but still raises would re-apply the usage merge and double the recorded token counts. --- backend/open_webui/models/chat_messages.py | 160 ++++++++++++--------- backend/open_webui/models/chats.py | 21 ++- 2 files changed, 104 insertions(+), 77 deletions(-) diff --git a/backend/open_webui/models/chat_messages.py b/backend/open_webui/models/chat_messages.py index 56c78d9b5b..4977df70c3 100644 --- a/backend/open_webui/models/chat_messages.py +++ b/backend/open_webui/models/chat_messages.py @@ -206,6 +206,66 @@ class ChatMessageModel(BaseModel): class ChatMessageTable: + @staticmethod + def _apply_message_data(message: ChatMessage, data: dict, now: int) -> None: + """Overwrite only the fields the payload carries.""" + if 'role' in data: + message.role = data['role'] + if 'parent_id' in data or 'parentId' in data: + message.parent_id = data.get('parent_id') or data.get('parentId') + if 'content' in data: + message.content = data.get('content') + if 'output' in data: + message.output = data.get('output') + if 'model_id' in data or 'model' in data: + message.model_id = data.get('model_id') or data.get('model') + if 'files' in data: + message.files = data.get('files') + if 'sources' in data: + message.sources = data.get('sources') + if 'embeds' in data: + message.embeds = data.get('embeds') + if 'meta' in data: + message.meta = data.get('meta') + if 'done' in data: + message.done = data['done'] + if 'status_history' in data or 'statusHistory' in data: + message.status_history = data.get('status_history') or data.get('statusHistory') + if 'error' in data: + message.error = data.get('error') + if 'context_summary' in data or 'contextSummary' in data: + message.context_summary = data.get('context_summary') or data.get('contextSummary') + + usage = get_usage(data) + if usage: + existing_usage = normalize_usage(message.usage) + message.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage) + message.updated_at = now + + @staticmethod + def _build_message(composite_id: str, chat_id: str, user_id: str, data: dict, now: int) -> ChatMessage: + return ChatMessage( + id=composite_id, + chat_id=chat_id, + user_id=user_id, + role=data.get('role', 'user'), + parent_id=data.get('parent_id') or data.get('parentId'), + content=data.get('content'), + output=data.get('output'), + model_id=data.get('model_id') or data.get('model'), + files=data.get('files'), + sources=data.get('sources'), + embeds=data.get('embeds'), + meta=data.get('meta'), + done=data.get('done', True), + status_history=data.get('status_history') or data.get('statusHistory'), + error=data.get('error'), + usage=get_usage(data), + context_summary=data.get('context_summary') or data.get('contextSummary'), + created_at=data.get('timestamp', now), + updated_at=now, + ) + async def upsert_message( self, message_id: str, @@ -217,76 +277,46 @@ class ChatMessageTable: """Insert or update a chat message.""" async with get_async_db_context(db) as db: now = int(time.time()) - timestamp = data.get('timestamp', now) - # Use composite ID: {chat_id}-{message_id} composite_id = f'{chat_id}-{message_id}' - existing = await db.get(ChatMessage, composite_id) - if existing: - # Update existing - if 'role' in data: - existing.role = data['role'] - if 'parent_id' in data or 'parentId' in data: - existing.parent_id = data.get('parent_id') or data.get('parentId') - if 'content' in data: - existing.content = data.get('content') - if 'output' in data: - existing.output = data.get('output') - if 'model_id' in data or 'model' in data: - existing.model_id = data.get('model_id') or data.get('model') - if 'files' in data: - existing.files = data.get('files') - if 'sources' in data: - existing.sources = data.get('sources') - if 'embeds' in data: - existing.embeds = data.get('embeds') - if 'meta' in data: - existing.meta = data.get('meta') - if 'done' in data: - existing.done = data.get('done', True) - if 'status_history' in data or 'statusHistory' in data: - existing.status_history = data.get('status_history') or data.get('statusHistory') - if 'error' in data: - existing.error = data.get('error') - if 'context_summary' in data or 'contextSummary' in data: - existing.context_summary = data.get('context_summary') or data.get('contextSummary') - # Extract and normalize usage - usage = get_usage(data) - if usage: - existing_usage = normalize_usage(existing.usage or {}) if existing.usage else {} - existing.usage = existing_usage if usage == existing_usage else merge_usage(existing_usage, usage) - existing.updated_at = now - await db.commit() - return ChatMessageModel.model_validate(existing) + message = await db.get(ChatMessage, composite_id) + if message: + self._apply_message_data(message, data, now) else: - # Insert new - # Extract and normalize usage - usage = get_usage(data) - message = ChatMessage( - id=composite_id, - chat_id=chat_id, - user_id=user_id, - role=data.get('role', 'user'), - parent_id=data.get('parent_id') or data.get('parentId'), - content=data.get('content'), - output=data.get('output'), - model_id=data.get('model_id') or data.get('model'), - files=data.get('files'), - sources=data.get('sources'), - embeds=data.get('embeds'), - meta=data.get('meta'), - done=data.get('done', True), - status_history=data.get('status_history') or data.get('statusHistory'), - error=data.get('error'), - usage=usage, - context_summary=data.get('context_summary') or data.get('contextSummary'), - created_at=timestamp, - updated_at=now, - ) + message = self._build_message(composite_id, chat_id, user_id, data, now) db.add(message) - await db.commit() - return ChatMessageModel.model_validate(message) + + await db.commit() + return ChatMessageModel.model_validate(message) + + async def upsert_messages( + self, + chat_id: str, + user_id: str, + messages: dict[str, dict], + db: AsyncSession | None = None, + ) -> None: + """Insert or update the given messages of one chat.""" + if not messages: + return + + async with get_async_db_context(db) as db: + now = int(time.time()) + result = await db.execute( + select(ChatMessage).filter(ChatMessage.id.in_([f'{chat_id}-{message_id}' for message_id in messages])) + ) + existing_by_id = {row.id: row for row in result.scalars().all()} + + for message_id, data in messages.items(): + composite_id = f'{chat_id}-{message_id}' + message = existing_by_id.get(composite_id) + if message: + self._apply_message_data(message, data, now) + else: + db.add(self._build_message(composite_id, chat_id, user_id, data, now)) + + await db.commit() async def get_message_by_id(self, id: str, db: Optional[AsyncSession] = None) -> Optional[ChatMessageModel]: async with get_async_db_context(db) as db: diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index 258d5b7cc6..0d918729c9 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -959,18 +959,15 @@ class ChatTable: """Write messages to the ``chat_message`` table so future lookups use the fast path. Errors are logged but never raised. """ - for message_id, message in messages.items(): - if not isinstance(message, dict) or not message.get('role'): - continue - try: - await ChatMessages.upsert_message( - message_id=message_id, - chat_id=chat_id, - user_id=user_id, - data=message, - ) - except Exception as e: - log.warning('Backfill failed for message %s in chat %s: %s', message_id, chat_id, e) + writable = { + message_id: message + for message_id, message in messages.items() + if isinstance(message, dict) and message.get('role') + } + try: + await ChatMessages.upsert_messages(chat_id, user_id, writable) + except Exception as e: + log.warning('Backfill failed for chat %s: %s', chat_id, e) async def reconcile_messages_by_chat_id(self, chat_id: str, user_id: str, messages: dict[str, dict]) -> None: """Sync ``chat_message`` rows with the committed JSON blob.