diff --git a/backend/open_webui/models/chats.py b/backend/open_webui/models/chats.py index af999f21cd..51f2f121b3 100644 --- a/backend/open_webui/models/chats.py +++ b/backend/open_webui/models/chats.py @@ -459,24 +459,89 @@ class ChatTable: return None return row[0] or 'New Chat' + @staticmethod + def get_unresolved_parent_ids(messages_map: dict) -> set[str]: + """Return parent IDs referenced by messages but absent from the map. + + An empty set means the message graph is fully connected. + """ + return { + msg['parentId'] + for msg in messages_map.values() + if msg.get('parentId') and msg['parentId'] not in messages_map + } + + async def backfill_messages_by_chat_id( + self, chat_id: str, user_id: str, messages: dict[str, dict] + ) -> None: + """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) + async def get_messages_map_by_chat_id(self, id: str) -> Optional[dict]: """Message map for walking history (see ``get_message_list``). - Prefer ``chat_message`` rows to avoid loading the large ``chat`` - JSON blob; fall back to embedded history when no rows exist - (legacy chats). + Prefer ``chat_message`` rows to avoid loading the large embedded + history; fall back to the legacy JSON when no rows exist. + When rows exist but the parent-link graph has gaps (e.g. migration + failures), missing messages are merged from the legacy history + and backfilled so future requests self-heal. """ # Fast path: build from normalized chat_message rows. messages_map = await ChatMessages.get_messages_map_by_chat_id(id) + if messages_map is not None: + unresolved_ids = self.get_unresolved_parent_ids(messages_map) + if not unresolved_ids: + return messages_map + + # Graph has gaps — enrich from the legacy embedded history. + log.info( + 'Chat %s: %d unresolved parent reference(s) in chat_message — ' + 'enriching from legacy history', + id, len(unresolved_ids), + ) + chat = await self.get_chat_by_id(id) + if chat: + history_messages = chat.chat.get('history', {}).get('messages', {}) or {} + missing_messages = { + message_id: history_messages[message_id] + for message_id in unresolved_ids + if message_id in history_messages + } + + if missing_messages: + messages_map.update(missing_messages) + + # Backfill so future requests use the fast path. + await self.backfill_messages_by_chat_id(id, chat.user_id, missing_messages) + return messages_map - # No rows — fall back to the embedded JSON blob for legacy chats. + # No rows — fall back to the legacy embedded history. chat = await self.get_chat_by_id(id) if chat is None: return None - return chat.chat.get('history', {}).get('messages', {}) or {} + history_messages = chat.chat.get('history', {}).get('messages', {}) or {} + + # Backfill so future requests use the fast path. + if history_messages: + await self.backfill_messages_by_chat_id(id, chat.user_id, history_messages) + + return history_messages async def get_message_by_id_and_message_id(self, id: str, message_id: str) -> Optional[dict]: chat = await self.get_chat_by_id(id)