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

bundlesize.config.json

@@ -10,11 +10,11 @@
     },
     {
       "path": "./packages/instantsearch.js/dist/instantsearch.production.min.js",
-      "maxSize": "121.25 kB"
+      "maxSize": "122 kB"
     },
     {
       "path": "./packages/instantsearch.js/dist/instantsearch.development.js",
-      "maxSize": "253 kB"
+      "maxSize": "254.25 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,17 @@ export function createChatMessageComponent({ createElement }: Renderer) {
               toolCallId: toolMessage.toolCallId,
             });
 
+          const showBaseLoader =
+            tool.showLoaderDuringStreaming &&
+            toolMessage.state === 'input-streaming' &&
+            LoaderComponent;
+
+          // If the tool is still streaming and has indicated to show the base loader,
+          // we don't render the tool layout component.
+          if (showBaseLoader) {
+            return null;
+          }
+
           if (!ToolLayoutComponent) {
             return null;
           }
@@ -252,8 +269,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

@@ -1,10 +1,12 @@
 /** @jsx createElement */
 
 import { cx } from '../../lib';
+import { startsWith } from '../../lib';
 import {
   getTextContent,
   hasTextContent,
   isPartText,
+  isPartTool,
 } from '../../lib/utils/chat';
 import { createButtonComponent } from '../Button';
 
@@ -30,7 +32,7 @@ import type {
 } from './ChatMessage';
 import type { ChatMessageErrorProps } from './ChatMessageError';
 import type { ChatMessageLoaderProps } from './ChatMessageLoader';
-import type { ChatMessageBase, ChatStatus, ClientSideTools } from './types';
+import type { ChatMessageBase, ChatStatus, ClientSideTool, ClientSideTools } from './types';
 
 export type ChatMessagesTranslations = {
   /**
@@ -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}
@@ -413,13 +418,32 @@ export function createChatMessagesComponent({
     const lastPart = lastMessage?.parts?.[lastMessage.parts.length - 1];
     const isWaitingForResponse = status === 'submitted';
     const isStreamingWithNoContent = status === 'streaming' && !lastPart;
+
+    // Check if the last tool part has showLoaderDuringStreaming enabled and is in input-streaming state
+    // if so, it renders nothing inline and we should show the global loader instead
+    const isToolDeferringToGlobalLoader =
+      lastPart &&
+      isPartTool(lastPart) &&
+      lastPart.state === 'input-streaming' &&
+      (() => {
+        const toolName = lastPart.type.replace('tool-', '');
+        let tool: ClientSideTool | undefined = tools[toolName];
+        if (!tool) {
+          tool = Object.entries(tools).find(([key]) =>
+            startsWith(toolName, `${key}_`)
+          )?.[1];
+        }
+        return tool?.showLoaderDuringStreaming;
+      })();
+
     const isStreamingNonTextContent =
       status === 'streaming' && lastPart && !isPartText(lastPart);
 
     const showLoader =
       isWaitingForResponse ||
       isStreamingWithNoContent ||
-      isStreamingNonTextContent;
+      isStreamingNonTextContent ||
+      isToolDeferringToGlobalLoader;
 
     const DefaultMessage = MessageComponent || DefaultMessageComponent;
     const DefaultLoader = LoaderComponent || DefaultLoaderComponent;
@@ -463,6 +487,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';
@@ -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;
+  showLoaderDuringStreaming?: boolean;
   addToolResult: AddToolResult;
   sendEvent?: SendEventForHits;
   onToolCall?: (

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,9 @@ 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-');
+};

packages/instantsearch.js/src/connectors/chat/__tests__/connectChat-test.ts

@@ -506,6 +506,65 @@ data: [DONE]`,
         );
       });
     });
+
+    it('streams tool input parts from tool-input-delta without tool-input-available', async () => {
+      const { widget } = getInitializedWidget({
+        agentId: undefined,
+        transport: {
+          fetch: () =>
+            Promise.resolve(
+              new Response(
+                `data: {"type": "start", "messageId": "test-id"}
+
+data: {"type": "start-step"}
+
+data: {"type": "tool-input-start", "toolCallId": "call_1", "toolName": "displayResults"}
+
+data: {"type": "tool-input-delta", "toolCallId": "call_1", "toolName": "displayResults", "inputTextDelta": "{}"}
+
+data: {"type": "finish-step"}
+
+data: {"type": "finish"}
+
+data: [DONE]`,
+                {
+                  headers: { 'Content-Type': 'text/event-stream' },
+                }
+              )
+            ),
+        },
+      });
+
+      const { chatInstance } = widget;
+
+      await chatInstance.sendMessage({
+        id: 'message-id',
+        role: 'user',
+        parts: [{ type: 'text', text: 'Show me product groups' }],
+      });
+
+      await waitFor(() => {
+        const lastMessage = chatInstance.messages[chatInstance.messages.length - 1];
+        expect(lastMessage?.role).toBe('assistant');
+
+        const toolPart = lastMessage?.parts.find(
+          (part) =>
+            'type' in part &&
+            part.type === 'tool-displayResults' &&
+            'toolCallId' in part &&
+            part.toolCallId === 'call_1'
+        ) as
+          | {
+              state: string;
+              rawInput?: string;
+              input?: Record<string, unknown>;
+            }
+          | undefined;
+
+        expect(toolPart?.state).toBe('input-streaming');
+        expect(toolPart?.input).toEqual({});
+      });
+    });
   });
 
   describe('transport configuration', () => {