feat (ai/react): add experimental_maxAutomaticRoundtrips to useChat (#1681)

Author: lgrammel

.changeset/great-bottles-burn.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai/react): add experimental_maxAutomaticRoundtrips to useChat

examples/next-openai/app/api/use-chat-tool-result-roundtrip/route.ts

@@ -0,0 +1,37 @@
+import { openai } from '@ai-sdk/openai';
+import { convertToCoreMessages, streamText } from 'ai';
+import { z } from 'zod';
+
+export const dynamic = 'force-dynamic';
+export const maxDuration = 60;
+
+export async function POST(req: Request) {
+  const { messages } = await req.json();
+
+  const result = await streamText({
+    model: openai('gpt-4-turbo'),
+    system:
+      'You are a weather bot that can use the weather tool to get the weather in a given city. ' +
+      'Respond to the user with weather information in a friendly and helpful manner.',
+    messages: convertToCoreMessages(messages),
+    tools: {
+      weather: {
+        description: 'show the weather in a given city to the user',
+        parameters: z.object({ city: z.string() }),
+        execute: async ({}: { city: string }) => {
+          // Random delay between 1000ms (1s) and 3000ms (3s):
+          const delay = Math.floor(Math.random() * (3000 - 1000 + 1)) + 1000;
+          await new Promise(resolve => setTimeout(resolve, delay));
+
+          // Random weather:
+          const weatherOptions = ['sunny', 'cloudy', 'rainy', 'snowy', 'windy'];
+          return weatherOptions[
+            Math.floor(Math.random() * weatherOptions.length)
+          ];
+        },
+      },
+    },
+  });
+
+  return result.toAIStreamResponse();
+}

examples/next-openai/app/use-chat-tool-result-roundtrip/page.tsx

@@ -0,0 +1,34 @@
+'use client';
+
+import { useChat } from 'ai/react';
+
+export default function Chat() {
+  const { messages, input, handleInputChange, handleSubmit } = useChat({
+    api: '/api/use-chat-tool-result-roundtrip',
+    experimental_maxAutomaticRoundtrips: 2,
+  });
+
+  return (
+    <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
+      {messages
+        .filter(m => m.content) // filter out empty messages
+        .map(m => (
+          <div key={m.id} className="whitespace-pre-wrap">
+            <strong>{`${m.role}: `}</strong>
+            {m.content}
+            <br />
+            <br />
+          </div>
+        ))}
+
+      <form onSubmit={handleSubmit}>
+        <input
+          className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
+          value={input}
+          placeholder="Say something..."
+          onChange={handleInputChange}
+        />
+      </form>
+    </div>
+  );
+}

packages/core/react/use-chat.ts

@@ -219,6 +219,7 @@ export function useChat({
   sendExtraMessageFields,
   experimental_onFunctionCall,
   experimental_onToolCall,
+  experimental_maxAutomaticRoundtrips = 0,
   streamMode,
   onResponse,
   onFinish,
@@ -230,6 +231,19 @@ export function useChat({
 }: Omit<UseChatOptions, 'api'> & {
   api?: string | StreamingReactResponseAction;
   key?: string;
+  /**
+Maximal number of automatic roundtrips for tool calls.
+
+An automatic tool call roundtrip is a call to the server with the 
+tool call results when all tool calls in the last assistant 
+message have results.
+
+A maximum number is required to prevent infinite loops in the
+case of misconfigured tools.
+
+By default, it's set to 0, which will disable the feature.
+   */
+  experimental_maxAutomaticRoundtrips?: number;
 } = {}): UseChatHelpers & {
   experimental_addToolResult: ({
     toolCallId,
@@ -343,6 +357,23 @@ export function useChat({
       } finally {
         mutateLoading(false);
       }
+
+      // auto-submit when all tool calls in the last assistant message have results:
+      const messages = messagesRef.current;
+      const lastMessage = messages[messages.length - 1];
+      if (
+        // ensure there is a last message:
+        lastMessage != null &&
+        // check if the feature is enabled:
+        experimental_maxAutomaticRoundtrips > 0 &&
+        // check that roundtrip is possible:
+        isAssistantMessageWithCompletedToolCalls(lastMessage) &&
+        // limit the number of automatic roundtrips:
+        countTrailingAssistantMessages(messages) <=
+          experimental_maxAutomaticRoundtrips
+      ) {
+        await triggerRequest({ messages });
+      }
     },
     [
       mutate,
@@ -359,6 +390,7 @@ export function useChat({
       sendExtraMessageFields,
       experimental_onFunctionCall,
       experimental_onToolCall,
+      experimental_maxAutomaticRoundtrips,
       messagesRef,
       abortControllerRef,
       generateId,
@@ -526,16 +558,38 @@ export function useChat({
 
       // auto-submit when all tool calls in the last assistant message have results:
       const lastMessage = updatedMessages[updatedMessages.length - 1];
-      if (
-        lastMessage.role === 'assistant' &&
-        lastMessage.toolInvocations &&
-        lastMessage.toolInvocations.length > 0 &&
-        lastMessage.toolInvocations.every(
-          toolInvocation => 'result' in toolInvocation,
-        )
-      ) {
+      if (isAssistantMessageWithCompletedToolCalls(lastMessage)) {
         triggerRequest({ messages: updatedMessages });
       }
     },
   };
 }
+
+/**
+Check if the message is an assistant message with completed tool calls. 
+The message must have at least one tool invocation and all tool invocations
+must have a result.
+ */
+function isAssistantMessageWithCompletedToolCalls(message: Message) {
+  return (
+    message.role === 'assistant' &&
+    message.toolInvocations &&
+    message.toolInvocations.length > 0 &&
+    message.toolInvocations.every(toolInvocation => 'result' in toolInvocation)
+  );
+}
+
+/**
+Returns the number of trailing assistant messages in the array.
+ */
+function countTrailingAssistantMessages(messages: Message[]) {
+  let count = 0;
+  for (let i = messages.length - 1; i >= 0; i--) {
+    if (messages[i].role === 'assistant') {
+      count++;
+    } else {
+      break;
+    }
+  }
+  return count;
+}