feat: support streaming output

Signed-off-by: SuZhou-Joe <suzhou@amazon.com>

Author: SuZhou-Joe

common/types/chat_saved_object_attributes.ts

@@ -87,3 +87,17 @@ export interface SendResponse {
   messages: IMessage[];
   interactions: Interaction[];
 }
+
+export type StreamChunk =
+  | {
+      type: 'patch';
+      body: Pick<SendResponse, 'interactions' | 'messages'>;
+    }
+  | {
+      type: 'error';
+      body: string;
+    }
+  | {
+      type: 'metadata';
+      body: Partial<SendResponse>;
+    };

common/utils/stream/serializer.ts

@@ -0,0 +1,19 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { StreamChunk } from '../../types/chat_saved_object_attributes';
+
+const separators = `\n`;
+
+export const streamSerializer = (chunk: StreamChunk): string => {
+  return `${JSON.stringify(chunk)}${separators}`;
+};
+
+export const sreamDeserializer = (content: string): StreamChunk[] => {
+  return content
+    .split(separators)
+    .filter((item) => item)
+    .map((item) => JSON.parse(item));
+};

public/chat_header_button.test.tsx

@@ -180,7 +180,9 @@ describe('<HeaderChatButton />', () => {
     expect(screen.queryByLabelText('chat flyout mock')).not.toBeVisible();
     // sidecar hide
     fireEvent.click(toggleButton);
-    expect(screen.queryByLabelText('chat flyout mock')).toBeVisible();
+    await waitFor(() => {
+      expect(screen.queryByLabelText('chat flyout mock')).toBeVisible();
+    });
   });
 
   it('should call send message without active conversation id after input text submitted', async () => {
@@ -246,11 +248,12 @@ describe('<HeaderChatButton />', () => {
       code: 'Escape',
       charCode: 27,
     });
+
     expect(screen.getByLabelText('chat input')).not.toHaveFocus();
     expect(screen.getByTitle('press Ctrl + / to start typing')).toBeInTheDocument();
   });
 
