diff --git a/libs/mobile/chat/features/use-edit-message/src/use-edit-message.ts b/libs/mobile/chat/features/use-edit-message/src/use-edit-message.ts index 1b4e75b..63bf719 100644 --- a/libs/mobile/chat/features/use-edit-message/src/use-edit-message.ts +++ b/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, diff --git a/libs/shared/data-access/api/src/lib/chats/utils/index.ts b/libs/shared/data-access/api/src/lib/chats/utils/index.ts index 100137c..8cb9013 100644 --- a/libs/shared/data-access/api/src/lib/chats/utils/index.ts +++ b/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'; diff --git a/libs/shared/data-access/api/src/lib/chats/utils/prepare-copy-edited-message-payload.ts b/libs/shared/data-access/api/src/lib/chats/utils/prepare-copy-edited-message-payload.ts new file mode 100644 index 0000000..4202525 --- /dev/null +++ b/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), + }, + }; +} diff --git a/libs/shared/data-access/api/src/lib/chats/utils/prepare-edit-assistant-message-payload.ts b/libs/shared/data-access/api/src/lib/chats/utils/prepare-edit-assistant-message-payload.ts new file mode 100644 index 0000000..b971e24 --- /dev/null +++ b/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, + }, + }; +}