Merge pull request #8 from RonasIT/PRD-2007-edit-ai-message-logic

PRD-2007: Edit AI message logic

Author: veliseev93

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

@@ -4,10 +4,14 @@ import { FormValues } from '@open-webui-react-native/mobile/shared/utils/form';
 import {
   chatApi,
   ChatResponse,
-  prepareUpdateMessageInChatPayload,
   prepareCompleteChatPayload,
   prepareUpdateMessageToSendPayload,
+  prepareEditAssistantMessagePayload,
+  prepareCopyEditedMessagePayload,
+  prepareUpdateMessageInChatPayload,
+  History,
 } 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 {
@@ -32,6 +36,12 @@ export const useEditMessage = ({ chat, modelId }: UseEditMessageProps): typeof r
   const { mutateAsync: updateChat, isPending: isChatUpdating } = chatApi.useUpdate();
   const { mutate: completeChat } = chatApi.useCompleteChat();
 
+  const isAssistantMessage = (chatHistory: History, messageId: string): boolean => {
+    const editedMessage = chatHistory.messages[messageId];
+
+    return editedMessage.role === Role.ASSISTANT;
+  };
+
   const startEditing = (messageId: string, content: string): void => {
     setEditingMessageId(messageId);
     reset({ editMessageInputValue: content });
@@ -47,7 +57,11 @@ export const useEditMessage = ({ chat, modelId }: UseEditMessageProps): typeof r
       return;
     }
 
-    const preparedChat = prepareUpdateMessageInChatPayload(chat, editingMessageId, message);
+    const chatHistory = chat.chat.history;
+
+    const preparedChat: ChatResponse = isAssistantMessage(chatHistory, editingMessageId)
+      ? prepareCopyEditedMessagePayload(chat, editingMessageId, message)
+      : prepareUpdateMessageInChatPayload(chat, editingMessageId, message);
 
     await updateChat(preparedChat);
     cancelEditing();
@@ -60,10 +74,14 @@ 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);
+    const preparedChat: ChatResponse = isAssistantMessage(chatHistory, editingMessageId)
+      ? prepareEditAssistantMessagePayload(chat, editingMessageId, message)
+      : 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

@@ -16,3 +16,5 @@ 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,55 @@
+import uuid from 'react-native-uuid';
+import { Role } from '@open-webui-react-native/shared/data-access/common';
+import { ChatResponse, History, Message } from '../models';
+import { createMessagesList } from './create-messages-list';
+
+export function prepareCopyEditedMessagePayload(
+  oldData: ChatResponse,
+  messageId: string,
+  newContent: string,
+): ChatResponse {
+  const history = oldData.chat.history;
+  const messagesMap = { ...history.messages };
+
+  const original = messagesMap[messageId];
+  if (!original || original.role !== Role.ASSISTANT) return oldData;
+
+  const newAssistantMessageId = uuid.v4();
+
+  const parentId = original.parentId ?? original.id;
+
+  const newAssistantMessage: Message = {
+    ...original,
+    id: newAssistantMessageId,
+    parentId,
+    childrenIds: [],
+    content: newContent,
+    done: true,
+  };
+
+  if (parentId && messagesMap[parentId]) {
+    const parent = messagesMap[parentId];
+    messagesMap[parentId] = {
+      ...parent,
+      childrenIds: [...(parent.childrenIds ?? []), newAssistantMessageId],
+    };
+  }
+
+  messagesMap[newAssistantMessageId] = newAssistantMessage;
+
+  const updatedHistory: History = {
+    ...history,
+    currentId: newAssistantMessageId,
+    messages: messagesMap,
+    lastAssistantMessage: newAssistantMessage,
+  };
+
+  return {
+    ...oldData,
+    chat: {
+      ...oldData.chat,
+      history: updatedHistory,
+      messages: createMessagesList(updatedHistory, newAssistantMessageId),
+    },
+  };
+}

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,
+    },
+  };
+}