From 8e535a799f62040384fee362e1db46e9d85447cd Mon Sep 17 00:00:00 2001
From: Ilya Bondar <ilya_bondar@epam.com>
Date: Wed, 11 Dec 2024 21:55:39 +0100
Subject: [PATCH] feat(chat): regenerate on empty assistant message (#2783)
 (#2788)

---
 .../chat-e2e/src/tests/workWithModels.test.ts | 66 +++++--------------
 .../src/ui/webElements/chatMessages.ts        |  2 +-
 .../src/ui/webElements/sendMessage.ts         |  1 +
 .../Chat/ChatInput/ChatInputMessage.tsx       | 16 +++--
 .../Chat/ChatInput/SendMessageButton.tsx      | 12 +++-
 5 files changed, 40 insertions(+), 57 deletions(-)

diff --git a/apps/chat-e2e/src/tests/workWithModels.test.ts b/apps/chat-e2e/src/tests/workWithModels.test.ts
index 59ae33143e..cdb6d96c2f 100644
--- a/apps/chat-e2e/src/tests/workWithModels.test.ts
+++ b/apps/chat-e2e/src/tests/workWithModels.test.ts
@@ -11,7 +11,7 @@ import {
   MockedChatApiResponseBodies,
   Theme,
 } from '@/src/testData';
-import { Cursors, Overflow, Styles } from '@/src/ui/domData';
+import { Overflow } from '@/src/ui/domData';
 import { keys } from '@/src/ui/keyboard';
 import { GeneratorUtil, ModelsUtil } from '@/src/utils';
 import { expect } from '@playwright/test';
@@ -121,7 +121,7 @@ dialTest(
 
         await expect
           .soft(
-            chatMessages.regenerate.getElementLocator(),
+            sendMessage.regenerate.getElementLocator(),
             ExpectedMessages.regenerateIsAvailable,
           )
           .toBeVisible();
@@ -384,8 +384,8 @@ dialTest(
     chat,
     setTestIds,
     chatMessages,
+    baseAssertion,
     sendMessage,
-    page,
     tooltip,
     localStorageManager,
     iconApiHelper,
@@ -431,59 +431,25 @@ dialTest(
           undefined,
           expectedModelIcon,
         );
-        await expect
-          .soft(
-            chatMessages.regenerate.getElementLocator(),
-            ExpectedMessages.regenerateIsAvailable,
-          )
-          .toBeVisible();
+        await baseAssertion.assertElementState(
+          chatMessages.regenerate,
+          'visible',
+        );
+        await baseAssertion.assertElementState(
+          sendMessage.regenerate,
+          'visible',
+        );
       },
     );
 
     await dialTest.step(
       'Hover over Send button and verify it is disabled and tooltip is shown',
       async () => {
-        for (let i = 1; i <= 2; i++) {
-          if (i === 2) {
-            const messagesCountBefore =
-              await chatMessages.chatMessages.getElementsCount();
-            await sendMessage.messageInput.fillInInput('   ');
-            await page.keyboard.press(keys.enter);
-            const messagesCountAfter =
-              await chatMessages.chatMessages.getElementsCount();
-            expect
-              .soft(
-                messagesCountBefore === messagesCountAfter,
-                ExpectedMessages.messageCountIsCorrect,
-              )
-              .toBeTruthy();
-          }
-          const isSendMessageBtnEnabled =
-            await sendMessage.sendMessageButton.isElementEnabled();
-          expect
-            .soft(
-              isSendMessageBtnEnabled,
-              ExpectedMessages.sendMessageButtonDisabled,
-            )
-            .toBeFalsy();
-
-          await sendMessage.sendMessageButton.hoverOver();
-          const sendBtnCursor =
-            await sendMessage.sendMessageButton.getComputedStyleProperty(
-              Styles.cursor,
-            );
-          expect
-            .soft(
-              sendBtnCursor[0],
-              ExpectedMessages.sendButtonCursorIsNotAllowed,
-            )
-            .toBe(Cursors.notAllowed);
-
-          const tooltipContent = await tooltip.getContent();
-          expect
-            .soft(tooltipContent, ExpectedMessages.tooltipContentIsValid)
-            .toBe(ExpectedConstants.regenerateResponseTooltip);
-        }
+        await sendMessage.regenerate.hoverOver();
+        const tooltipContent = await tooltip.getContent();
+        expect
+          .soft(tooltipContent, ExpectedMessages.tooltipContentIsValid)
+          .toBe(ExpectedConstants.regenerateResponseTooltip);
       },
     );
 
diff --git a/apps/chat-e2e/src/ui/webElements/chatMessages.ts b/apps/chat-e2e/src/ui/webElements/chatMessages.ts
index f68c39c63e..9a14c0c308 100644
--- a/apps/chat-e2e/src/ui/webElements/chatMessages.ts
+++ b/apps/chat-e2e/src/ui/webElements/chatMessages.ts
@@ -39,7 +39,7 @@ export class ChatMessages extends BaseElement {
       ChatSelectors.chatMessage,
     );
 
-  public regenerate = new BaseElement(this.page, ChatSelectors.regenerate);
+  public regenerate = this.getChildElementBySelector(ChatSelectors.regenerate);
   private inputAttachments!: InputAttachments;
 
   getInputAttachments(): InputAttachments {
diff --git a/apps/chat-e2e/src/ui/webElements/sendMessage.ts b/apps/chat-e2e/src/ui/webElements/sendMessage.ts
index 6619471933..076ea10e04 100644
--- a/apps/chat-e2e/src/ui/webElements/sendMessage.ts
+++ b/apps/chat-e2e/src/ui/webElements/sendMessage.ts
@@ -62,6 +62,7 @@ export class SendMessage extends BaseElement {
   public stopGenerating = this.getChildElementBySelector(
     SendMessageSelectors.stopGenerating,
   );
+  public regenerate = this.getChildElementBySelector(ChatSelectors.regenerate);
 
   public proceedGenerating = this.getChildElementBySelector(
     SendMessageSelectors.proceedGenerating,
diff --git a/apps/chat/src/components/Chat/ChatInput/ChatInputMessage.tsx b/apps/chat/src/components/Chat/ChatInput/ChatInputMessage.tsx
index 26fe2c3ea9..e3cbc9b63f 100644
--- a/apps/chat/src/components/Chat/ChatInput/ChatInputMessage.tsx
+++ b/apps/chat/src/components/Chat/ChatInput/ChatInputMessage.tsx
@@ -116,7 +116,8 @@ export const ChatInputMessage = ({
   const isModelsLoaded = useAppSelector(ModelsSelectors.selectIsModelsLoaded);
   const isChatFullWidth = useAppSelector(UISelectors.selectIsChatFullWidth);
 
-  const isError = isLastAssistantMessageEmpty || isMessageError;
+  const shouldRegenerate =
+    isLastMessageError || (isLastAssistantMessageEmpty && !messageIsStreaming);
 
   const selectedModels = useAppSelector(
     ConversationsSelectors.selectSelectedConversationsModels,
@@ -166,7 +167,7 @@ export const ChatInputMessage = ({
   ]);
   const isSendDisabled =
     isReplay ||
-    isError ||
+    isMessageError ||
     isInputEmpty ||
     !isModelsLoaded ||
     isUploadingFilePresent ||
@@ -205,6 +206,11 @@ export const ChatInputMessage = ({
       return;
     }
 
+    if (shouldRegenerate) {
+      onRegenerate();
+      return;
+    }
+
     if (isSendDisabled) {
       return;
     }
@@ -234,6 +240,7 @@ export const ChatInputMessage = ({
     }
   }, [
     messageIsStreaming,
+    shouldRegenerate,
     isSendDisabled,
     dispatch,
     onSend,
@@ -244,6 +251,7 @@ export const ChatInputMessage = ({
     setContent,
     textareaRef,
     onStopConversation,
+    onRegenerate,
   ]);
 
   const handleKeyDown = useCallback(
@@ -380,7 +388,7 @@ export const ChatInputMessage = ({
     if (isReplay) {
       return t('Please continue replay to continue working with conversation');
     }
-    if (isError) {
+    if (shouldRegenerate) {
       return t('Regenerate response');
     }
     if (isUploadingFilePresent) {
@@ -437,7 +445,7 @@ export const ChatInputMessage = ({
         />
         <ChatControls
           showReplayControls={showReplayControls}
-          onSend={isLastMessageError ? onRegenerate : handleSend}
+          onSend={shouldRegenerate ? onRegenerate : handleSend}
           tooltip={tooltipContent()}
           isLastMessageError={isLastMessageError}
           isLoading={isLoading}
diff --git a/apps/chat/src/components/Chat/ChatInput/SendMessageButton.tsx b/apps/chat/src/components/Chat/ChatInput/SendMessageButton.tsx
index c73ed2e893..b26f5fa31e 100644
--- a/apps/chat/src/components/Chat/ChatInput/SendMessageButton.tsx
+++ b/apps/chat/src/components/Chat/ChatInput/SendMessageButton.tsx
@@ -37,11 +37,19 @@ export const SendMessageButton = ({
     ConversationsSelectors.selectIsConversationsStreaming,
   );
 
-  if (isLastMessageError) {
+  const isLastAssistantMessageEmpty = useAppSelector(
+    ConversationsSelectors.selectIsLastAssistantMessageEmpty,
+  );
+
+  if (
+    isLastMessageError ||
+    (isLastAssistantMessageEmpty && !messageIsStreaming)
+  ) {
     return (
       <button
         className={classNames(
-          'absolute top-[calc(50%_-_12px)] rounded text-error hover:text-accent-primary',
+          'absolute top-[calc(50%_-_12px)] rounded hover:text-accent-primary',
+          isLastMessageError && 'text-error',
           isOverlay ? 'right-3' : 'right-4',
         )}
         onClick={onSend}