streamText Multi-Step Cookbook
Learn how to create several streamText steps with different settings
You may want to have different steps in your stream where each step has different settings, e.g. models, tools, or system prompts. With createUIMessageStream and sendFinish / sendStart options when merging into the UIMessageStream, you can control when the finish and start events are sent to the client, allowing you to have different steps in a single assistant UI message.
Run it locally
$ npm install aiServer
import {
convertToModelMessages,
createUIMessageStream,
createUIMessageStreamResponse,
streamText,
tool,
} from 'ai-toolkit';
import { z } from 'zod';
export async function POST(req: Request) {
const { messages } = await req.json();
const stream = createUIMessageStream({
execute: async ({ writer }) => {
// step 1 example: forced tool call
const result1 = streamText({
model: 'openai/gpt-4o-mini',
system: 'Extract the user goal from the conversation.',
messages,
toolChoice: 'required', // force the model to call a tool
tools: {
extractGoal: tool({
inputSchema: z.object({ goal: z.string() }),
execute: async ({ goal }) => goal, // no-op extract tool
}),
},
});
// forward the initial result to the client without the finish event:
writer.merge(result1.toUIMessageStream({ sendFinish: false }));
// note: you can use any programming construct here, e.g. if-else, loops, etc.
// workflow programming is normal programming with this approach.
// example: continue stream with forced tool call from previous step
const result2 = streamText({
// different system prompt, different model, no tools:
model: 'openai/gpt-4o',
system:
'You are a helpful assistant with a different system prompt. Repeat the extract user goal in your answer.',
// continue the workflow stream with the messages from the previous step:
messages: [
...convertToModelMessages(messages),
...(await result1.response).messages,
],
});
// forward the 2nd result to the client (incl. the finish event):
writer.merge(result2.toUIMessageStream({ sendStart: false }));
},
});
return createUIMessageStreamResponse({ stream });
}
- 1import {
- 2 convertToModelMessages,
- 3 createUIMessageStream,
- 4 createUIMessageStreamResponse,
- 5 streamText,
- 6 tool,
- 7} from 'ai-toolkit';
- 8import { z } from 'zod';
- 9export async function POST(req: Request) {
- 10 const { messages } = await req.json();
- 11 const stream = createUIMessageStream({
- 12 execute: async ({ writer }) => {
- 13 // step 1 example: forced tool call
- 14 const result1 = streamText({
- 15 model: 'openai/gpt-4o-mini',
- 16 system: 'Extract the user goal from the conversation.',
- 17 messages,
- 18 toolChoice: 'required', // force the model to call a tool
- 19 tools: {
- 20 extractGoal: tool({
- 21 inputSchema: z.object({ goal: z.string() }),
- 22 execute: async ({ goal }) => goal, // no-op extract tool
- 23 }),
- 24 },
- 25 });
- 26 // forward the initial result to the client without the finish event:
- 27 writer.merge(result1.toUIMessageStream({ sendFinish: false }));
- 28 // note: you can use any programming construct here, e.g. if-else, loops, etc.
- 29 // workflow programming is normal programming with this approach.
- 30 // example: continue stream with forced tool call from previous step
- 31 const result2 = streamText({
- 32 // different system prompt, different model, no tools:
- 33 model: 'openai/gpt-4o',
- 34 system:
- 35 'You are a helpful assistant with a different system prompt. Repeat the extract user goal in your answer.',
- 36 // continue the workflow stream with the messages from the previous step:
- 37 messages: [
- 38 ...convertToModelMessages(messages),
- 39 ...(await result1.response).messages,
- 40 ],
- 41 });
- 42 // forward the 2nd result to the client (incl. the finish event):
- 43 writer.merge(result2.toUIMessageStream({ sendStart: false }));
- 44 },
- 45 });
- 46 return createUIMessageStreamResponse({ stream });
- 47}
Client
'use client';
import { useChat } from '@ai-toolkit/react';
import { useState } from 'react';
export default function Chat() {
const [input, setInput] = useState('');
const { messages, sendMessage } = useChat();
return (
<div>
{messages?.map(message => (
<div key={message.id}>
<strong>{${message.role}: }</strong>
{message.parts.map((part, index) => {
switch (part.type) {
case 'text':
return <span key={index}>{part.text}</span>;
case 'tool-extractGoal': {
return <pre key={index}>{JSON.stringify(part, null, 2)}</pre>;
}
}
})}
</div>
))}
<form
onSubmit={e => {
e.preventDefault();
sendMessage({ text: input });
setInput('');
}}
>
<input value={input} onChange={e => setInput(e.currentTarget.value)} />
</form>
</div>
);
}
- 1'use client';
- 2import { useChat } from '@ai-toolkit/react';
- 3import { useState } from 'react';
- 4export default function Chat() {
- 5 const [input, setInput] = useState('');
- 6 const { messages, sendMessage } = useChat();
- 7 return (
- 8 <div>
- 9 {messages?.map(message => (
- 10 <div key={message.id}>
- 11 <strong>{`${message.role}: `}</strong>
- 12 {message.parts.map((part, index) => {
- 13 switch (part.type) {
- 14 case 'text':
- 15 return <span key={index}>{part.text}</span>;
- 16 case 'tool-extractGoal': {
- 17 return <pre key={index}>{JSON.stringify(part, null, 2)}</pre>;
- 18 }
- 19 }
- 20 })}
- 21 </div>
- 22 ))}
- 23 <form
- 24 onSubmit={e => {
- 25 e.preventDefault();
- 26 sendMessage({ text: input });
- 27 setInput('');
- 28 }}
- 29 >
- 30 <input value={input} onChange={e => setInput(e.currentTarget.value)} />
- 31 </form>
- 32 </div>
- 33 );
- 34}