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.