PRD-2089: Regenerate AI message functionality

libs/mobile/chat/features/ai-message-actions/src/lib/component.tsx

@@ -15,6 +15,9 @@ interface AiMessageActionsProps {
   onEditPress: (messageId: string, content: string) => void;
   onSuggestPress: (messageId: string) => void;
   onContinueResponsePress: (messageId: string, content: string) => void;
+  onTryAgain: (messageId: string) => void;
+  onAddDetails: (messageId: string) => void;
+  onMoreConcise: (messageId: string) => void;
   isLast: boolean;
 }
 
@@ -24,6 +27,9 @@ export function AiMessageActions({
   onEditPress,
   onSuggestPress,
   onContinueResponsePress,
+  onTryAgain,
+  onAddDetails,
+  onMoreConcise,
   isLast,
   children,
 }: PropsWithChildren<AiMessageActionsProps>): ReactElement {
@@ -43,22 +49,40 @@ export function AiMessageActions({
     actionsSheetRef.current?.dismiss();
   };
 
-  const handleSuggestPress = (): void => {
+  const handleContinueResponsePress = (): void => {
+    onContinueResponsePress(message.id, message.content);
     actionsSheetRef.current?.dismiss();
+  };
+
+  const openRegenerateActions = (): void => {
+    regenerateActionsSheetRef.current?.present();
+  };
+
+  const runRegenerateAction = (action?: (messageId: string) => void): void => {
+    actionsSheetRef.current?.dismiss();
+
     //NOTE: Small delay ensures sheet is fully closed before showing input
     setTimeout(() => {
-      onSuggestPress(message.id);
+      action?.(message.id);
     }, 100);
+
     regenerateActionsSheetRef.current?.dismiss();
   };
 
-  const handleContinueResponsePress = (): void => {
-    onContinueResponsePress(message.id, message.content);
-    actionsSheetRef.current?.dismiss();
+  const handleSuggestPress = (): void => {
+    runRegenerateAction(onSuggestPress);
   };
 
-  const openRegenerateActions = (): void => {
-    regenerateActionsSheetRef.current?.present();
+  const handleTryAgainPress = (): void => {
+    runRegenerateAction(onTryAgain);
+  };
+
+  const handleAddDetailsPress = (): void => {
+    runRegenerateAction(onAddDetails);
+  };
+
+  const handleMoreConcisePress = (): void => {
+    runRegenerateAction(onMoreConcise);
   };
 
   const actions: Array<ActionSheetItemProps> = compact([
@@ -94,14 +118,17 @@ export function AiMessageActions({
     {
       title: translate('REGENERATE_MESSAGE_ACTION_SHEET.TEXT_TRY_AGAIN'),
       iconName: 'refresh',
+      onPress: handleTryAgainPress,
     },
     {
       title: translate('REGENERATE_MESSAGE_ACTION_SHEET.TEXT_ADD_DETAILS'),
       iconName: 'moreText',
+      onPress: handleAddDetailsPress,
     },
     {
       title: translate('REGENERATE_MESSAGE_ACTION_SHEET.TEXT_MORE_CONCISE'),
       iconName: 'lessText',
+      onPress: handleMoreConcisePress,
     },
   ]);
 

libs/mobile/chat/features/chat/src/lib/component.tsx

@@ -35,6 +35,7 @@ interface ChatProps {
 
 export function Chat({ chatId, selectedModelId, isNewChat, resetToChatsList }: ChatProps): ReactElement {
   const translate = useTranslation('CHAT.CHAT');
+  const translateRegeneratePrompt = useTranslation('CHAT.AI_MESSAGE_ACTIONS.REGENERATE_MESSAGE_ACTION_SHEET');
 
   const [isInputFocusing, setIsInputFocusing] = useState(false); //NOTE: Needs to avoid ChatBottomButton jumping when auto-scrolling after focus
 
@@ -71,7 +72,8 @@ export function Chat({ chatId, selectedModelId, isNewChat, resetToChatsList }: C
     cancelSuggesting,
     control: suggestMessageControl,
     submitSuggestion,
-  } = useSuggestChange();
+    regenerateWithSuggestion,
+  } = useSuggestChange({ chat, modelId: selectedModelId });
 
   const history = chat?.chat.history;
   const isResponseGenerating = !history?.messages[history.currentId].done;
@@ -136,6 +138,23 @@ export function Chat({ chatId, selectedModelId, isNewChat, resetToChatsList }: C
     setActiveInputMode(null);
   };
 
+  const handleQuickSuggestion = (messageId: string, message: string): void => {
+    // NOTE: Quick suggestions should not open the suggest input, they should immediately trigger regeneration
+    void regenerateWithSuggestion(messageId, message);
+  };
+
+  const handleTryAgain = (messageId: string): void => {
+    handleQuickSuggestion(messageId, '');
+  };
+
+  const handleAddDetails = (messageId: string): void => {
+    handleQuickSuggestion(messageId, translateRegeneratePrompt('TEXT_ADD_DETAILS'));
+  };
+
+  const handleMoreConcise = (messageId: string): void => {
+    handleQuickSuggestion(messageId, translateRegeneratePrompt('TEXT_MORE_CONCISE'));
+  };
+
   const onSubmit = (options: Array<ChatGenerationOption>): Promise<void> =>
     handleSubmit(({ inputValue }: FormValues<FormChatInputSchema>): void => {
       if (!selectedModelId) {
@@ -175,6 +194,9 @@ export function Chat({ chatId, selectedModelId, isNewChat, resetToChatsList }: C
           <LazyChatMessagesList
             onEditPress={handleStartEditing}
             onSuggestPress={handleStartSuggesting}
+            onTryAgain={handleTryAgain}
+            onAddDetails={handleAddDetails}
+            onMoreConcise={handleMoreConcise}
             chatId={chatId}
             isInputFocusing={isInputFocusing}
             messages={chat?.chat.messages ?? []}

libs/mobile/chat/features/chat/src/lib/components/messages-list/component.tsx

@@ -31,6 +31,9 @@ interface ChatMessagesListProps {
   isInputFocusing: boolean;
   onEditPress: (messageId: string, content: string) => void;
   onSuggestPress: (messageId: string) => void;
+  onTryAgain: (messageId: string) => void;
+  onAddDetails: (messageId: string) => void;
+  onMoreConcise: (messageId: string) => void;
   history?: ChatHistory;
   messages?: Array<Message>;
   editingMessageId?: string;
@@ -45,6 +48,9 @@ export default function ChatMessagesList({
   isInputFocusing,
   onEditPress,
   onSuggestPress,
+  onTryAgain,
+  onAddDetails,
+  onMoreConcise,
   editingMessageId,
 }: ChatMessagesListProps): ReactElement {
   const listRef = useRef<FlashList<Message>>(null);
@@ -174,6 +180,9 @@ export default function ChatMessagesList({
           onEditPress={onEditPress}
           onSuggestPress={onSuggestPress}
           onContinueResponsePress={handleContinueResponsePress}
+          onTryAgain={onTryAgain}
+          onAddDetails={onAddDetails}
+          onMoreConcise={onMoreConcise}
           isLast={isLast}>
           <ChatAiMessage
             message={message}
@@ -205,6 +214,9 @@ export default function ChatMessagesList({
       getSiblingsInfo,
       modelId,
       onSuggestPress,
+      onTryAgain,
+      onAddDetails,
+      onMoreConcise,
     ],
   );
 

libs/mobile/chat/features/use-suggest-change/src/use-suggest-change.ts

@@ -1,27 +1,46 @@
+import dayjs from 'dayjs';
 import { useState } from 'react';
 import { Control, useForm, UseFormHandleSubmit } from 'react-hook-form';
+import uuid from 'react-native-uuid';
 import { FormValues } from '@open-webui-react-native/mobile/shared/utils/form';
+import {
+  chatApi,
+  ChatResponse,
+  Message,
+  prepareRegeneratePayload,
+  createMessagesList,
+  patchChatWithSelectedMessages,
+} from '@open-webui-react-native/shared/data-access/api';
+import { Role } from '@open-webui-react-native/shared/data-access/common';
 
 export interface SuggestChangeSchema {
   suggestionInputValue: string;
 }
 
+interface UseSuggestChangeProps {
+  chat?: ChatResponse;
+  modelId?: string;
+}
+
 export interface UseSuggestChangeReturn {
   suggestingMessageId?: string;
   control: Control<FormValues<SuggestChangeSchema>>;
   handleSubmit: UseFormHandleSubmit<FormValues<SuggestChangeSchema>>;
   startSuggesting: (messageId: string) => void;
   cancelSuggesting: () => void;
-  submitSuggestion: () => Promise<void>;
+  submitSuggestion: (message: string) => Promise<void>;
+  regenerateWithSuggestion: (messageId: string, message: string) => Promise<void>;
 }
 
-export const useSuggestChange = (): UseSuggestChangeReturn => {
+export const useSuggestChange = ({ chat, modelId }: UseSuggestChangeProps): UseSuggestChangeReturn => {
   const [suggestingMessageId, setSuggestingMessageId] = useState<string>();
 
   const { control, handleSubmit, reset } = useForm<FormValues<SuggestChangeSchema>>({
     defaultValues: { suggestionInputValue: '' },
   });
 
+  const { mutate: completeChat } = chatApi.useCompleteChat();
+
   const startSuggesting = (messageId: string): void => {
     setSuggestingMessageId(messageId);
     reset({ suggestionInputValue: '' });
@@ -32,17 +51,59 @@ export const useSuggestChange = (): UseSuggestChangeReturn => {
     reset({ suggestionInputValue: '' });
   };
 
-  const submitSuggestion = async (): Promise<void> => {
-    //TODO: implement suggestion sending logic
+  const regenerateForMessage = async (messageId: string, message: string): Promise<void> => {
+    if (!chat?.chat || !modelId) return;
+
+    const history = chat.chat.history;
+    const baseAssistant = history.messages[messageId];
+    const parentUserId = baseAssistant.parentId!;
+    const newAssistantId = uuid.v4();
+    const now = dayjs();
+    const timestampSec = Math.floor(now.unix());
+
+    history.messages[newAssistantId] = new Message({
+      id: newAssistantId,
+      timestamp: timestampSec,
+      parentId: parentUserId,
+      role: Role.ASSISTANT,
+      content: '',
+      model: modelId,
+      modelName: modelId,
+      childrenIds: [],
+    });
+
+    history.messages[parentUserId].childrenIds?.push(newAssistantId);
+    history.currentId = newAssistantId;
+
+    const newMessagesList = createMessagesList(history, newAssistantId);
+    patchChatWithSelectedMessages(chat.id, newAssistantId, newMessagesList);
+
+    const payload = prepareRegeneratePayload({
+      chat,
+      messageId: newAssistantId,
+      model: modelId,
+      suggestionPrompt: message,
+    });
+
+    completeChat(payload);
     cancelSuggesting();
   };
 
+  const submitSuggestion = async (message: string): Promise<void> => {
+    if (!suggestingMessageId) return;
+    await regenerateForMessage(suggestingMessageId, message);
+  };
+
+  const regenerateWithSuggestion = async (messageId: string, message: string): Promise<void> =>
+    regenerateForMessage(messageId, message);
+
   return {
     suggestingMessageId,
     control,
     handleSubmit,
     startSuggesting,
     cancelSuggesting,
     submitSuggestion,
+    regenerateWithSuggestion,
   };
 };

libs/shared/data-access/api/src/lib/chats/utils/create-regeneration-messages.ts

@@ -0,0 +1,25 @@
+import { History, Message } from '../models';
+import { createMessagesList } from './create-messages-list';
+
+export function createRegenerationMessages(
+  history: History,
+  userMessageId: string,
+  suggestionPrompt?: string,
+): Array<Message> {
+  const base = createMessagesList(history, userMessageId);
+
+  if (!suggestionPrompt) {
+    return base;
+  }
+
+  return base.map((msg, idx) => {
+    if (idx === base.length - 1 && msg.role === 'user') {
+      return {
+        ...msg,
+        content: `${msg.content}\n\n${suggestionPrompt}`,
+      };
+    }
+
+    return msg;
+  });
+}

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

@@ -18,3 +18,5 @@ 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';
+export * from './prepare-regenerate-payload';
+export * from './create-regeneration-messages';

libs/shared/data-access/api/src/lib/chats/utils/prepare-regenerate-payload.ts

@@ -0,0 +1,30 @@
+import { socketService } from '@open-webui-react-native/shared/data-access/websocket';
+import { ChatGenerationOption } from '../enums';
+import { ChatResponse, CompleteChatRequest } from '../models';
+import { createRegenerationMessages } from './create-regeneration-messages';
+import { prepareCompleteChatPayload } from './prepare-complete-chat-payload';
+
+export function prepareRegeneratePayload({
+  chat,
+  messageId,
+  model,
+  suggestionPrompt,
+  generationOptions,
+}: {
+  chat: ChatResponse;
+  messageId: string;
+  model: string;
+  suggestionPrompt?: string;
+  generationOptions?: Array<ChatGenerationOption>;
+}): CompleteChatRequest {
+  const history = chat.chat.history;
+
+  return prepareCompleteChatPayload({
+    chatId: chat.id,
+    messageId,
+    messages: createRegenerationMessages(history, messageId, suggestionPrompt),
+    sessionId: socketService.socketSessionId,
+    model,
+    generationOptions,
+  });
+}