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": "122 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

@@ -237,6 +237,13 @@ export function createChatMessageComponent({ createElement }: Renderer) {
               toolCallId: toolMessage.toolCallId,
             });
 
+          if (
+            toolMessage.state === 'input-streaming' &&
+            !tool.streamInput
+          ) {
+            return null;
+          }
+
           if (!ToolLayoutComponent) {
             return null;
           }

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';
 
@@ -411,15 +413,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;
@@ -506,3 +500,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/types.ts

@@ -494,6 +494,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;
+};

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

@@ -506,6 +506,194 @@ 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({});
+      });
+    });
+
+    it('skips JSON repair for tools without streamInput (default)', async () => {
+      const { widget } = getInitializedWidget({
+        agentId: undefined,
+        tools: {
+          myTool: {},
+        },
+        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": "myTool"}
+
+data: {"type": "tool-input-delta", "toolCallId": "call_1", "toolName": "myTool", "inputTextDelta": "{\\"query\\": \\"sho"}
+
+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: 'search' }],
+      });
+
+      await waitFor(() => {
+        const lastMessage =
+          chatInstance.messages[chatInstance.messages.length - 1];
+        const toolPart = lastMessage?.parts.find(
+          (part) =>
+            'type' in part &&
+            part.type === 'tool-myTool' &&
+            'toolCallId' in part &&
+            part.toolCallId === 'call_1'
+        ) as
+          | {
+              state: string;
+              rawInput?: string;
+              input?: unknown;
+            }
+          | undefined;
+
+        expect(toolPart?.state).toBe('input-streaming');
+        // Input is not repaired since streamInput is not set (default)
+        expect(toolPart?.input).toBeUndefined();
+        // Raw input is still accumulated
+        expect(toolPart?.rawInput).toBe('{"query": "sho');
+      });
+    });
+
+    it('repairs JSON for tools with streamInput set to true', async () => {
+      const { widget } = getInitializedWidget({
+        agentId: undefined,
+        tools: {
+          myTool: {
+            streamInput: true,
+          },
+        },
+        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": "myTool"}
+
+data: {"type": "tool-input-delta", "toolCallId": "call_1", "toolName": "myTool", "inputTextDelta": "{\\"query\\": \\"sho"}
+
+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: 'search' }],
+      });
+
+      await waitFor(() => {
+        const lastMessage =
+          chatInstance.messages[chatInstance.messages.length - 1];
+        const toolPart = lastMessage?.parts.find(
+          (part) =>
+            'type' in part &&
+            part.type === 'tool-myTool' &&
+            'toolCallId' in part &&
+            part.toolCallId === 'call_1'
+        ) as
+          | {
+              state: string;
+              rawInput?: string;
+              input?: unknown;
+            }
+          | undefined;
+
+        expect(toolPart?.state).toBe('input-streaming');
+        // Input is repaired since streamInput is true
+        expect(toolPart?.input).toEqual({ query: 'sho' });
+        expect(toolPart?.rawInput).toBe('{"query": "sho');
+      });
+    });
   });
 
   describe('transport configuration', () => {

packages/instantsearch.js/src/connectors/chat/connectChat.ts

@@ -420,6 +420,14 @@ export default (function connectChat<TWidgetParams extends UnknownWidgetParams>(
         ...options,
         transport,
         sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithToolCalls,
+        shouldRepairToolInput(toolName) {
+          let tool = tools[toolName];
+          if (!tool && toolName.startsWith(`${SearchIndexToolType}_`)) {
+            tool = tools[SearchIndexToolType];
+          }
+          if (!tool) return true;
+          return Boolean(tool.streamInput);
+        },
         onToolCall({ toolCall }) {
           let tool = tools[toolCall.toolName];