Auto-generate: Error handling and monitoring (#75468)

* Error handling for the hook to expose `error` as part of the state
* Monitor errors at the hook level to catch common OpenAI issues
* Allow retry when the request fails
* Optimize onGenerate number of calls
* Re-enable component tests

.betterer.results

@@ -3129,9 +3129,6 @@ exports[`better eslint`] = {
     "public/app/features/dashboard/components/DeleteDashboard/DeleteDashboardModal.tsx:5381": [
       [0, 0, 0, "Styles should be written using objects.", "0"]
     ],
-    "public/app/features/dashboard/components/GenAI/GenAIButton.tsx:5381": [
-      [0, 0, 0, "Styles should be written using objects.", "0"]
-    ],
     "public/app/features/dashboard/components/GenAI/llms/openai.ts:5381": [
       [0, 0, 0, "Do not use any type assertions.", "0"],
       [0, 0, 0, "Do not use any type assertions.", "1"]

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

@@ -2,11 +2,13 @@ import { fireEvent, render, screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import React from 'react';
 import { Router } from 'react-router-dom';
+import { Observable } from 'rxjs';
 
 import { selectors } from '@grafana/e2e-selectors';
 import { locationService } from '@grafana/runtime';
 
 import { GenAIButton, GenAIButtonProps } from './GenAIButton';
+import { useOpenAIStream } from './hooks';
 import { Role, isLLMPluginEnabled } from './utils';
 
 jest.mock('./utils', () => ({
@@ -41,7 +43,7 @@ describe('GenAIButton', () => {
       jest.mocked(isLLMPluginEnabled).mockResolvedValue(false);
     });
 
-    it('should renders text ', async () => {
+    it('should render text ', async () => {
       const { getByText } = setup();
       waitFor(() => expect(getByText('Auto-generate')).toBeInTheDocument());
     });
@@ -68,12 +70,23 @@ describe('GenAIButton', () => {
     });
   });
 
-  describe('when LLM plugin is properly configured', () => {
+  describe('when LLM plugin is properly configured, so it is enabled', () => {
+    const setMessagesMock = jest.fn();
     beforeEach(() => {
       jest.mocked(isLLMPluginEnabled).mockResolvedValue(true);
+      jest.mocked(useOpenAIStream).mockReturnValue({
+        error: undefined,
+        isGenerating: false,
+        reply: 'Some completed genereated text',
+        setMessages: setMessagesMock,
+        value: {
+          enabled: true,
+          stream: new Observable().subscribe(),
+        },
+      });
     });
 
-    it('should renders text ', async () => {
+    it('should render text ', async () => {
       setup();
 
       waitFor(async () => expect(await screen.findByText('Auto-generate')).toBeInTheDocument());
@@ -84,47 +97,133 @@ describe('GenAIButton', () => {
       waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled());
     });
 
-    it.skip('disables the button while generating', async () => {
-      const { getByText, getByRole } = setup();
-      const generateButton = getByText('Auto-generate');
-
-      // Click the button
-      await fireEvent.click(generateButton);
-
-      // The loading text should be visible and the button disabled
-      expect(await screen.findByText('Generating')).toBeVisible();
-      await waitFor(() => expect(getByRole('button')).toBeDisabled());
-    });
-
-    it.skip('handles the response and re-enables the button', async () => {
-      const onGenerate = jest.fn();
-      setup({ onGenerate, messages: [] });
+    it('should send the configured messages', async () => {
+      setup({ onGenerate, messages: [{ content: 'Generate X', role: 'system' as Role }] });
       const generateButton = await screen.findByRole('button');
 
       // Click the button
       await fireEvent.click(generateButton);
       await waitFor(() => expect(generateButton).toBeEnabled());
-      await waitFor(() => expect(onGenerate).toHaveBeenCalledTimes(1));
 
       // Wait for the loading state to be resolved
-      expect(onGenerate).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledWith([{ content: 'Generate X', role: 'system' as Role }]);
     });
 
-    it.skip('should call the LLM service with the messages configured and the right temperature', async () => {
+    it('should call the onClick callback', async () => {
       const onGenerate = jest.fn();
+      const onClick = jest.fn();
       const messages = [{ content: 'Generate X', role: 'system' as Role }];
-      setup({ onGenerate, messages, temperature: 3 });
+      setup({ onGenerate, messages, temperature: 3, onClick });
 
       const generateButton = await screen.findByRole('button');
       await fireEvent.click(generateButton);
 
-      await waitFor(() => expect(mockedUseOpenAiStreamState.setMessages).toHaveBeenCalledTimes(1));
-      await waitFor(() =>
-        expect(mockedUseOpenAiStreamState.setMessages).toHaveBeenCalledWith(messages, expect.any(Function), 3)
-      );
+      await waitFor(() => expect(onClick).toHaveBeenCalledTimes(1));
+    });
+  });
+
+  describe('when it is generating data', () => {
+    beforeEach(() => {
+      jest.mocked(isLLMPluginEnabled).mockResolvedValue(true);
+      jest.mocked(useOpenAIStream).mockReturnValue({
+        error: undefined,
+        isGenerating: true,
+        reply: 'Some incompleted generated text',
+        setMessages: jest.fn(),
+        value: {
+          enabled: true,
+          stream: new Observable().subscribe(),
+        },
+      });
     });
 
-    it.skip('should call the onClick callback', async () => {
+    it('should render loading text ', async () => {
+      setup();
+
+      waitFor(async () => expect(await screen.findByText('Auto-generate')).toBeInTheDocument());
+    });
+
+    it('should enable the button', async () => {
+      setup();
+      waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled());
+    });
+
+    it('disables the button while generating', async () => {
+      const { getByText, getByRole } = setup();
+      const generateButton = getByText('Generating');
+
+      // The loading text should be visible and the button disabled
+      expect(generateButton).toBeVisible();
+      await waitFor(() => expect(getByRole('button')).toBeDisabled());
+    });
+
+    it('should call onGenerate when the text is generating', async () => {
+      const onGenerate = jest.fn();
+      setup({ onGenerate, messages: [] });
+
+      await waitFor(() => expect(onGenerate).toHaveBeenCalledTimes(1));
+
+      expect(onGenerate).toHaveBeenCalledWith('Some incompleted generated text');
+    });
+  });
+
+  describe('when there is an error generating data', () => {
+    const setMessagesMock = jest.fn();
+    beforeEach(() => {
+      jest.mocked(isLLMPluginEnabled).mockResolvedValue(true);
+      jest.mocked(useOpenAIStream).mockReturnValue({
+        error: new Error('Something went wrong'),
+        isGenerating: false,
+        reply: '',
+        setMessages: setMessagesMock,
+        value: {
+          enabled: true,
+          stream: new Observable().subscribe(),
+        },
+      });
+    });
+
+    it('should render error state text', async () => {
+      setup();
+
+      waitFor(async () => expect(await screen.findByText('Retry')).toBeInTheDocument());
+    });
+
+    it('should enable the button', async () => {
+      setup();
+      waitFor(async () => expect(await screen.findByRole('button')).toBeEnabled());
+    });
+
+    it('should retry when clicking', async () => {
+      const onGenerate = jest.fn();
+      const messages = [{ content: 'Generate X', role: 'system' as Role }];
+      const { getByText } = setup({ onGenerate, messages, temperature: 3 });
+      const generateButton = getByText('Retry');
+
+      await fireEvent.click(generateButton);
+
+      expect(setMessagesMock).toHaveBeenCalledTimes(1);
+      expect(setMessagesMock).toHaveBeenCalledWith(messages);
+    });
+
+    it('should display the error message as tooltip', async () => {
+      const { getByRole, getByTestId } = setup();
+
+      // Wait for the check to be completed
+      const button = getByRole('button');
+      await userEvent.hover(button);
+
+      const tooltip = await waitFor(() => getByTestId(selectors.components.Tooltip.container));
+      expect(tooltip).toBeVisible();
+
+      // The tooltip keeps interactive to be able to click the link
+      await userEvent.hover(tooltip);
+      expect(tooltip).toBeVisible();
+      expect(tooltip).toHaveTextContent('Something went wrong');
+    });
+
+    it('should call the onClick callback', async () => {
       const onGenerate = jest.fn();
       const onClick = jest.fn();
       const messages = [{ content: 'Generate X', role: 'system' as Role }];

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

@@ -1,5 +1,5 @@
 import { css } from '@emotion/css';
-import React from 'react';
+import React, { useEffect } from 'react';
 
 import { GrafanaTheme2 } from '@grafana/data';
 import { Button, Spinner, useStyles2, Link, Tooltip } from '@grafana/ui';
@@ -34,43 +34,66 @@ export const GenAIButton = ({
   const styles = useStyles2(getStyles);
 
   // TODO: Implement error handling (use error object from hook)
-  const { setMessages, reply, isGenerating, value } = useOpenAIStream(OPEN_AI_MODEL, temperature);
+  const { setMessages, reply, isGenerating, value, error } = useOpenAIStream(OPEN_AI_MODEL, temperature);
 
   const onClick = (e: React.MouseEvent<HTMLButtonElement>) => {
     onClickProp?.(e);
     setMessages(messages);
   };
 
-  // Todo: Consider other options for `"` sanitation
-  if (isGenerating) {
-    onGenerate(reply.replace(/^"|"$/g, ''));
-  }
+  useEffect(() => {
+    // Todo: Consider other options for `"` sanitation
+    if (isGenerating && reply) {
+      onGenerate(reply.replace(/^"|"$/g, ''));
+    }
+  }, [isGenerating, reply, onGenerate]);
 
   const getIcon = () => {
+    if (error || !value?.enabled) {
+      return 'exclamation-circle';
+    }
     if (isGenerating) {
       return undefined;
     }
-    if (!value?.enabled) {
-      return 'exclamation-circle';
-    }
     return 'ai';
   };
 
+  const getTooltipContent = () => {
+    if (error) {
+      return `Unexpected error: ${error.message}`;
+    }
+    if (!value?.enabled) {
+      return (
+        <span>
+          The LLM plugin is not correctly configured. See your <Link href={`/plugins/grafana-llm-app`}>settings</Link>{' '}
+          and enable your plugin.
+        </span>
+      );
+    }
+    return '';
+  };
+
+  const getText = () => {
+    if (error) {
+      return 'Retry';
+    }
+
+    return !isGenerating ? text : loadingText;
+  };
+
   return (
     <div className={styles.wrapper}>
       {isGenerating && <Spinner size={14} />}
-      <Tooltip
-        show={value?.enabled ? false : undefined}
-        interactive
-        content={
-          <span>
-            The LLM plugin is not correctly configured. See your <Link href={`/plugins/grafana-llm-app`}>settings</Link>{' '}
-            and enable your plugin.
-          </span>
-        }
-      >
-        <Button icon={getIcon()} onClick={onClick} fill="text" size="sm" disabled={isGenerating || !value?.enabled}>
-          {!isGenerating ? text : loadingText}
+      <Tooltip show={value?.enabled && !error ? false : undefined} interactive content={getTooltipContent()}>
+        <Button
+          icon={getIcon()}
+          onClick={onClick}
+          fill="text"
+          size="sm"
+          disabled={isGenerating || (!value?.enabled && !error)}
+          variant={error ? 'destructive' : 'primary'}
+        >
+          {getText()}
         </Button>
       </Tooltip>
     </div>
@@ -78,7 +101,7 @@ export const GenAIButton = ({
 };
 
 const getStyles = (theme: GrafanaTheme2) => ({
-  wrapper: css`
-    display: flex;
-  `,
+  wrapper: css({
+    display: 'flex',
+  }),
 });

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

@@ -2,6 +2,9 @@ import { useState } from 'react';
 import { useAsync } from 'react-use';
 import { Subscription } from 'rxjs';
 
+import { logError } from '@grafana/runtime';
+import { useAppNotification } from 'app/core/copy/appNotification';
+
 import { openai } from './llms';
 import { isLLMPluginEnabled, OPEN_AI_MODEL } from './utils';
 
@@ -20,11 +23,11 @@ export function useOpenAIStream(
   error: Error | undefined;
   value:
     | {
-        enabled: boolean;
+        enabled: boolean | undefined;
         stream?: undefined;
       }
     | {
-        enabled: boolean;
+        enabled: boolean | undefined;
         stream: Subscription;
       }
     | undefined;
@@ -33,10 +36,11 @@ export function useOpenAIStream(
   const [messages, setMessages] = useState<Message[]>([]);
   // The latest reply from the LLM.
   const [reply, setReply] = useState('');
-
   const [isGenerating, setIsGenerating] = useState(false);
+  const [error, setError] = useState<Error>();
+  const { error: notifyError } = useAppNotification();
 
-  const { error, value } = useAsync(async () => {
+  const { error: asyncError, value } = useAsync(async () => {
     // Check if the LLM plugin is enabled and configured.
     // If not, we won't be able to make requests, so return early.
     const enabled = await isLLMPluginEnabled();
@@ -48,6 +52,7 @@ export function useOpenAIStream(
     }
 
     setIsGenerating(true);
+    setError(undefined);
     // Stream the completions. Each element is the next stream chunk.
     const stream = openai
       .streamChatCompletions({
@@ -63,29 +68,31 @@ export function useOpenAIStream(
         // functionality to update state, e.g. recording when the stream
         // has completed.
         // The operator decision tree on the rxjs website is a useful resource:
-        // https://rxjs.dev/operator-decision-tree.
+        // https://rxjs.dev/operator-decision-tree.)
       );
     // Subscribe to the stream and update the state for each returned value.
     return {
       enabled,
       stream: stream.subscribe({
         next: setReply,
-        error: (e) => {
-          console.log('The backend for the stream returned an error and nobody has implemented error handling yet!');
-          console.log(e);
+        error: (e: Error) => {
+          setIsGenerating(false);
+          setMessages([]);
+          setError(e);
+          notifyError('OpenAI Error', `${e.message}`);
+          logError(e, { messages: JSON.stringify(messages), model, temperature: String(temperature) });
         },
         complete: () => {
           setIsGenerating(false);
           setMessages([]);
+          setError(undefined);
         },
       }),
     };
   }, [messages]);
 
-  if (error) {
-    // TODO: handle errors.
-    console.log('An error occurred');
-    console.log(error.message);
+  if (asyncError) {
+    setError(asyncError);
   }
 
   return {