AI TOOLKIT
Get started
Recipes/Next.js

Call Tools in Multiple Steps

Learn how to call tools in multiple steps using the AI TOOLKIT and Next.js

3 min readnextstreamingtool useView source

Some language models are great at calling tools in multiple steps to achieve a more complex task. This is particularly useful when the tools are dependent on each other and need to be executed in sequence during the same generation step.

Run it locally

$ npm install ai

Client

Let's create a React component that imports the useChat hook from the @ai-toolkit/react module. The useChat hook will call the /api/chat endpoint when the user sends a message. The endpoint will generate the assistant's response based on the conversation history and stream it to the client. If the assistant responds with a tool call, the hook will automatically display them as well.

'use client';

import { useChat } from '@ai-toolkit/react';

import { DefaultChatTransport } from 'ai-toolkit';

import { useState } from 'react';

import type { ChatMessage } from './api/chat/route';

export default function Page() {

const [input, setInput] = useState('');

const { messages, sendMessage } = useChat<ChatMessage>({

transport: new DefaultChatTransport({

api: '/api/chat',

}),

});

return (

<div>

<input

className="border"

value={input}

onChange={event => {

setInput(event.target.value);

}}

onKeyDown={async event => {

if (event.key === 'Enter') {

sendMessage({

text: input,

});

setInput('');

}

}}

/>

{messages.map((message, index) => (

<div key={index}>

{message.parts.map((part, i) => {

switch (part.type) {

case 'text':

return <div key={${message.id}-text}>{part.text}</div>;

case 'tool-getLocation':

case 'tool-getWeather':

return (

<div key={${message.id}-weather-${i}}>

{JSON.stringify(part, null, 2)}

</div>

);

}

})}

</div>

))}

</div>

);

}

app/page.tsx
tsx
  1. 1'use client';
  2. 2import { useChat } from '@ai-toolkit/react';
  3. 3import { DefaultChatTransport } from 'ai-toolkit';
  4. 4import { useState } from 'react';
  5. 5import type { ChatMessage } from './api/chat/route';
  6. 6export default function Page() {
  7. 7 const [input, setInput] = useState('');
  8. 8 const { messages, sendMessage } = useChat<ChatMessage>({
  9. 9 transport: new DefaultChatTransport({
  10. 10 api: '/api/chat',
  11. 11 }),
  12. 12 });
  13. 13 return (
  14. 14 <div>
  15. 15 <input
  16. 16 className="border"
  17. 17 value={input}
  18. 18 onChange={event => {
  19. 19 setInput(event.target.value);
  20. 20 }}
  21. 21 onKeyDown={async event => {
  22. 22 if (event.key === 'Enter') {
  23. 23 sendMessage({
  24. 24 text: input,
  25. 25 });
  26. 26 setInput('');
  27. 27 }
  28. 28 }}
  29. 29 />
  30. 30 {messages.map((message, index) => (
  31. 31 <div key={index}>
  32. 32 {message.parts.map((part, i) => {
  33. 33 switch (part.type) {
  34. 34 case 'text':
  35. 35 return <div key={`${message.id}-text`}>{part.text}</div>;
  36. 36 case 'tool-getLocation':
  37. 37 case 'tool-getWeather':
  38. 38 return (
  39. 39 <div key={`${message.id}-weather-${i}`}>
  40. 40 {JSON.stringify(part, null, 2)}
  41. 41 </div>
  42. 42 );
  43. 43 }
  44. 44 })}
  45. 45 </div>
  46. 46 ))}
  47. 47 </div>
  48. 48 );
  49. 49}

Server

You will create a new route at /api/chat that will use the streamText function from the ai module to generate the assistant's response based on the conversation history.

You will use the `tools` parameter to specify two tools called getLocation and getWeather that will first get the user's location and then use it to get the weather.

You will add the two functions mentioned earlier and use zod to specify the schema for its parameters.

To call tools in multiple steps, you can use the stopWhen option to specify the stopping conditions for when the model generates a tool call. In this example, you will set it to stepCountIs(5) to allow for multiple consecutive tool calls (steps).

import {

type InferUITools,

type ToolSet,

type UIDataTypes,

type UIMessage,

convertToModelMessages,

stepCountIs,

streamText,

tool,

} from 'ai-toolkit';

import { z } from 'zod';

const tools = {

getLocation: tool({

description: 'Get the location of the user',

inputSchema: z.object({}),

execute: async () => {

const location = { lat: 37.7749, lon: -122.4194 };

return Your location is at latitude ${location.lat} and longitude ${location.lon};

},

}),

getWeather: tool({

description: 'Get the weather for a location',

inputSchema: z.object({

city: z.string().describe('The city to get the weather for'),

unit: z

.enum(['C', 'F'])

.describe('The unit to display the temperature in'),

}),

execute: async ({ city, unit }) => {

const weather = {

value: 24,

description: 'Sunny',

};

return It is currently ${weather.value}°${unit} and ${weather.description} in ${city}!;

},

}),

} satisfies ToolSet;

export type ChatTools = InferUITools<typeof tools>;

export type ChatMessage = UIMessage<never, UIDataTypes, ChatTools>;

export async function POST(req: Request) {

const { messages }: { messages: ChatMessage[] } = await req.json();

const result = streamText({

model: 'openai/gpt-4o',

system: 'You are a helpful assistant.',

messages: await convertToModelMessages(messages),

stopWhen: stepCountIs(5),

tools,

});

return result.toUIMessageStreamResponse();

}

app/api/chat/route.ts
ts
  1. 1import {
  2. 2 type InferUITools,
  3. 3 type ToolSet,
  4. 4 type UIDataTypes,
  5. 5 type UIMessage,
  6. 6 convertToModelMessages,
  7. 7 stepCountIs,
  8. 8 streamText,
  9. 9 tool,
  10. 10} from 'ai-toolkit';
  11. 11import { z } from 'zod';
  12. 12const tools = {
  13. 13 getLocation: tool({
  14. 14 description: 'Get the location of the user',
  15. 15 inputSchema: z.object({}),
  16. 16 execute: async () => {
  17. 17 const location = { lat: 37.7749, lon: -122.4194 };
  18. 18 return `Your location is at latitude ${location.lat} and longitude ${location.lon}`;
  19. 19 },
  20. 20 }),
  21. 21 getWeather: tool({
  22. 22 description: 'Get the weather for a location',
  23. 23 inputSchema: z.object({
  24. 24 city: z.string().describe('The city to get the weather for'),
  25. 25 unit: z
  26. 26 .enum(['C', 'F'])
  27. 27 .describe('The unit to display the temperature in'),
  28. 28 }),
  29. 29 execute: async ({ city, unit }) => {
  30. 30 const weather = {
  31. 31 value: 24,
  32. 32 description: 'Sunny',
  33. 33 };
  34. 34 return `It is currently ${weather.value}°${unit} and ${weather.description} in ${city}!`;
  35. 35 },
  36. 36 }),
  37. 37} satisfies ToolSet;
  38. 38export type ChatTools = InferUITools<typeof tools>;
  39. 39export type ChatMessage = UIMessage<never, UIDataTypes, ChatTools>;
  40. 40export async function POST(req: Request) {
  41. 41 const { messages }: { messages: ChatMessage[] } = await req.json();
  42. 42 const result = streamText({
  43. 43 model: 'openai/gpt-4o',
  44. 44 system: 'You are a helpful assistant.',
  45. 45 messages: await convertToModelMessages(messages),
  46. 46 stopWhen: stepCountIs(5),
  47. 47 tools,
  48. 48 });
  49. 49 return result.toUIMessageStreamResponse();
  50. 50}
Browse all recipes