feat(chat): add tool input streaming support and chat loader component

bundlesize.config.json

@@ -10,11 +10,11 @@
     },
     {
       "path": "./packages/instantsearch.js/dist/instantsearch.production.min.js",
-      "maxSize": "121.8 kB"
+      "maxSize": "122.75 kB"
     },
     {
       "path": "./packages/instantsearch.js/dist/instantsearch.development.js",
-      "maxSize": "253.4 kB"
+      "maxSize": "255 kB"
     },
     {
       "path": "packages/react-instantsearch-core/dist/umd/ReactInstantSearchCore.min.js",

packages/instantsearch-ui-components/src/components/chat/ChatMessage.tsx

@@ -7,6 +7,7 @@ import { createButtonComponent } from '../Button';
 import { MenuIcon } from './icons';
 
 import type { ComponentProps, Renderer, VNode } from '../../types';
+import type { ChatMessageLoaderProps } from './ChatMessageLoader';
 import type {
   AddToolResultWithOutput,
   ChatMessageBase,
@@ -135,6 +136,10 @@ export type ChatMessageProps = ComponentProps<'article'> & {
    * Array of tools available for the assistant (for tool messages)
    */
   tools: ClientSideTools;
+  /**
+   * Loader component passed to tool layout components
+   */
+  loaderComponent: (props: ChatMessageLoaderProps) => JSX.Element;
   /**
    * Optional suggestions element
    */
@@ -171,6 +176,7 @@ export function createChatMessageComponent({ createElement }: Renderer) {
       indexUiState,
       setIndexUiState,
       onClose,
+      loaderComponent: LoaderComponent,
       translations: userTranslations,
       suggestionsElement,
       ...props
@@ -237,6 +243,13 @@ export function createChatMessageComponent({ createElement }: Renderer) {
               toolCallId: toolMessage.toolCallId,
             });
 
+          if (
+            toolMessage.state === 'input-streaming' &&
+            !tool.streamInput
+          ) {
+            return null;
+          }
+
           if (!ToolLayoutComponent) {
             return null;
           }
@@ -252,8 +265,9 @@ export function createChatMessageComponent({ createElement }: Renderer) {
                 setIndexUiState={setIndexUiState}
                 addToolResult={boundAddToolResult}
                 applyFilters={tool.applyFilters}
-                sendEvent={tool.sendEvent || (() => {})}
+                sendEvent={tool.sendEvent || (() => { })}
                 onClose={onClose}
+                loaderComponent={LoaderComponent}
               />
             </div>
           );

packages/instantsearch-ui-components/src/components/chat/ChatMessageLoader.tsx

@@ -24,7 +24,7 @@ export function createChatMessageLoaderComponent({
   return function ChatMessageLoader(userProps: ChatMessageLoaderProps) {
     const { translations: userTranslations, ...props } = userProps;
     const translations: Required<ChatMessageLoaderTranslations> = {
-      loaderText: 'Thinking...',
+      loaderText: '',
       ...userTranslations,
     };
 

packages/instantsearch-ui-components/src/components/chat/ChatMessages.tsx

@@ -2,9 +2,11 @@
 
 import { cx } from '../../lib';
 import {
+  findTool,
   getTextContent,
   hasTextContent,
   isPartText,
+  isPartTool,
 } from '../../lib/utils/chat';
 import { createButtonComponent } from '../Button';
 
@@ -225,6 +227,7 @@ function createDefaultMessageComponent<
     onFeedback,
     feedbackState,
     actionsComponent,
+    loaderComponent,
     classNames,
     messageTranslations,
     translations,
@@ -243,6 +246,7 @@ function createDefaultMessageComponent<
     onFeedback?: (messageId: string, vote: 0 | 1) => void;
     feedbackState?: Record<string, 'sending' | 0 | 1>;
     actionsComponent?: ChatMessageProps['actionsComponent'];
+    loaderComponent: ChatMessageProps['loaderComponent'];
     translations: ChatMessagesTranslations;
     classNames?: Partial<ChatMessageClassNames>;
     messageTranslations?: Partial<ChatMessageTranslations>;
@@ -326,6 +330,7 @@ function createDefaultMessageComponent<
         onClose={onClose}
         actions={defaultActions}
         actionsComponent={actionsComponent}
+        loaderComponent={loaderComponent}
         data-role={message.role}
         classNames={classNames}
         translations={messageTranslations}
@@ -411,15 +416,7 @@ export function createChatMessagesComponent({
 
     const lastMessage = messages[messages.length - 1];
     const lastPart = lastMessage?.parts?.[lastMessage.parts.length - 1];
-    const isWaitingForResponse = status === 'submitted';
-    const isStreamingWithNoContent = status === 'streaming' && !lastPart;
-    const isStreamingNonTextContent =
-      status === 'streaming' && lastPart && !isPartText(lastPart);
-
-    const showLoader =
-      isWaitingForResponse ||
-      isStreamingWithNoContent ||
-      isStreamingNonTextContent;
+    const showLoader = getShowLoader(status, lastPart, tools);
 
     const DefaultMessage = MessageComponent || DefaultMessageComponent;
     const DefaultLoader = LoaderComponent || DefaultLoaderComponent;
@@ -463,6 +460,7 @@ export function createChatMessagesComponent({
                 onFeedback={onFeedback}
                 feedbackState={feedbackState}
                 actionsComponent={ActionsComponent}
+                loaderComponent={DefaultLoader}
                 onClose={onClose}
                 translations={translations}
                 classNames={messageClassNames}
@@ -506,3 +504,27 @@ export function createChatMessagesComponent({
     );
   };
 }
+
+const getShowLoader = (
+  status: ChatStatus,
+  lastPart: ChatMessageBase['parts'][number] | undefined,
+  tools: ClientSideTools
+): boolean => {
+  if (status !== 'submitted' && status !== 'streaming') return false;
+  if (status === 'submitted') return true;
+
+  if (!lastPart) return true;
+  if (isPartText(lastPart)) return false;
+
+  if (isPartTool(lastPart)) {
+    if (lastPart.state === 'output-available') return false;
+    if (lastPart.state === 'input-streaming') {
+      const tool = findTool(lastPart.type, tools);
+      return !tool?.streamInput;
+    }
+    return true;
+  }
+
+  return true;
+};
+

packages/instantsearch-ui-components/src/components/chat/__tests__/ChatMessage.test.tsx

@@ -21,6 +21,7 @@ describe('ChatMessage', () => {
         message={{ role: 'user', id: '1', parts: [] }}
         status="ready"
         tools={{}}
+        loaderComponent={jest.fn()}
         onClose={jest.fn()}
       />
     );
@@ -66,6 +67,7 @@ describe('ChatMessage', () => {
           actions: 'actions',
         }}
         tools={{}}
+        loaderComponent={jest.fn()}
         onClose={jest.fn()}
       />
     );
@@ -104,6 +106,7 @@ describe('ChatMessage', () => {
           }}
           status="ready"
           tools={{}}
+          loaderComponent={jest.fn()}
           onClose={jest.fn()}
         />
         <ChatMessage
@@ -116,6 +119,7 @@ describe('ChatMessage', () => {
           }}
           status="ready"
           tools={{}}
+          loaderComponent={jest.fn()}
           onClose={jest.fn()}
         />
         <ChatMessage
@@ -128,6 +132,7 @@ describe('ChatMessage', () => {
           }}
           status="ready"
           tools={{}}
+          loaderComponent={jest.fn()}
           onClose={jest.fn()}
         />
       </div>
@@ -235,6 +240,7 @@ describe('ChatMessage', () => {
             applyFilters: jest.fn(),
           },
         }}
+        loaderComponent={jest.fn()}
         onClose={jest.fn()}
       />
     );

packages/instantsearch-ui-components/src/components/chat/types.ts

@@ -1,4 +1,5 @@
 import type { ComponentProps, SendEventForHits } from '../../types';
+import type { ChatMessageLoaderProps } from './ChatMessageLoader';
 import type { SearchParameters } from 'algoliasearch-helper';
 
 export type ChatStatus = 'ready' | 'submitted' | 'streaming' | 'error';
@@ -485,6 +486,7 @@ export type ClientSideToolComponentProps = {
   onClose: () => void;
   addToolResult: AddToolResultWithOutput;
   applyFilters: (params: ApplyFiltersParams) => SearchParameters;
+  loaderComponent: (props: ChatMessageLoaderProps) => JSX.Element;
   sendEvent: SendEventForHits;
 };
 
@@ -494,6 +496,7 @@ export type ClientSideToolComponent = (
 
 export type ClientSideTool = {
   layoutComponent?: ClientSideToolComponent;
+  streamInput?: boolean;
   addToolResult: AddToolResult;
   sendEvent?: SendEventForHits;
   onToolCall?: (

packages/instantsearch-ui-components/src/lib/utils/chat.ts

@@ -1,4 +1,11 @@
+import { startsWith } from './startsWith';
+
 import type { ChatMessageBase } from '../../components';
+import type {
+  ChatToolMessage,
+  ClientSideTool,
+  ClientSideTools,
+} from '../../components/chat/types';
 
 export const getTextContent = (message: ChatMessageBase) => {
   return message.parts
@@ -15,3 +22,23 @@ export const isPartText = (
 ): part is Extract<ChatMessageBase['parts'][number], { type: 'text' }> => {
   return part.type === 'text';
 };
+
+export const isPartTool = (
+  part: ChatMessageBase['parts'][number]
+): part is ChatToolMessage => {
+  return startsWith(part.type, 'tool-');
+};
+
+export const findTool = (
+  partType: string,
+  tools: ClientSideTools
+): ClientSideTool | undefined => {
+  const toolName = partType.replace('tool-', '');
+  let tool: ClientSideTool | undefined = tools[toolName];
+  if (!tool) {
+    tool = Object.entries(tools).find(([key]) =>
+      startsWith(toolName, `${key}_`)
+    )?.[1];
+  }
+  return tool;
+};