Dashboards: Auto-generate get stuck and quick feedback actions doesn't respond (#83879)

* Update the component only when the response is fully generated

* Fix quick feedback action doesn't respond

* Fix history not displaying after the second click

* Fix the history that moves when regenerating

---------

Co-authored-by: Adela Almasan <88068998+adela-almasan@users.noreply.github.com>

public/app/features/dashboard/components/GenAI/GenAIButton.test.tsx

@@ -173,13 +173,11 @@ describe('GenAIButton', () => {
       await waitFor(() => expect(getByRole('button')).toBeEnabled());
     });
 
-    it('should call onGenerate when the text is generating', async () => {
+    it('should not call onGenerate when the text is generating', async () => {
       const onGenerate = jest.fn();
       setup({ onGenerate, messages: [], eventTrackingSrc: eventTrackingSrc });
 
-      await waitFor(() => expect(onGenerate).toHaveBeenCalledTimes(1));
-
-      expect(onGenerate).toHaveBeenCalledWith('Some incomplete generated text');
+      await waitFor(() => expect(onGenerate).not.toHaveBeenCalledTimes(1));
     });
 
     it('should stop generating when clicking the button', async () => {
@@ -191,6 +189,45 @@ describe('GenAIButton', () => {
 
       expect(setShouldStopMock).toHaveBeenCalledTimes(1);
       expect(setShouldStopMock).toHaveBeenCalledWith(true);
+      expect(onGenerate).not.toHaveBeenCalled();
+    });
+  });
+
+  describe('when it is completed from generating data', () => {
+    const setShouldStopMock = jest.fn();
+
+    beforeEach(() => {
+      jest.mocked(useOpenAIStream).mockReturnValue({
+        messages: [],
+        error: undefined,
+        streamStatus: StreamStatus.COMPLETED,
+        reply: 'Some completed generated text',
+        setMessages: jest.fn(),
+        setStopGeneration: setShouldStopMock,
+        value: {
+          enabled: true,
+          stream: new Observable().subscribe(),
+        },
+      });
+    });
+
+    it('should render improve text ', async () => {
+      setup();
+
+      waitFor(async () => expect(await screen.findByText('Improve')).toBeInTheDocument());
+    });
+
+    it('should enable the button', async () => {
+      setup();
+      waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled());
+    });
+
+    it('should call onGenerate when the text is completed', async () => {
+      const onGenerate = jest.fn();
+      setup({ onGenerate, messages: [], eventTrackingSrc: eventTrackingSrc });
+
+      await waitFor(() => expect(onGenerate).toHaveBeenCalledTimes(1));
+      expect(onGenerate).toHaveBeenCalledWith('Some completed generated text');
     });
   });
 

public/app/features/dashboard/components/GenAI/GenAIButton.tsx

@@ -54,10 +54,11 @@ export const GenAIButton = ({
   } = useOpenAIStream(model, temperature);
 
   const [history, setHistory] = useState<string[]>([]);
-  const [showHistory, setShowHistory] = useState(true);
+  const [showHistory, setShowHistory] = useState(false);
 
   const hasHistory = history.length > 0;
-  const isFirstHistoryEntry = streamStatus === StreamStatus.GENERATING && !hasHistory;
+  const isGenerating = streamStatus === StreamStatus.GENERATING;
+  const isFirstHistoryEntry = !hasHistory;
   const isButtonDisabled = disabled || (value && !value.enabled && !error);
   const reportInteraction = (item: AutoGenerateItem) => reportAutoGenerateInteraction(eventTrackingSrc, item);
 
@@ -69,19 +70,17 @@ export const GenAIButton = ({
         onClickProp?.(e);
         setMessages(typeof messages === 'function' ? messages() : messages);
       } else {
-        if (setShowHistory) {
-          setShowHistory(true);
-        }
+        setShowHistory(true);
       }
     }
 
     const buttonItem = error
       ? AutoGenerateItem.erroredRetryButton
-      : isFirstHistoryEntry
+      : isGenerating
         ? AutoGenerateItem.stopGenerationButton
-        : hasHistory
-          ? AutoGenerateItem.improveButton
-          : AutoGenerateItem.autoGenerateButton;
+        : isFirstHistoryEntry
+          ? AutoGenerateItem.autoGenerateButton
+          : AutoGenerateItem.improveButton;
     reportInteraction(buttonItem);
   };
 
@@ -96,10 +95,10 @@ export const GenAIButton = ({
 
   useEffect(() => {
     // Todo: Consider other options for `"` sanitation
-    if (isFirstHistoryEntry && reply) {
+    if (streamStatus === StreamStatus.COMPLETED && reply) {
       onGenerate(sanitizeReply(reply));
     }
-  }, [streamStatus, reply, onGenerate, isFirstHistoryEntry]);
+  }, [streamStatus, reply, onGenerate]);
 
   useEffect(() => {
     if (streamStatus === StreamStatus.COMPLETED) {
@@ -119,7 +118,7 @@ export const GenAIButton = ({
   };
 
   const getIcon = () => {
-    if (isFirstHistoryEntry) {
+    if (isGenerating) {
       return undefined;
     }
     if (error || (value && !value?.enabled)) {
@@ -135,7 +134,7 @@ export const GenAIButton = ({
       buttonText = 'Retry';
     }
 
-    if (isFirstHistoryEntry) {
+    if (isGenerating) {
       buttonText = STOP_GENERATION_TEXT;
     }
 
@@ -175,9 +174,11 @@ export const GenAIButton = ({
               eventTrackingSrc={eventTrackingSrc}
             />
           }
-          placement="bottom-start"
+          placement="left-start"
           fitContent={true}
-          show={showHistory ? undefined : false}
+          show={showHistory}
+          onClose={() => setShowHistory(false)}
+          onOpen={() => setShowHistory(true)}
         >
           {button}
         </Toggletip>
@@ -189,8 +190,8 @@ export const GenAIButton = ({
 
   return (
     <div className={styles.wrapper}>
-      {isFirstHistoryEntry && <Spinner size="sm" className={styles.spinner} />}
-      {!hasHistory && (
+      {isGenerating && <Spinner size="sm" className={styles.spinner} />}
+      {isFirstHistoryEntry ? (
         <Tooltip
           show={error ? undefined : false}
           interactive
@@ -200,8 +201,9 @@ export const GenAIButton = ({
         >
           {button}
         </Tooltip>
+      ) : (
+        renderButtonWithToggletip()
       )}
-      {hasHistory && renderButtonWithToggletip()}
     </div>
   );
 };

public/app/features/dashboard/components/GenAI/GenAIHistory.tsx

@@ -2,18 +2,7 @@ import { css } from '@emotion/css';
 import React, { useEffect, useState } from 'react';
 
 import { GrafanaTheme2 } from '@grafana/data';
-import {
-  Alert,
-  Button,
-  HorizontalGroup,
-  Icon,
-  IconButton,
-  Input,
-  Text,
-  TextLink,
-  useStyles2,
-  VerticalGroup,
-} from '@grafana/ui';
+import { Alert, Button, Icon, IconButton, Input, Stack, Text, TextLink, useStyles2 } from '@grafana/ui';
 
 import { STOP_GENERATION_TEXT } from './GenAIButton';
 import { GenerationHistoryCarousel } from './GenerationHistoryCarousel';
@@ -100,7 +89,9 @@ export const GenAIHistory = ({
 
   const onGenerateWithFeedback = (suggestion: string | QuickFeedbackType) => {
     if (suggestion !== QuickFeedbackType.Regenerate) {
-      messages = [...messages, ...getFeedbackMessage(history[currentIndex], suggestion)];
+      messages = [...messages, ...getFeedbackMessage(history[currentIndex - 1], suggestion)];
+    } else {
+      messages = [...messages, ...getFeedbackMessage(history[currentIndex - 1], 'Please, regenerate')];
     }
 
     setMessages(messages);
@@ -122,13 +113,11 @@ export const GenAIHistory = ({
   return (
     <div className={styles.container}>
       {showError && (
-        <div>
-          <Alert title="">
-            <VerticalGroup>
-              <div>Sorry, I was unable to complete your request. Please try again.</div>
-            </VerticalGroup>
-          </Alert>
-        </div>
+        <Alert title="">
+          <Stack direction={'column'}>
+            <p>Sorry, I was unable to complete your request. Please try again.</p>
+          </Stack>
+        </Alert>
       )}
 
       <Input
@@ -157,11 +146,11 @@ export const GenAIHistory = ({
         />
       </div>
       <div className={styles.applySuggestion}>
-        <HorizontalGroup justify={'flex-end'}>
+        <Stack justifyContent={'flex-end'} direction={'row'}>
           <Button icon={!isStreamGenerating ? 'check' : 'fa fa-spinner'} onClick={onApply}>
             {isStreamGenerating ? STOP_GENERATION_TEXT : 'Apply'}
           </Button>
-        </HorizontalGroup>
+        </Stack>
       </div>
       <div className={styles.footer}>
         <Icon name="exclamation-circle" aria-label="exclamation-circle" className={styles.infoColor} />
@@ -186,7 +175,7 @@ const getStyles = (theme: GrafanaTheme2) => ({
     display: 'flex',
     flexDirection: 'column',
     width: 520,
-    height: 250,
+    maxHeight: 350,
     // This is the space the footer height
     paddingBottom: 35,
   }),

public/app/features/dashboard/components/GenAI/hooks.ts

@@ -52,6 +52,8 @@ export function useOpenAIStream(
   const [streamStatus, setStreamStatus] = useState<StreamStatus>(StreamStatus.IDLE);
   const [error, setError] = useState<Error>();
   const { error: notifyError } = useAppNotification();
+  // Accumulate response and it will only update the state of the attatched component when the stream is completed.
+  let partialReply = '';
 
   const onError = useCallback(
     (e: Error) => {
@@ -69,6 +71,12 @@ export function useOpenAIStream(
     [messages, model, temperature, notifyError]
   );
 
+  useEffect(() => {
+    if (messages.length > 0) {
+      setReply('');
+    }
+  }, [messages]);
+
   const { error: enabledError, value: enabled } = useAsync(
     async () => await isLLMPluginEnabled(),
     [isLLMPluginEnabled]
@@ -102,9 +110,12 @@ export function useOpenAIStream(
     return {
       enabled,
       stream: stream.subscribe({
-        next: setReply,
+        next: (reply) => {
+          partialReply = reply;
+        },
         error: onError,
         complete: () => {
+          setReply(partialReply);
           setStreamStatus(StreamStatus.COMPLETED);
           setTimeout(() => {
             setStreamStatus(StreamStatus.IDLE);