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

bundlesize.config.json

@@ -14,7 +14,7 @@
     },
     {
       "path": "./packages/instantsearch.js/dist/instantsearch.development.js",
-      "maxSize": "251.25 kB"
+      "maxSize": "252.75 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
@@ -252,8 +258,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/ChatMessages.tsx

@@ -5,6 +5,7 @@ import {
   getTextContent,
   hasTextContent,
   isPartText,
+  isToolPartActivelyRendering,
 } from '../../lib/utils/chat';
 import { createButtonComponent } from '../Button';
 
@@ -225,6 +226,7 @@ function createDefaultMessageComponent<
     onFeedback,
     feedbackState,
     actionsComponent,
+    loaderComponent,
     classNames,
     messageTranslations,
     translations,
@@ -243,6 +245,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 +329,7 @@ function createDefaultMessageComponent<
         onClose={onClose}
         actions={defaultActions}
         actionsComponent={actionsComponent}
+        loaderComponent={loaderComponent}
         data-role={message.role}
         classNames={classNames}
         translations={messageTranslations}
@@ -414,7 +418,7 @@ export function createChatMessagesComponent({
     const isWaitingForResponse = status === 'submitted';
     const isStreamingWithNoContent = status === 'streaming' && !lastPart;
     const isStreamingNonTextContent =
-      status === 'streaming' && lastPart && !isPartText(lastPart);
+      status === 'streaming' && lastPart && !(isPartText(lastPart) || isToolPartActivelyRendering(lastPart));
 
     const showLoader =
       isWaitingForResponse ||
@@ -463,6 +467,7 @@ export function createChatMessagesComponent({
                 onFeedback={onFeedback}
                 feedbackState={feedbackState}
                 actionsComponent={ActionsComponent}
+                loaderComponent={DefaultLoader}
                 onClose={onClose}
                 translations={translations}
                 classNames={messageClassNames}

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';
@@ -120,38 +121,39 @@ export type ToolUIPart<TTools extends UITools = UITools> = ValueOf<{
     toolCallId: string;
   } & (
     | {
-        state: 'input-streaming';
-        input: DeepPartial<TTools[NAME]['input']> | undefined;
-        providerExecuted?: boolean;
-        output?: never;
-        errorText?: never;
-      }
+      state: 'input-streaming';
+      input: DeepPartial<TTools[NAME]['input']> | undefined;
+      rawInput?: string;
+      providerExecuted?: boolean;
+      output?: never;
+      errorText?: never;
+    }
     | {
-        state: 'input-available';
-        input: TTools[NAME]['input'];
-        providerExecuted?: boolean;
-        output?: never;
-        errorText?: never;
-        callProviderMetadata?: ProviderMetadata;
-      }
+      state: 'input-available';
+      input: TTools[NAME]['input'];
+      providerExecuted?: boolean;
+      output?: never;
+      errorText?: never;
+      callProviderMetadata?: ProviderMetadata;
+    }
     | {
-        state: 'output-available';
-        input: TTools[NAME]['input'];
-        output: TTools[NAME]['output'];
-        errorText?: never;
-        providerExecuted?: boolean;
-        callProviderMetadata?: ProviderMetadata;
-        preliminary?: boolean;
-      }
+      state: 'output-available';
+      input: TTools[NAME]['input'];
+      output: TTools[NAME]['output'];
+      errorText?: never;
+      providerExecuted?: boolean;
+      callProviderMetadata?: ProviderMetadata;
+      preliminary?: boolean;
+    }
     | {
-        state: 'output-error';
-        input: TTools[NAME]['input'] | undefined;
-        rawInput?: unknown;
-        output?: never;
-        errorText: string;
-        providerExecuted?: boolean;
-        callProviderMetadata?: ProviderMetadata;
-      }
+      state: 'output-error';
+      input: TTools[NAME]['input'] | undefined;
+      rawInput?: unknown;
+      output?: never;
+      errorText: string;
+      providerExecuted?: boolean;
+      callProviderMetadata?: ProviderMetadata;
+    }
   );
 }>;
 
