fix (ai/core): only append assistant response at the end when there i…

…s a final user message (#4623)
vercel · Jan 31, 2025 · ca89615 · ca89615
1 parent d99c7fb
commit ca89615
Show file tree

Hide file tree

Showing 3 changed files with 171 additions and 14 deletions.
diff --git a/.changeset/six-jokes-join.md b/.changeset/six-jokes-join.md
@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+fix (ai/core): only append assistant response at the end when there is a final user message
diff --git a/packages/ai/core/prompt/convert-to-core-messages.test.ts b/packages/ai/core/prompt/convert-to-core-messages.test.ts
@@ -230,6 +230,10 @@ describe('assistant message', () => {
       {
         role: 'assistant',
         content: [
+          {
+            type: 'text',
+            text: 'Let me calculate that for you.',
+          },
           {
             type: 'tool-call',
             toolCallId: 'call1',
@@ -249,10 +253,6 @@ describe('assistant message', () => {
           },
         ],
       },
-      {
-        role: 'assistant',
-        content: 'Let me calculate that for you.',
-      },
     ]);
   });
 
@@ -290,6 +290,10 @@ describe('assistant message', () => {
       {
         role: 'assistant',
         content: [
+          {
+            type: 'text',
+            text: 'Let me calculate that for you.',
+          },
           {
             type: 'tool-call',
             toolCallId: 'call1',
@@ -310,10 +314,6 @@ describe('assistant message', () => {
           },
         ],
       },
-      {
-        role: 'assistant',
-        content: 'Let me calculate that for you.',
-      },
     ]);
   });
 
@@ -400,6 +400,10 @@ describe('assistant message', () => {
       {
         role: 'assistant',
         content: [
+          {
+            type: 'text',
+            text: 'response',
+          },
           {
             type: 'tool-call',
             toolCallId: 'call-1',
@@ -475,10 +479,6 @@ describe('assistant message', () => {
           },
         ],
       },
-      {
-        role: 'assistant',
-        content: 'response',
-      },
     ]);
   });
 });
@@ -525,6 +525,153 @@ describe('multiple messages', () => {
       },
     ]);
   });
+
+  it('should handle conversation with multiple tool invocations and user message at the end', () => {
+    const tools = {
+      screenshot: tool({
+        parameters: z.object({ value: z.string() }),
+        execute: async () => 'imgbase64',
+      }),
+    };
+
+    const result = convertToCoreMessages(
+      [
+        {
+          role: 'assistant',
+          content: 'response',
+          toolInvocations: [
+            {
+              state: 'result',
+              toolCallId: 'call-1',
+              toolName: 'screenshot',
+              args: { value: 'value-1' },
+              result: 'result-1',
+              step: 0,
+            },
+            {
+              state: 'result',
+              toolCallId: 'call-2',
+              toolName: 'screenshot',
+              args: { value: 'value-2' },
+              result: 'result-2',
+              step: 1,
+            },
+
+            {
+              state: 'result',
+              toolCallId: 'call-3',
+              toolName: 'screenshot',
+              args: { value: 'value-3' },
+              result: 'result-3',
+              step: 1,
+            },
+            {
+              state: 'result',
+              toolCallId: 'call-4',
+              toolName: 'screenshot',
+              args: { value: 'value-4' },
+              result: 'result-4',
+              step: 2,
+            },
+          ],
+        },
+        {
+          role: 'user',
+          content: 'Thanks!',
+        },
+      ],
+      { tools }, // separate tools to ensure that types are inferred correctly
+    );
+
+    expect(result).toEqual([
+      {
+        role: 'assistant',
+        content: [
+          {
+            type: 'tool-call',
+            toolCallId: 'call-1',
+            toolName: 'screenshot',
+            args: { value: 'value-1' },
+          },
+        ],
+      },
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'screenshot',
+            result: 'result-1',
+          },
+        ],
+      },
+      {
+        role: 'assistant',
+        content: [
+          {
+            type: 'tool-call',
+            toolCallId: 'call-2',
+            toolName: 'screenshot',
+            args: { value: 'value-2' },
+          },
+          {
+            type: 'tool-call',
+            toolCallId: 'call-3',
+            toolName: 'screenshot',
+            args: { value: 'value-3' },
+          },
+        ],
+      },
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-2',
+            toolName: 'screenshot',
+            result: 'result-2',
+          },
+          {
+            type: 'tool-result',
+            toolCallId: 'call-3',
+            toolName: 'screenshot',
+            result: 'result-3',
+          },
+        ],
+      },
+      {
+        role: 'assistant',
+        content: [
+          {
+            type: 'tool-call',
+            toolCallId: 'call-4',
+            toolName: 'screenshot',
+            args: { value: 'value-4' },
+          },
+        ],
+      },
+      {
+        role: 'tool',
+        content: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-4',
+            toolName: 'screenshot',
+            result: 'result-4',
+          },
+        ],
+      },
+      {
+        role: 'assistant',
+        content: 'response',
+      },
+      {
+        role: 'user',
+        content: 'Thanks!',
+      },
+    ]);
+  });
 });
 
 describe('error handling', () => {

diff --git a/packages/ai/core/prompt/convert-to-core-messages.ts b/packages/ai/core/prompt/convert-to-core-messages.ts
@@ -15,7 +15,9 @@ export function convertToCoreMessages<TOOLS extends ToolSet = never>(
   const tools = options?.tools ?? ({} as TOOLS);
   const coreMessages: CoreMessage[] = [];
 
-  for (const message of messages) {
+  for (let i = 0; i < messages.length; i++) {
+    const message = messages[i];
+    const isLastMessage = i === messages.length - 1;
     const { role, content, toolInvocations, experimental_attachments } =
       message;
 
@@ -64,6 +66,9 @@ export function convertToCoreMessages<TOOLS extends ToolSet = never>(
           coreMessages.push({
             role: 'assistant',
             content: [
+              ...(isLastMessage && content && i === 0
+                ? [{ type: 'text' as const, text: content }]
+                : []),
               ...stepInvocations.map(
                 ({ toolCallId, toolName, args }): ToolCallPart => ({
                   type: 'tool-call' as const,
@@ -110,7 +115,7 @@ export function convertToCoreMessages<TOOLS extends ToolSet = never>(
           });
         }
 
-        if (content) {
+        if (content && !isLastMessage) {
           coreMessages.push({ role: 'assistant', content });
         }