📄 ai-sdk/cookbook/rsc/stream-ui-record-token-usage

======================================================================================================================================================================

When you're streaming structured data with streamUI , you may want to record the token usage for billing purposes.

onFinish Callback

You can use the onFinish callback to record token usage. It is called when the stream is finished.

app/page.tsx

'use client';
import { useState } from 'react';import { ClientMessage } from './actions';import { useActions, useUIState } from '@ai-sdk/rsc';import { generateId } from 'ai';
// Allow streaming responses up to 30 secondsexport const maxDuration = 30;
export default function Home() {  const [input, setInput] = useState<string>('');  const [conversation, setConversation] = useUIState();  const { continueConversation } = useActions();
  return (    <div>      <div>        {conversation.map((message: ClientMessage) => (          <div key={message.id}>            {message.role}: {message.display}          </div>        ))}      </div>
      <div>        <input          type="text"          value={input}          onChange={event => {            setInput(event.target.value);          }}        />        <button          onClick={async () => {            setConversation((currentConversation: ClientMessage[]) => [              ...currentConversation,              { id: generateId(), role: 'user', display: input },            ]);
            const message = await continueConversation(input);
            setConversation((currentConversation: ClientMessage[]) => [              ...currentConversation,              message,            ]);          }}        >          Send Message        </button>      </div>    </div>  );}

Server

app/actions.tsx

'use server';
import { createAI, getMutableAIState, streamUI } from '@ai-sdk/rsc';import { openai } from '@ai-sdk/openai';import { ReactNode } from 'react';import { z } from 'zod';import { generateId } from 'ai';
export interface ServerMessage {  role: 'user' | 'assistant';  content: string;}
export interface ClientMessage {  id: string;  role: 'user' | 'assistant';  display: ReactNode;}
export async function continueConversation(  input: string,): Promise<ClientMessage> {  'use server';
  const history = getMutableAIState();
  const result = await streamUI({    model: openai('gpt-3.5-turbo'),    messages: [...history.get(), { role: 'user', content: input }],    text: ({ content, done }) => {      if (done) {        history.done((messages: ServerMessage[]) => [          ...messages,          { role: 'assistant', content },        ]);      }
      return <div>{content}</div>;    },    tools: {      deploy: {        description: 'Deploy repository to vercel',        inputSchema: z.object({          repositoryName: z            .string()            .describe('The name of the repository, example: vercel/ai-chatbot'),        }),        generate: async function* ({ repositoryName }) {          yield <div>Cloning repository {repositoryName}...</div>; // [!code highlight:5]          await new Promise(resolve => setTimeout(resolve, 3000));          yield <div>Building repository {repositoryName}...</div>;          await new Promise(resolve => setTimeout(resolve, 2000));          return <div>{repositoryName} deployed!</div>;        },      },    },    onFinish: ({ usage }) => {      const { promptTokens, completionTokens, totalTokens } = usage;      // your own logic, e.g. for saving the chat history or recording usage      console.log('Prompt tokens:', promptTokens);      console.log('Completion tokens:', completionTokens);      console.log('Total tokens:', totalTokens);    },  });
  return {    id: generateId(),    role: 'assistant',    display: result.value,  };}

app/ai.ts

import { createAI } from '@ai-sdk/rsc';import { ServerMessage, ClientMessage, continueConversation } from './actions';
export const AI = createAI<ServerMessage[], ClientMessage[]>({  actions: {    continueConversation,  },  initialAIState: [],  initialUIState: [],});

On this page

Record Token Usage after Streaming User Interfaces

onFinish Callback

Server

Deploy and Scale AI Apps with Vercel.

Vercel delivers the infrastructure and developer experience you need to ship reliable AI-powered applications at scale.

Trusted by industry leaders:

OpenAI
Photoroom

Talk to an expert