@@ -163,35 +165,36 @@ export type DynamicToolUIPart = {
   toolName: string;
   toolCallId: string;
 } & (
-  | {
+    | {
       state: 'input-streaming';
       input: unknown | undefined;
+      rawInput?: string;
       output?: never;
       errorText?: never;
     }
-  | {
+    | {
       state: 'input-available';
       input: unknown;
       output?: never;
       errorText?: never;
       callProviderMetadata?: ProviderMetadata;
     }
-  | {
+    | {
       state: 'output-available';
       input: unknown;
       output: unknown;
       errorText?: never;
       callProviderMetadata?: ProviderMetadata;
       preliminary?: boolean;
     }
-  | {
+    | {
       state: 'output-error';
       input: unknown;
       output?: never;
       errorText: string;
       callProviderMetadata?: ProviderMetadata;
     }
-);
+  );
 
 /**
  * All possible message part types.
@@ -294,23 +297,23 @@ export type ChatOnErrorCallback = (error: Error) => void;
  */
 export type InferUIMessageToolCall<TUIMessage extends UIMessage> =
   | ValueOf<{
-      [NAME in keyof InferUIMessageTools<TUIMessage>]: {
-        toolName: NAME & string;
-        toolCallId: string;
-        input: InferUIMessageTools<TUIMessage>[NAME] extends {
-          input: infer INPUT;
-        }
-          ? INPUT
-          : never;
-        dynamic?: false;
-      };
-    }>
-  | {
-      toolName: string;
+    [NAME in keyof InferUIMessageTools<TUIMessage>]: {
+      toolName: NAME & string;
       toolCallId: string;
-      input: unknown;
-      dynamic: true;
+      input: InferUIMessageTools<TUIMessage>[NAME] extends {
+        input: infer INPUT;
+      }
+      ? INPUT
+      : never;
+      dynamic?: false;
     };
+  }>
+  | {
+    toolName: string;
+    toolCallId: string;
+    input: unknown;
+    dynamic: true;
+  };
 
 /**
  * Optional callback function that is invoked when a tool call is received.
@@ -398,25 +401,25 @@ export interface AbstractChat<TUIMessage extends UIMessage> {
   sendMessage: (
     message?:
       | (Omit<TUIMessage, 'id' | 'role'> & {
-          id?: TUIMessage['id'];
-          role?: TUIMessage['role'];
-          text?: never;
-          files?: never;
-          messageId?: string;
-        })
+        id?: TUIMessage['id'];
+        role?: TUIMessage['role'];
+        text?: never;
+        files?: never;
+        messageId?: string;
+      })
       | {
-          text: string;
-          files?: FileList | FileUIPart[];
-          metadata?: InferUIMessageMetadata<TUIMessage>;
-          parts?: never;
-          messageId?: string;
-        }
+        text: string;
+        files?: FileList | FileUIPart[];
+        metadata?: InferUIMessageMetadata<TUIMessage>;
+        parts?: never;
+        messageId?: string;
+      }
       | {
-          files: FileList | FileUIPart[];
-          metadata?: InferUIMessageMetadata<TUIMessage>;
-          parts?: never;
-          messageId?: string;
-        },
+        files: FileList | FileUIPart[];
+        metadata?: InferUIMessageMetadata<TUIMessage>;
+        parts?: never;
+        messageId?: string;
+      },
     options?: { headers?: Record<string, string> | Headers; body?: object }
   ) => Promise<void>;
 
@@ -485,6 +488,7 @@ export type ClientSideToolComponentProps = {
   onClose: () => void;
   addToolResult: AddToolResultWithOutput;
   applyFilters: (params: ApplyFiltersParams) => SearchParameters;
+  loaderComponent: (props: ChatMessageLoaderProps) => JSX.Element;
   sendEvent: SendEventForHits;
 };
 

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

@@ -1,4 +1,7 @@
+import { startsWith } from './startsWith';
+
 import type { ChatMessageBase } from '../../components';
+import type { ChatToolMessage } from '../../components/chat/types';
 
 export const getTextContent = (message: ChatMessageBase) => {
   return message.parts
@@ -15,3 +18,20 @@ 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 isToolPartActivelyRendering = (
+  part: ChatMessageBase['parts'][number]
+): boolean => {
+  return (
+    isPartTool(part) &&
+    (part.state === 'input-streaming' ||
+      part.state === 'output-available' ||
+      part.state === 'output-error')
+  );
+};