Skip to main content

Improving UX with Multi-Step Calls

It would be nice if the model could summarize the action too. However, technically, once the model calls a tool, it has completed its generation as it ‘generated’ a tool call. How could you achieve this desired behavior?
The AI SDK has a feature called stopWhen which allows stopping conditions when the model generates a tool call. If those stopping conditions haven’t been hit, the AI SDK will automatically send tool call results back to the model!

Configuring Multi-Step Processing

Open your route handler (app/api/chat/route.ts) and add the following key to the streamText configuration object:
app/api/chat/route.ts
import { createResource } from "@/lib/actions/resources";
import { openai } from "@ai-sdk/openai";
import {
  convertToModelMessages,
  streamText,
  tool,
  UIMessage,
  stepCountIs,
} from "ai";
import { z } from "zod";

// Allow streaming responses up to 30 seconds
export const maxDuration = 30;

export async function POST(req: Request) {
  const { messages }: { messages: UIMessage[] } = await req.json();

  const result = streamText({
    model: openai("gpt-4o"),
    system: `You are a helpful assistant. Check your knowledge base before answering any questions.
    Only respond to questions using information from tool calls.
    if no relevant information is found in the tool calls, respond, "Sorry, I don't know."`,
    messages: convertToModelMessages(messages),
    stopWhen: stepCountIs(5),
    tools: {
      addResource: tool({
        description: `add a resource to your knowledge base.
          If the user provides a random piece of knowledge unprompted, use this tool without asking for confirmation.`,
        inputSchema: z.object({
          content: z
            .string()
            .describe("the content or resource to add to the knowledge base"),
        }),
        execute: async ({ content }) => createResource({ content }),
      }),
    },
  });

  return result.toUIMessageStreamResponse();
}

How stopWhen Works

1

Tool Call Generation

The model generates a tool call and completes its initial generation.
2

Tool Execution

The AI SDK automatically executes the tool with the provided parameters.
3

Result Processing

Tool results are sent back to the model for additional processing.
4

Continued Generation

The model can now provide additional context, confirmation, or summary.

Understanding Step Count

The stepCountIs(5) configuration means the model can continue generating for up to 5 steps after the initial tool call. This allows for:

Action Confirmation

User Feedback: The model can confirm what action was taken and provide context.

Additional Information

Enhanced Responses: The model can add relevant details or suggestions.

Error Handling

Graceful Recovery: The model can handle tool execution errors and provide alternatives.

Context Building

Knowledge Expansion: The model can build upon the tool results with additional insights.

Testing Multi-Step Functionality

1

Test the Enhanced Behavior

Head back to the browser and tell the model your favorite pizza topping (note: pineapple is not an option).
2

Observe the Response

You should see a follow-up response from the model confirming the action and providing additional context.
3

Verify Tool Call Flow

The model will first call the tool, then provide a summary or confirmation message.
With multi-step processing, the AI can now provide a much better user experience by explaining what it did and why, rather than just executing the tool silently.

Alternative Stopping Conditions

You can configure different stopping conditions based on your needs:
// Stop after a specific number of steps
stopWhen: stepCountIs(3);

// Stop when the model generates a specific pattern
stopWhen: textMatches(/^I have|^The tool/);

// Stop when the model reaches a certain token limit
stopWhen: tokenCountIs(100);

// Combine multiple conditions
stopWhen: [stepCountIs(5), textMatches(/^Summary:/)];

Benefits of Multi-Step Processing

Better User Experience

Clear Communication: Users understand what actions were taken and why.

Error Recovery

Graceful Handling: The model can handle tool failures and provide alternatives.

Context Building

Rich Responses: The model can build upon tool results with additional insights.

Action Confirmation

User Confidence: Users receive confirmation that their requests were processed correctly.

Testing Different Scenarios

1

Test Resource Addition

Try telling the model various pieces of information to see how it handles different types of content.
2

Observe Response Patterns

Notice how the model confirms actions and provides context after tool execution.
3

Check Database Updates

Verify that resources are being properly added to your knowledge base.
Remember that the model will still respond “Sorry, I don’t know” to questions it can’t answer, as we haven’t implemented the retrieval tool yet.

Extension task