feat: save/save as copy edited AI message functionality

Author: tfomkin

libs/mobile/chat/features/use-edit-message/src/use-edit-message.ts

@@ -4,10 +4,12 @@ import { FormValues } from '@open-webui-react-native/mobile/shared/utils/form';
 import {
   chatApi,
   ChatResponse,
-  prepareUpdateMessageInChatPayload,
   prepareCompleteChatPayload,
   prepareUpdateMessageToSendPayload,
+  prepareEditAssistantMessagePayload,
+  prepareCopyEditedMessagePayload,
 } from '@open-webui-react-native/shared/data-access/api';
+import { Role } from '@open-webui-react-native/shared/data-access/common';
 import { socketService } from '@open-webui-react-native/shared/data-access/websocket';
 
 interface UseEditMessageProps {
@@ -47,7 +49,7 @@ export const useEditMessage = ({ chat, modelId }: UseEditMessageProps): typeof r
       return;
     }
 
-    const preparedChat = prepareUpdateMessageInChatPayload(chat, editingMessageId, message);
+    const preparedChat = prepareCopyEditedMessagePayload(chat, editingMessageId, message);
 
     await updateChat(preparedChat);
     cancelEditing();
@@ -60,10 +62,18 @@ export const useEditMessage = ({ chat, modelId }: UseEditMessageProps): typeof r
     const chatHistory = chat.chat.history;
     const editedMessage = chatHistory.messages[editingMessageId];
 
-    const preparedChat = prepareUpdateMessageToSendPayload(chat, message, modelId, editedMessage.parentId);
+    let preparedChat: ChatResponse;
+
+    if (editedMessage.role === Role.ASSISTANT) {
+      preparedChat = prepareEditAssistantMessagePayload(chat, editingMessageId, message);
+    } else {
+      preparedChat = prepareUpdateMessageToSendPayload(chat, message, modelId, editedMessage.parentId);
+    }
 
     await updateChat(preparedChat, {
       onSuccess: (data) => {
+        if (editedMessage.role === Role.ASSISTANT) return;
+
         const completePayload = prepareCompleteChatPayload({
           chatId: data.id,
           messageId: data.chat!.history.currentId,

libs/shared/data-access/api/src/lib/chats/utils/index.ts

@@ -14,5 +14,6 @@ export * from './invalidate-archived-chat-list-query';
 export * from './invalidate-search-chats-query';
 export * from './create-messages-list';
 export * from './patch-chat-with-selected-messages';
-export * from './prepare-update-message-in-chat-payload';
 export * from './prepare-update-message-to-send-payload';
+export * from './prepare-edit-assistant-message-payload';
+export * from './prepare-copy-edited-message-payload';

libs/shared/data-access/api/src/lib/chats/utils/prepare-copy-edited-message-payload.ts

@@ -0,0 +1,64 @@
+import { Role } from '@open-webui-react-native/shared/data-access/common';
+import { ChatResponse, Message } from '../models';
+
+export function prepareCopyEditedMessagePayload(
+  oldData: ChatResponse,
+  messageId: string,
+  newContent: string,
+): ChatResponse {
+  const history = oldData.chat.history;
+  const messagesMap = { ...history.messages };
+
+  const target = messagesMap[messageId];
+  if (!target) return oldData;
+
+  const isAIMessage = target.role === Role.ASSISTANT;
+
+  const updatedMessage: Message = {
+    ...target,
+    content: newContent,
+  };
+  messagesMap[messageId] = updatedMessage;
+
+  if (!isAIMessage) {
+    const patchedMessagesList = oldData.chat.messages.map((msg) => (msg.id === messageId ? updatedMessage : msg));
+    const lastAssistantMessage = oldData.chat.history.lastAssistantMessage;
+
+    return {
+      ...oldData,
+      chat: {
+        ...oldData.chat,
+        history: {
+          ...history,
+          messages: messagesMap,
+          lastAssistantMessage,
+        },
+        messages: patchedMessagesList,
+      },
+    };
+  }
+
+  const chain: Array<Message> = [];
+  let pointer: Message | undefined = updatedMessage;
+
+  while (pointer) {
+    chain.unshift(pointer);
+    pointer = pointer.parentId ? messagesMap[pointer.parentId] : undefined;
+  }
+
+  const newCurrentId = updatedMessage.id;
+
+  return {
+    ...oldData,
+    chat: {
+      ...oldData.chat,
+      history: {
+        ...history,
+        messages: messagesMap,
+        currentId: newCurrentId,
+        lastAssistantMessage: updatedMessage,
+      },
+      messages: chain,
+    },
+  };
+}

libs/shared/data-access/api/src/lib/chats/utils/prepare-edit-assistant-message-payload.ts

@@ -0,0 +1,40 @@
+import { Role } from '@open-webui-react-native/shared/data-access/common';
+import { ChatResponse } from '../models';
+
+export function prepareEditAssistantMessagePayload(
+  oldData: ChatResponse,
+  messageId: string,
+  newContent: string,
+): ChatResponse {
+  const history = oldData.chat.history;
+  const messagesMap = { ...history.messages };
+
+  const target = messagesMap[messageId];
+  if (!target || target.role !== Role.ASSISTANT) return oldData;
+
+  const updatedMessage = {
+    ...target,
+    content: newContent,
+    done: true,
+  };
+
+  messagesMap[messageId] = updatedMessage;
+
+  const updatedMessagesList = oldData.chat.messages.map((m) => (m.id === messageId ? updatedMessage : m));
+
+  const lastAssistantMessage =
+    history.lastAssistantMessage?.id === messageId ? updatedMessage : history.lastAssistantMessage;
+
+  return {
+    ...oldData,
+    chat: {
+      ...oldData.chat,
+      history: {
+        ...history,
+        messages: messagesMap,
+        lastAssistantMessage,
+      },
+      messages: updatedMessagesList,
+    },
+  };
+}