reset chat or reload history after data source change

CHANGELOG.md

@@ -15,4 +15,5 @@ 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))
 - Refactor default data source retriever ([#197](https://github.com/opensearch-project/dashboards-assistant/pull/197))

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/services/__tests__/data_source_service.test.ts

@@ -111,6 +111,22 @@ describe('DataSourceService', () => {
       });
       expect(await dataSource.getDataSourceId$().pipe(first()).toPromise()).toBe('foo');
     });
+
+    it('should not fire change for same data source id', async () => {
+      const { dataSource, defaultDataSourceSelection$ } = setup({
+        dataSourceSelection: new Map(),
+        defaultDataSourceId: 'foo',
+      });
+      const observerFn = jest.fn();
+      dataSource.getDataSourceId$().subscribe(observerFn);
+
+      expect(observerFn).toHaveBeenCalledTimes(1);
+      dataSource.setDataSourceId('foo');
+      expect(observerFn).toHaveBeenCalledTimes(1);
+
+      defaultDataSourceSelection$.next('foo');
+      expect(observerFn).toHaveBeenCalledTimes(1);
+    });
   });
 
   describe('isMDSEnabled', () => {

public/services/data_source_service.ts

@@ -4,7 +4,7 @@
  */
 
 import { BehaviorSubject, Subscription, combineLatest, of } from 'rxjs';
-import { first, map } from 'rxjs/operators';
+import { distinctUntilChanged, first, map } from 'rxjs/operators';
 
 import type { IUiSettingsClient } from '../../../../src/core/public';
 import type { DataSourceOption } from '../../../../src/plugins/data_source_management/public/components/data_source_menu/types';
@@ -74,24 +74,23 @@ export class DataSourceService {
   }
 
   setDataSourceId(newDataSourceId: string | null) {
-    if (this.dataSourceId$.getValue() === newDataSourceId) {
-      return;
-    }
     this.dataSourceId$.next(newDataSourceId);
   }
 
   getDataSourceId$() {
     return combineLatest([
       this.dataSourceId$,
       this.dataSourceManagement?.getDefaultDataSourceId$?.(this.uiSettings) ?? of(null),
-    ]).pipe(
-      map(([selectedDataSourceId, defaultDataSourceId]) => {
-        if (selectedDataSourceId !== null) {
-          return selectedDataSourceId;
-        }
-        return defaultDataSourceId;
-      })
-    );
+    ])
+      .pipe(
+        map(([selectedDataSourceId, defaultDataSourceId]) => {
+          if (selectedDataSourceId !== null) {
+            return selectedDataSourceId;
+          }
+          return defaultDataSourceId;
+        })
+      )
+      .pipe(distinctUntilChanged());
   }
 
   setup({

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();
+  });
 });