feat (ai): support changing the system prompt in prepareSteps (#6503)

## Background

Changing the system prompt for specific steps is necessary for
swarm-like multi agent behavior.

## Summary

Add system option to prepareStep result.

Author: lgrammel

.changeset/proud-dancers-doubt.md

@@ -0,0 +1,5 @@
+---
+'ai': patch
+---
+
+feat (ai): support changing the system prompt in prepareSteps

examples/ai-core/src/stream-text/openai-swarm.ts

@@ -0,0 +1,41 @@
+import { openai } from '@ai-sdk/openai';
+import { stepCountIs, streamText, tool } from 'ai';
+import 'dotenv/config';
+import { z } from 'zod';
+
+main().catch(console.error);
+
+async function main() {
+  const agentA = {
+    system: 'You are a helpful agent.',
+    activeTools: ['transferToAgentB'] as 'transferToAgentB'[],
+  };
+
+  const agentB = {
+    system: 'Only speak in Haikus.',
+    activeTools: [],
+  };
+
+  let activeAgent = agentA;
+
+  const result = streamText({
+    model: openai('gpt-4o'),
+    tools: {
+      transferToAgentB: tool({
+        description: 'Transfer to agent B.',
+        parameters: z.object({}),
+        execute: async () => {
+          activeAgent = agentB;
+          return 'Transferred to agent B.';
+        },
+      }),
+    },
+    stopWhen: stepCountIs(5),
+    prepareStep: () => activeAgent,
+    prompt: 'I want to talk to agent B.',
+  });
+
+  for await (const chunk of result.textStream) {
+    process.stdout.write(chunk);
+  }
+}

packages/ai/core/generate-text/generate-text.test.ts

@@ -828,6 +828,10 @@ describe('options.stopWhen', () => {
               ]);
 
               expect(prompt).toStrictEqual([
+                {
+                  role: 'system',
+                  content: 'system-message-0',
+                },
                 {
                   role: 'user',
                   content: [{ type: 'text', text: 'test-input' }],
@@ -873,6 +877,10 @@ describe('options.stopWhen', () => {
               expect(toolChoice).toStrictEqual({ type: 'auto' });
 
               expect(prompt).toStrictEqual([
+                {
+                  role: 'system',
+                  content: 'system-message-1',
+                },
                 {
                   role: 'user',
                   content: [{ type: 'text', text: 'test-input' }],
@@ -955,12 +963,17 @@ describe('options.stopWhen', () => {
                 type: 'tool',
                 toolName: 'tool1' as const,
               },
+              system: 'system-message-0',
             };
           }
 
           if (stepNumber === 1) {
             expect(steps.length).toStrictEqual(1);
-            return { model: trueModel, activeTools: [] };
+            return {
+              model: trueModel,
+              activeTools: [],
+              system: 'system-message-1',
+            };
           }
         },
       });

packages/ai/core/generate-text/generate-text.ts

@@ -280,7 +280,7 @@ A function that attempts to repair a tool call that failed to parse.
 
         const promptMessages = await convertToLanguageModelPrompt({
           prompt: {
-            system: initialPrompt.system,
+            system: prepareStepResult?.system ?? initialPrompt.system,
             messages: stepInputMessages,
           },
           supportedUrls: await model.supportedUrls,

packages/ai/core/generate-text/index.ts

@@ -6,7 +6,7 @@ export type {
   GeneratedFile,
 } from './generated-file';
 export * as Output from './output';
-export type { PrepareStepFunction } from './prepare-step';
+export type { PrepareStepFunction, PrepareStepResult } from './prepare-step';
 export { smoothStream, type ChunkDetector } from './smooth-stream';
 export type { StepResult } from './step-result';
 export { hasToolCall, stepCountIs, type StopCondition } from './stop-condition';

packages/ai/core/generate-text/prepare-step.ts

@@ -19,18 +19,15 @@ export type PrepareStepFunction<
   steps: Array<StepResult<NoInfer<TOOLS>>>;
   stepNumber: number;
   model: LanguageModel;
-}) =>
-  | PromiseLike<
-      | {
-          model?: LanguageModel;
-          toolChoice?: ToolChoice<NoInfer<TOOLS>>;
-          activeTools?: Array<keyof NoInfer<TOOLS>>;
-        }
-      | undefined
-    >
+}) => PromiseLike<PrepareStepResult<TOOLS>> | PrepareStepResult<TOOLS>;
+
+export type PrepareStepResult<
+  TOOLS extends Record<string, Tool> = Record<string, Tool>,
+> =
   | {
       model?: LanguageModel;
       toolChoice?: ToolChoice<NoInfer<TOOLS>>;
       activeTools?: Array<keyof NoInfer<TOOLS>>;
+      system?: string;
     }
   | undefined;

packages/ai/core/generate-text/stream-text.test.ts

@@ -2887,11 +2887,15 @@ describe('streamText', () => {
                   type: 'tool',
                   toolName: 'tool1' as const,
                 },
+                system: 'system-message-0',
               };
             }
 
             if (stepNumber === 1) {
-              return { activeTools: [] };
+              return {
+                activeTools: [],
+                system: 'system-message-1',
+              };
             }
           },
         });
@@ -2908,6 +2912,10 @@ describe('streamText', () => {
               "maxOutputTokens": undefined,
               "presencePenalty": undefined,
               "prompt": [
+                {
+                  "content": "system-message-0",
+                  "role": "system",
+                },
                 {
                   "content": [
                     {
@@ -2958,6 +2966,10 @@ describe('streamText', () => {
               "maxOutputTokens": undefined,
               "presencePenalty": undefined,
               "prompt": [
+                {
+                  "content": "system-message-1",
+                  "role": "system",
+                },
                 {
                   "content": [
                     {

packages/ai/core/generate-text/stream-text.ts

@@ -856,7 +856,7 @@ class DefaultStreamTextResult<TOOLS extends ToolSet, OUTPUT, PARTIAL_OUTPUT>
 
           const promptMessages = await convertToLanguageModelPrompt({
             prompt: {
-              system: initialPrompt.system,
+              system: prepareStepResult?.system ?? initialPrompt.system,
               messages: stepInputMessages,
             },
             supportedUrls: await model.supportedUrls,