-  it('should focus on chat input when pressing global shortcut', () => {
+  it('should focus on chat input when pressing global shortcut', async () => {
     render(
       <HeaderChatButton
         application={applicationServiceMock.createStartContract()}
@@ -267,10 +270,12 @@ describe('<HeaderChatButton />', () => {
       charCode: 111,
       ctrlKey: true,
     });
-    expect(screen.getByLabelText('chat input')).toHaveFocus();
+    await waitFor(() => {
+      expect(screen.getByLabelText('chat input')).toHaveFocus();
+    });
   });
 
-  it('should not focus on chat input when no access and pressing global shortcut', () => {
+  it('should not focus on chat input when no access and pressing global shortcut', async () => {
     render(
       <HeaderChatButton
         application={applicationServiceMock.createStartContract()}
@@ -287,7 +292,9 @@ describe('<HeaderChatButton />', () => {
       charCode: 111,
       metaKey: true,
     });
-    expect(screen.getByLabelText('chat input')).not.toHaveFocus();
+    await waitFor(() => {
+      expect(screen.getByLabelText('chat input')).not.toHaveFocus();
+    });
   });
 
   it('should call sidecar hide and close when button unmount and chat flyout is visible', async () => {
@@ -307,11 +314,16 @@ describe('<HeaderChatButton />', () => {
 
     fireEvent.click(getByLabelText('toggle chat flyout icon'));
 
-    expect(sideCarHideMock).not.toHaveBeenCalled();
-    expect(sideCarRefMock.close).not.toHaveBeenCalled();
+    await waitFor(() => {
+      expect(sideCarHideMock).not.toHaveBeenCalled();
+      expect(sideCarRefMock.close).not.toHaveBeenCalled();
+    });
+
     unmount();
-    expect(sideCarHideMock).toHaveBeenCalled();
-    expect(sideCarRefMock.close).toHaveBeenCalled();
+    await waitFor(() => {
+      expect(sideCarHideMock).toHaveBeenCalled();
+      expect(sideCarRefMock.close).toHaveBeenCalled();
+    });
   });
 
   it('should render toggle chat flyout button icon', () => {

public/components/blink_cursor.scss

@@ -0,0 +1,8 @@
+.assistant_blinkCursor {
+  border-right: 2px solid black;
+  animation: cursor-blink 0.8s infinite;
+}
+
+@keyframes cursor-blink {
+  50% { border-right-color: transparent; }
+}

public/components/blink_cursor.tsx

@@ -0,0 +1,11 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import React from 'react';
+import './blink_cursor.scss';
+
+export const BlinkCursor = () => {
+  return <span className="assistant_blinkCursor" />;
+};

public/components/markdown_with_blink_cursor.tsx

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { EuiMarkdownFormat, EuiMarkdownFormatProps } from '@elastic/eui';
+import React from 'react';
+import { useRef } from 'react';
+import { useEffect } from 'react';
+
+function getLastTextNode(node: HTMLElement | Node): Node | null | undefined {
+  if (node.nodeType === node.TEXT_NODE) {
+    return node;
+  }
+  const children = node.childNodes;
+  for (let i = children.length - 1; i >= 0; i--) {
+    const child = children[i];
+    const result = getLastTextNode(child);
+    if (result) {
+      return result;
+    }
+    return null;
+  }
+}
+
+export const MarkdownWithBlinkCursor = (props: EuiMarkdownFormatProps & { loading?: boolean }) => {
+  const ref = useRef<HTMLDivElement | null>(null);
+  useEffect(() => {
+    let blinkCursorNode: undefined | HTMLSpanElement;
+    if (props.loading && ref.current) {
+      const lastNode = getLastTextNode(ref.current);
+      blinkCursorNode = document.createElement('span');
+      blinkCursorNode.classList.add('assistant_blinkCursor');
+      if (lastNode) {
+        lastNode.parentNode?.appendChild(blinkCursorNode);
+      }
+    }
+
+    return () => {
+      blinkCursorNode?.remove();
+    };
+  }, [props.children, props.loading]);
+  return (
+    <div className="markdown_with_blink_cursor" ref={ref}>
+      <EuiMarkdownFormat>{props.children}</EuiMarkdownFormat>
+    </div>
+  );
+};

public/hooks/use_chat_actions.test.tsx

@@ -61,6 +61,19 @@ const SEND_MESSAGE_RESPONSE = {
   ],
 };
 
+const mockPureFetchResponse = (props: { headers?: Headers; responseJson?: unknown } = {}) => {
+  const {
+    headers = new Headers({ 'Content-Type': 'application/json' }),
+    responseJson = {},
+  } = props;
+  return {
+    response: {
+      headers,
+      json: () => Promise.resolve(responseJson),
+    },
+  };
+};
+
 describe('useChatActions hook', () => {
   const httpMock = httpServiceMock.createStartContract();
   const chatStateDispatchMock = jest.fn();
@@ -112,7 +125,11 @@ describe('useChatActions hook', () => {
   });
 
   it('should send message correctly', async () => {
-    httpMock.post.mockResolvedValueOnce(SEND_MESSAGE_RESPONSE);
+    httpMock.post.mockImplementationOnce(async () =>
+      mockPureFetchResponse({
+        responseJson: SEND_MESSAGE_RESPONSE,
+      })
+    );
     jest.spyOn(chatStateHookExports, 'useChatState').mockReturnValue({
       chatState: {
         messages: [SEND_MESSAGE_RESPONSE.messages[0] as IMessage],
@@ -134,6 +151,8 @@ describe('useChatActions hook', () => {
         input: INPUT_MESSAGE,
       }),
       query: dataSourceServiceMock.getDataSourceQuery(),
+      asResponse: true,
+      pureFetch: true,
     });
 
     // it should send dispatch `receive` action to remove the message without messageId
@@ -209,6 +228,8 @@ describe('useChatActions hook', () => {
         input: { type: 'input', content: 'message that send as input', contentType: 'text' },
       }),
       query: dataSourceServiceMock.getDataSourceQuery(),
+      pureFetch: true,
+      asResponse: true,
     });
   });
 
@@ -276,7 +297,11 @@ describe('useChatActions hook', () => {
   });
 
   it('should regenerate message', async () => {
-    httpMock.put.mockResolvedValue(SEND_MESSAGE_RESPONSE);
+    httpMock.put.mockImplementationOnce(async () =>
+      mockPureFetchResponse({
+        responseJson: SEND_MESSAGE_RESPONSE,
+      })
+    );
     jest
       .spyOn(chatContextHookExports, 'useChatContext')
       .mockReturnValue({ ...chatContextMock, conversationId: 'conversation_id_mock' });
@@ -300,14 +325,12 @@ describe('useChatActions hook', () => {
         interactionId: 'interaction_id_mock',
       }),
       query: dataSourceServiceMock.getDataSourceQuery(),
+      pureFetch: true,
+      asResponse: true,
     });
-    expect(chatStateDispatchMock).toHaveBeenCalledWith(
-      expect.objectContaining({ type: 'receive', payload: { messages: [], interactions: [] } })
-    );
-
     expect(chatStateDispatchMock).toHaveBeenCalledWith(
       expect.objectContaining({
-        type: 'patch',
+        type: 'receive',
         payload: {
           messages: SEND_MESSAGE_RESPONSE.messages,
           interactions: SEND_MESSAGE_RESPONSE.interactions,
@@ -317,12 +340,18 @@ 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(),
-    }));
+    const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => {
+      return {
+        signal: AbortSignal.abort(),
+        abort: jest.fn(),
+      };
+    });
 
-    httpMock.put.mockResolvedValue(SEND_MESSAGE_RESPONSE);
+    httpMock.put.mockImplementationOnce(async () =>
+      mockPureFetchResponse({
+        responseJson: SEND_MESSAGE_RESPONSE,
+      })
+    );
     jest
       .spyOn(chatContextHookExports, 'useChatContext')
       .mockReturnValue({ ...chatContextMock, conversationId: 'conversation_id_mock' });
@@ -337,6 +366,8 @@ describe('useChatActions hook', () => {
         interactionId: 'interaction_id_mock',
       }),
       query: dataSourceServiceMock.getDataSourceQuery(),
+      pureFetch: true,
+      asResponse: true,
     });
     expect(chatStateDispatchMock).not.toHaveBeenCalledWith(
       expect.objectContaining({ type: 'receive' })
@@ -359,10 +390,12 @@ 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(),
-    }));
+    const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => {
+      return {
+        signal: AbortSignal.abort(),
+        abort: jest.fn(),
+      };
+    });
     httpMock.put.mockImplementationOnce(() => {
       throw new Error();
     });
@@ -393,7 +426,7 @@ describe('useChatActions hook', () => {
   it('should abort send action after reset chat', async () => {
     const abortFn = jest.fn();
     const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => ({
-      signal: { aborted: true },
+      signal: AbortSignal.abort(),
       abort: abortFn,
     }));
     const { result } = renderHook(() => useChatActions());
@@ -407,7 +440,7 @@ describe('useChatActions hook', () => {
   it('should abort load action after reset chat', async () => {
     const abortFn = jest.fn();
     const AbortControllerMock = jest.spyOn(window, 'AbortController').mockImplementation(() => ({
-      signal: { aborted: true },
+      signal: AbortSignal.abort(),
       abort: abortFn,
     }));
     const { result } = renderHook(() => useChatActions());