reset chat or reload history after data source change

CHANGELOG.md

@@ -15,3 +15,4 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Integrate chatbot with sidecar service ([#164](https://github.com/opensearch-project/dashboards-assistant/pull/164))
 - Add data source service ([#191](https://github.com/opensearch-project/dashboards-assistant/pull/191))
 - Update router and controller to support MDS ([#190](https://github.com/opensearch-project/dashboards-assistant/pull/190))
+- Reset chat and reload history after data source change ([#194](https://github.com/opensearch-project/dashboards-assistant/pull/194))

public/chat_header_button.test.tsx

@@ -48,35 +48,44 @@ jest.mock('./services', () => {
   };
 });
 
+let dataSourceId$: BehaviorSubject<string | null>;
 // mock sidecar open,hide and show
-jest.spyOn(coreContextExports, 'useCore').mockReturnValue({
-  overlays: {
-    // @ts-ignore
-    sidecar: () => {
-      const attachElement = document.createElement('div');
-      attachElement.id = 'sidecar-mock-div';
-      return {
-        open: (mountPoint) => {
-          document.body.appendChild(attachElement);
-          render(<MountWrapper mount={mountPoint} />, {
-            container: attachElement,
-          });
-        },
-        hide: () => {
-          const element = document.getElementById('sidecar-mock-div');
-          if (element) {
-            element.style.display = 'none';
-          }
-        },
-        show: () => {
-          const element = document.getElementById('sidecar-mock-div');
-          if (element) {
-            element.style.display = 'block';
-          }
-        },
-      };
+jest.spyOn(coreContextExports, 'useCore').mockImplementation(() => {
+  dataSourceId$ = new BehaviorSubject<string | null>(null);
+  return {
+    services: {
+      dataSource: {
+        getDataSourceId$: jest.fn(() => dataSourceId$),
+      },
+    },
+    overlays: {
+      // @ts-ignore
+      sidecar: () => {
+        const attachElement = document.createElement('div');
+        attachElement.id = 'sidecar-mock-div';
+        return {
+          open: (mountPoint) => {
+            document.body.appendChild(attachElement);
+            render(<MountWrapper mount={mountPoint} />, {
+              container: attachElement,
+            });
+          },
+          hide: () => {
+            const element = document.getElementById('sidecar-mock-div');
+            if (element) {
+              element.style.display = 'none';
+            }
+          },
+          show: () => {
+            const element = document.getElementById('sidecar-mock-div');
+            if (element) {
+              element.style.display = 'block';
+            }
+          },
+        };
+      },
     },
-  },
+  };
 });
 
 describe('<HeaderChatButton />', () => {
@@ -217,4 +226,23 @@ describe('<HeaderChatButton />', () => {
     });
     expect(screen.getByLabelText('chat input')).not.toHaveFocus();
   });
+
+  it('should call resetChat after data source change', () => {
+    const resetChatMock = jest.fn();
+    expect(dataSourceId$).toBeTruthy();
+    dataSourceId$.next('foo');
+    render(
+      <HeaderChatButton
+        application={applicationServiceMock.createStartContract()}
+        userHasAccess={false}
+        messageRenderers={{}}
+        actionExecutors={{}}
+        assistantActions={{ resetChat: resetChatMock } as AssistantActions}
+        currentAccount={{ username: 'test_user', tenant: 'test_tenant' }}
+      />
+    );
+
+    dataSourceId$.next('bar');
+    expect(resetChatMock).toHaveBeenCalled();
+  });
 });

public/chat_header_button.tsx

@@ -7,6 +7,8 @@ import { EuiBadge, EuiFieldText, EuiIcon } from '@elastic/eui';
 import classNames from 'classnames';
 import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { useEffectOnce } from 'react-use';
+import { skip } from 'rxjs/operators';
+
 import { ApplicationStart, SIDECAR_DOCKED_MODE } from '../../../src/core/public';
 // TODO: Replace with getChrome().logos.Chat.url
 import chatIcon from './assets/chat.svg';
@@ -54,6 +56,9 @@ export const HeaderChatButton = (props: HeaderChatButtonProps) => {
   const flyoutFullScreen = sidecarDockedMode === SIDECAR_DOCKED_MODE.TAKEOVER;
   const flyoutMountPoint = useRef(null);
 
+  const resetChatRef = useRef(props.assistantActions.resetChat);
+  resetChatRef.current = props.assistantActions.resetChat;
+
   useEffectOnce(() => {
     const subscription = props.application.currentAppId$.subscribe((id) => setAppId(id));
     return () => subscription.unsubscribe();
@@ -194,6 +199,18 @@ export const HeaderChatButton = (props: HeaderChatButtonProps) => {
     };
   }, [appId, flyoutVisible, props.assistantActions, registry]);
 
+  useEffect(() => {
+    const subscription = core.services.dataSource
+      .getDataSourceId$()
+      .pipe(skip(1))
+      .subscribe(() => {
+        resetChatRef.current();
+      });
+    return () => {
+      subscription.unsubscribe();
+    };
+  }, [core.services.dataSource]);
+
   return (
     <>
       <div className={classNames('llm-chat-header-icon-wrapper')}>

public/hooks/use_chat_actions.test.tsx

@@ -26,7 +26,14 @@ jest.mock('../services/conversations_service', () => {
 jest.mock('../services/conversation_load_service', () => {
   return {
     ConversationLoadService: jest.fn().mockImplementation(() => {
-      return { load: jest.fn().mockReturnValue({ messages: [], interactions: [] }) };
+      const conversationLoadMock = {
+        abortController: new AbortController(),
+        load: jest.fn().mockImplementation(async () => {
+          conversationLoadMock.abortController = new AbortController();
+          return { messages: [], interactions: [] };
+        }),
+      };
+      return conversationLoadMock;
     }),
   };
 });
@@ -314,6 +321,7 @@ describe('useChatActions hook', () => {
   it('should not handle regenerate response if the regenerate operation has already aborted', async () => {
     const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => ({
       signal: { aborted: true },
+      abort: jest.fn(),
     }));
 
     httpMock.put.mockResolvedValue(SEND_MESSAGE_RESPONSE);
@@ -355,6 +363,7 @@ describe('useChatActions hook', () => {
   it('should not handle regenerate error if the regenerate operation has already aborted', async () => {
     const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => ({
       signal: { aborted: true },
+      abort: jest.fn(),
     }));
     httpMock.put.mockImplementationOnce(() => {
       throw new Error();
@@ -371,4 +380,43 @@ describe('useChatActions hook', () => {
     );
     AbortControllerMock.mockRestore();
   });
+
+  it('should clear chat title, conversation id, flyoutComponent and call reset action', async () => {
+    const { result } = renderHook(() => useChatActions());
+    result.current.resetChat();
+
+    expect(chatContextMock.setConversationId).toHaveBeenLastCalledWith(undefined);
+    expect(chatContextMock.setTitle).toHaveBeenLastCalledWith(undefined);
+    expect(chatContextMock.setFlyoutComponent).toHaveBeenLastCalledWith(null);
+
+    expect(chatStateDispatchMock).toHaveBeenLastCalledWith({ type: 'reset' });
+  });
+
+  it('should abort send action after reset chat', async () => {
+    const abortFn = jest.fn();
+    const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => ({
+      signal: { aborted: true },
+      abort: abortFn,
+    }));
+    const { result } = renderHook(() => useChatActions());
+    await result.current.send(INPUT_MESSAGE);
+    result.current.resetChat();
+
+    expect(abortFn).toHaveBeenCalled();
+    AbortControllerMock.mockRestore();
+  });
+
+  it('should abort load action after reset chat', async () => {
+    const abortFn = jest.fn();
+    const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => ({
+      signal: { aborted: true },
+      abort: abortFn,
+    }));
+    const { result } = renderHook(() => useChatActions());
+    await result.current.loadChat('conversation_id_mock');
+    result.current.resetChat();
+
+    expect(abortFn).toHaveBeenCalled();
+    AbortControllerMock.mockRestore();
+  });
 });

public/hooks/use_chat_actions.tsx

@@ -106,6 +106,15 @@ export const useChatActions = (): AssistantActions => {
     }
   };
 
+  const resetChat = () => {
+    abortControllerRef?.abort();
+    core.services.conversationLoad.abortController?.abort();
+    chatContext.setConversationId(undefined);
+    chatContext.setTitle(undefined);
+    chatContext.setFlyoutComponent(null);
+    chatStateDispatch({ type: 'reset' });
+  };
+
   const openChatUI = () => {
     chatContext.setFlyoutVisible(true);
     chatContext.setSelectedTabId(TAB_ID.CHAT);
@@ -225,5 +234,5 @@ export const useChatActions = (): AssistantActions => {
     }
   };
 
-  return { send, loadChat, executeAction, openChatUI, abortAction, regenerate };
+  return { send, loadChat, resetChat, executeAction, openChatUI, abortAction, regenerate };
 };

public/tabs/history/__tests__/chat_history_page.test.tsx

@@ -4,7 +4,7 @@
  */
 
 import React from 'react';
-import { act, fireEvent, render, waitFor } from '@testing-library/react';
+import { fireEvent, render, waitFor } from '@testing-library/react';
 import { I18nProvider } from '@osd/i18n/react';
 
 import { coreMock } from '../../../../../../src/core/public/mocks';
@@ -17,6 +17,7 @@ import { ConversationsService } from '../../../services/conversations_service';
 import { DataSourceServiceMock } from '../../../services/data_source_service.mock';
 
 import { ChatHistoryPage } from '../chat_history_page';
+import { BehaviorSubject } from 'rxjs';
 
 const mockGetConversationsHttp = () => {
   const http = coreMock.createStart().http;
@@ -35,11 +36,17 @@ const mockGetConversationsHttp = () => {
 const setup = ({
   http = mockGetConversationsHttp(),
   chatContext = {},
+  shouldRefresh = false,
 }: {
   http?: HttpStart;
   chatContext?: { flyoutFullScreen?: boolean };
+  shouldRefresh?: boolean;
 } = {}) => {
-  const dataSourceMock = new DataSourceServiceMock();
+  const dataSourceId$ = new BehaviorSubject<string | null>('foo');
+  const dataSourceMock = {
+    getDataSourceId$: jest.fn(() => dataSourceId$),
+    getDataSourceQuery: jest.fn(() => ({ dataSourceId: dataSourceId$.getValue() })),
+  };
   const useCoreMock = {
     services: {
       ...coreMock.createStart(),
@@ -65,7 +72,7 @@ const setup = ({
 
   const renderResult = render(
     <I18nProvider>
-      <ChatHistoryPage shouldRefresh={false} />
+      <ChatHistoryPage shouldRefresh={shouldRefresh} />
     </I18nProvider>
   );
 
@@ -74,6 +81,7 @@ const setup = ({
     useChatStateMock,
     useChatContextMock,
     renderResult,
+    dataSourceId$,
   };
 };
 
@@ -240,4 +248,30 @@ describe('<ChatHistoryPage />', () => {
       expect(abortMock).toHaveBeenCalled();
     });
   });
+
+  it('should call conversations.reload after data source changed', async () => {
+    const { useCoreMock, dataSourceId$ } = setup({ shouldRefresh: true });
+
+    jest.spyOn(useCoreMock.services.conversations, 'reload');
+
+    expect(useCoreMock.services.conversations.reload).not.toHaveBeenCalled();
+
+    dataSourceId$.next('bar');
+
+    await waitFor(() => {
+      expect(useCoreMock.services.conversations.reload).toHaveBeenCalled();
+    });
+  });
+
+  it('should not call conversations.reload after unmount', async () => {
+    const { useCoreMock, dataSourceId$, renderResult } = setup({ shouldRefresh: true });
+
+    jest.spyOn(useCoreMock.services.conversations, 'reload');
+
+    expect(useCoreMock.services.conversations.reload).not.toHaveBeenCalled();
+    renderResult.unmount();
+
+    dataSourceId$.next('bar');
+    expect(useCoreMock.services.conversations.reload).not.toHaveBeenCalled();
+  });
 });

public/tabs/history/chat_history_page.tsx

@@ -16,10 +16,12 @@ import {
   EuiFlexGroup,
   EuiFlexItem,
 } from '@elastic/eui';
-import React, { useCallback, useEffect, useMemo, useState } from 'react';
+import React, { useCallback, useEffect, useMemo, useRef, useState } from 'react';
 import { FormattedMessage } from '@osd/i18n/react';
-import { useDebounce, useObservable } from 'react-use';
+import { useDebounce, useObservable, useUpdateEffect } from 'react-use';
 import cs from 'classnames';
+import { skip } from 'rxjs/operators';
+
 import { useChatActions, useChatState } from '../../hooks';
 import { useChatContext, useCore } from '../../contexts';
 import { TAB_ID } from '../../utils/constants';
@@ -61,6 +63,8 @@ export const ChatHistoryPage: React.FC<ChatHistoryPageProps> = React.memo((props
   const chatHistories = useMemo(() => conversations?.objects || [], [conversations]);
   const hasNoConversations =
     !debouncedSearchName && !!conversations && conversations.total === 0 && !loading;
+  const shouldRefreshRef = useRef(props.shouldRefresh);
+  shouldRefreshRef.current = props.shouldRefresh;
 
   const handleSearchChange = useCallback((e) => {
     setSearchName(e.target.value);
@@ -96,14 +100,29 @@ export const ChatHistoryPage: React.FC<ChatHistoryPageProps> = React.memo((props
     [searchName]
   );
 
-  useEffect(() => {
-    if (props.shouldRefresh) services.conversations.reload();
+  useUpdateEffect(() => {
+    if (!props.shouldRefresh) {
+      return;
+    }
+    services.conversations.reload();
+    return () => {
+      services.conversations.abortController?.abort();
+    };
   }, [props.shouldRefresh, services.conversations]);
 
   useEffect(() => {
     services.conversations.load(bulkGetOptions);
+    const subscription = services.dataSource
+      .getDataSourceId$()
+      .pipe(skip(1))
+      .subscribe(() => {
+        if (shouldRefreshRef.current) {
+          services.conversations.reload();
+        }
+      });
     return () => {
       services.conversations.abortController?.abort();
+      subscription.unsubscribe();
     };
   }, [services.conversations, bulkGetOptions]);
 

public/types.ts

@@ -29,6 +29,7 @@ export type ActionExecutor = (params: Record<string, unknown>) => void;
 export interface AssistantActions {
   send: (input: IMessage) => Promise<void>;
   loadChat: (conversationId?: string, title?: string) => Promise<void>;
+  resetChat: () => void;
   openChatUI: (conversationId?: string) => void;
   executeAction: (suggestedAction: ISuggestedAction, message: IMessage) => Promise<void>;
   abortAction: (conversationId?: string) => Promise<void>;