Merge pull request elastic#16 from dgieselaar/abort-and-regenerate-re…

…sponses

x-pack/plugins/observability_ai_assistant/kibana.jsonc

@@ -8,7 +8,7 @@
     "browser": true,
     "configPath": ["xpack", "observabilityAIAssistant"],
     "requiredPlugins": ["triggersActionsUi", "actions", "security", "observabilityShared"],
-    "requiredBundles": ["kibanaReact"],
+    "requiredBundles": ["kibanaReact", "kibanaUtils"],
     "optionalPlugins": [],
     "extraPublicDirs": []
   }

x-pack/plugins/observability_ai_assistant/public/components/insight/insight.tsx

@@ -14,6 +14,8 @@ import { MessagePanel } from '../message_panel/message_panel';
 import { MessageText } from '../message_panel/message_text';
 import { InsightBase } from './insight_base';
 import { InsightMissingCredentials } from './insight_missing_credentials';
+import { StopGeneratingButton } from '../stop_generating_button';
+import { RegenerateResponseButton } from '../regenerate_response_button';
 
 function ChatContent({ messages, connectorId }: { messages: Message[]; connectorId: string }) {
   const chat = useChat({ messages, connectorId });
@@ -22,7 +24,21 @@ function ChatContent({ messages, connectorId }: { messages: Message[]; connector
     <MessagePanel
       body={<MessageText content={chat.content ?? ''} loading={chat.loading} />}
       error={chat.error}
-      controls={null}
+      controls={
+        chat.loading ? (
+          <StopGeneratingButton
+            onClick={() => {
+              chat.abort();
+            }}
+          />
+        ) : (
+          <RegenerateResponseButton
+            onClick={() => {
+              chat.regenerate();
+            }}
+          />
+        )
+      }
     />
   );
 }

x-pack/plugins/observability_ai_assistant/public/components/regenerate_response_button.tsx

@@ -11,7 +11,7 @@ import { i18n } from '@kbn/i18n';
 
 export function RegenerateResponseButton(props: Partial<EuiButtonEmptyProps>) {
   return (
-    <EuiButtonEmpty {...props} iconType="sparkles" size="s">
+    <EuiButtonEmpty size="s" {...props} iconType="sparkles">
       {i18n.translate('xpack.observabilityAiAssistant.regenerateResponseButtonLabel', {
         defaultMessage: 'Regenerate',
       })}

x-pack/plugins/observability_ai_assistant/public/components/stop_generating_button.stories.tsx

@@ -0,0 +1,19 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+import { ComponentMeta, ComponentStoryObj } from '@storybook/react';
+import { StopGeneratingButton as Component } from './stop_generating_button';
+
+const meta: ComponentMeta<typeof Component> = {
+  component: Component,
+  title: 'app/Atoms/StopGeneratingButton',
+};
+
+export default meta;
+
+export const StopGeneratingButton: ComponentStoryObj<typeof Component> = {
+  args: {},
+};

x-pack/plugins/observability_ai_assistant/public/components/stop_generating_button.tsx

@@ -0,0 +1,20 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { EuiButtonEmpty, EuiButtonEmptyProps } from '@elastic/eui';
+import React from 'react';
+import { i18n } from '@kbn/i18n';
+
+export function StopGeneratingButton(props: Partial<EuiButtonEmptyProps>) {
+  return (
+    <EuiButtonEmpty size="s" {...props} iconType="stop" color="text">
+      {i18n.translate('xpack.observabilityAiAssistant.stopGeneratingButtonLabel', {
+        defaultMessage: 'Stop generating',
+      })}
+    </EuiButtonEmpty>
+  );
+}

x-pack/plugins/observability_ai_assistant/public/hooks/use_chat.test.ts

@@ -0,0 +1,264 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { useKibana } from '@kbn/kibana-react-plugin/public';
+import { act, renderHook } from '@testing-library/react-hooks';
+import { ChatCompletionResponseMessage } from 'openai';
+import { Observable } from 'rxjs';
+import { AbortError } from '@kbn/kibana-utils-plugin/common';
+import { ObservabilityAIAssistantService } from '../types';
+import { useChat } from './use_chat';
+import { useObservabilityAIAssistant } from './use_observability_ai_assistant';
+
+jest.mock('@kbn/kibana-react-plugin/public');
+jest.mock('./use_observability_ai_assistant');
+
+const WAIT_OPTIONS = { timeout: 5000 };
+
+const mockUseKibana = useKibana as jest.MockedFunction<typeof useKibana>;
+const mockUseObservabilityAIAssistant = useObservabilityAIAssistant as jest.MockedFunction<
+  typeof useObservabilityAIAssistant
+>;
+
+function mockDeltas(deltas: Array<Partial<ChatCompletionResponseMessage>>) {
+  return mockResponse(
+    Promise.resolve(
+      new Observable((subscriber) => {
+        async function simulateDelays() {
+          for (const delta of deltas) {
+            await new Promise<void>((resolve) => {
+              setTimeout(() => {
+                subscriber.next({
+                  choices: [
+                    {
+                      role: 'assistant',
+                      delta,
+                    },
+                  ],
+                });
+                resolve();
+              }, 100);
+            });
+          }
+          subscriber.complete();
+        }
+
+        simulateDelays();
+      })
+    )
+  );
+}
+
+function mockResponse(response: Promise<any>) {
+  mockUseObservabilityAIAssistant.mockReturnValue({
+    chat: jest.fn().mockImplementation(() => {
+      return response;
+    }),
+  } as unknown as ObservabilityAIAssistantService);
+}
+
+describe('useChat', () => {
+  beforeEach(() => {
+    mockUseKibana.mockReturnValue({
+      services: { notifications: { showErrorDialog: jest.fn() } },
+    } as any);
+  });
+
+  it('returns the result of the chat API', async () => {
+    mockDeltas([{ content: 'testContent' }]);
+    const { result, waitFor } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    expect(result.current.loading).toBeTruthy();
+    expect(result.current.error).toBeUndefined();
+    expect(result.current.content).toBeUndefined();
+
+    await waitFor(() => result.current.loading === false, WAIT_OPTIONS);
+
+    expect(result.current.error).toBeUndefined();
+    expect(result.current.content).toBe('testContent');
+  });
+
+  it('handles 4xx and 5xx', async () => {
+    mockResponse(Promise.reject(new Error()));
+    const { result, waitFor } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitFor(() => result.current.loading === false, WAIT_OPTIONS);
+
+    expect(result.current.error).toBeInstanceOf(Error);
+    expect(result.current.content).toBeUndefined();
+
+    expect(mockUseKibana().services.notifications?.showErrorDialog).toHaveBeenCalled();
+  });
+
+  it('handles valid responses but generation errors', async () => {
+    mockResponse(
+      Promise.resolve(
+        new Observable((subscriber) => {
+          subscriber.next({ choices: [{ role: 'assistant', delta: { content: 'foo' } }] });
+          setTimeout(() => {
+            subscriber.error(new Error());
+          }, 100);
+        })
+      )
+    );
+
+    const { result, waitFor } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitFor(() => result.current.loading === false, WAIT_OPTIONS);
+
+    expect(result.current.loading).toBe(false);
+    expect(result.current.error).toBeInstanceOf(Error);
+    expect(result.current.content).toBe('foo');
+
+    expect(mockUseKibana().services.notifications?.showErrorDialog).toHaveBeenCalled();
+  });
+
+  it('handles aborted requests', async () => {
+    mockResponse(
+      Promise.resolve(
+        new Observable((subscriber) => {
+          subscriber.next({ choices: [{ role: 'assistant', delta: { content: 'foo' } }] });
+        })
+      )
+    );
+
+    const { result, waitFor, unmount } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitFor(() => result.current.content === 'foo', WAIT_OPTIONS);
+
+    unmount();
+
+    expect(mockUseKibana().services.notifications?.showErrorDialog).not.toHaveBeenCalled();
+  });
+
+  it('handles regenerations triggered by updates', async () => {
+    mockResponse(
+      Promise.resolve(
+        new Observable((subscriber) => {
+          subscriber.next({ choices: [{ role: 'assistant', delta: { content: 'foo' } }] });
+        })
+      )
+    );
+
+    const { result, waitFor, rerender } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitFor(() => result.current.content === 'foo', WAIT_OPTIONS);
+
+    mockDeltas([{ content: 'bar' }]);
+
+    rerender({ messages: [], connectorId: 'bar' });
+
+    await waitFor(() => result.current.loading === false);
+
+    expect(mockUseKibana().services.notifications?.showErrorDialog).not.toHaveBeenCalled();
+
+    expect(result.current.content).toBe('bar');
+  });
+
+  it('handles streaming updates', async () => {
+    mockDeltas([
+      {
+        content: 'my',
+      },
+      {
+        content: ' ',
+      },
+      {
+        content: 'update',
+      },
+    ]);
+
+    const { result, waitForNextUpdate } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitForNextUpdate(WAIT_OPTIONS);
+
+    expect(result.current.content).toBe('my');
+
+    await waitForNextUpdate(WAIT_OPTIONS);
+
+    expect(result.current.content).toBe('my ');
+
+    await waitForNextUpdate(WAIT_OPTIONS);
+
+    expect(result.current.content).toBe('my update');
+  });
+
+  it('handles user aborts', async () => {
+    mockResponse(
+      Promise.resolve(
+        new Observable((subscriber) => {
+          subscriber.next({ choices: [{ role: 'assistant', delta: { content: 'foo' } }] });
+        })
+      )
+    );
+
+    const { result, waitForNextUpdate } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitForNextUpdate(WAIT_OPTIONS);
+
+    act(() => {
+      result.current.abort();
+    });
+
+    expect(mockUseKibana().services.notifications?.showErrorDialog).not.toHaveBeenCalled();
+
+    expect(result.current.content).toBe('foo');
+    expect(result.current.loading).toBe(false);
+    expect(result.current.error).toBeInstanceOf(AbortError);
+  });
+
+  it('handles user regenerations', async () => {
+    mockResponse(
+      Promise.resolve(
+        new Observable((subscriber) => {
+          subscriber.next({ choices: [{ role: 'assistant', delta: { content: 'foo' } }] });
+        })
+      )
+    );
+
+    const { result, waitForNextUpdate } = renderHook(
+      ({ messages, connectorId }) => useChat({ messages, connectorId }),
+      { initialProps: { messages: [], connectorId: 'myConnectorId' } }
+    );
+
+    await waitForNextUpdate(WAIT_OPTIONS);
+
+    act(() => {
+      mockDeltas([{ content: 'bar' }]);
+      result.current.regenerate();
+    });
+
+    await waitForNextUpdate(WAIT_OPTIONS);
+
+    expect(mockUseKibana().services.notifications?.showErrorDialog).not.toHaveBeenCalled();
+
+    expect(result.current.content).toBe('bar');
+    expect(result.current.loading).toBe(false);
+    expect(result.current.error).toBeUndefined();
+  });
+});