Record Token Usage After Streaming Object
Learn how to record token usage when streaming structured data using the AI TOOLKIT and Node
When you're streaming structured data with `streamObject`, you may want to record the token usage for billing purposes.
Run it locally
$ npm install ai`onFinish` Callback
You can use the onFinish callback to record token usage.
It is called when the stream is finished.
import { streamObject } from 'ai-toolkit';
import { z } from 'zod';
const result = streamObject({
model: 'openai/gpt-4.1',
schema: z.object({
recipe: z.object({
name: z.string(),
ingredients: z.array(z.string()),
steps: z.array(z.string()),
}),
}),
prompt: 'Generate a lasagna recipe.',
onFinish({ usage }) {
console.log('Token usage:', usage);
},
});
- 1import { streamObject } from 'ai-toolkit';
- 2import { z } from 'zod';
- 3const result = streamObject({
- 4 model: 'openai/gpt-4.1',
- 5 schema: z.object({
- 6 recipe: z.object({
- 7 name: z.string(),
- 8 ingredients: z.array(z.string()),
- 9 steps: z.array(z.string()),
- 10 }),
- 11 }),
- 12 prompt: 'Generate a lasagna recipe.',
- 13 onFinish({ usage }) {
- 14 console.log('Token usage:', usage);
- 15 },
- 16});
`usage` Promise
The `streamObject` result contains a usage promise that resolves to the total token usage.
import { streamObject, LanguageModelUsage } from 'ai-toolkit';
import { z } from 'zod';
const result = streamObject({
model: 'openai/gpt-4.1',
schema: z.object({
recipe: z.object({
name: z.string(),
ingredients: z.array(z.string()),
steps: z.array(z.string()),
}),
}),
prompt: 'Generate a lasagna recipe.',
});
// your custom function to record token usage:
function recordUsage({
inputTokens,
outputTokens,
totalTokens,
}: LanguageModelUsage) {
console.log('Prompt tokens:', inputTokens);
console.log('Completion tokens:', outputTokens);
console.log('Total tokens:', totalTokens);
}
// use as promise:
result.usage.then(recordUsage);
// use with async/await:
recordUsage(await result.usage);
// note: the stream needs to be consumed because of backpressure
for await (const partialObject of result.partialObjectStream) {
}
- 1import { streamObject, LanguageModelUsage } from 'ai-toolkit';
- 2import { z } from 'zod';
- 3const result = streamObject({
- 4 model: 'openai/gpt-4.1',
- 5 schema: z.object({
- 6 recipe: z.object({
- 7 name: z.string(),
- 8 ingredients: z.array(z.string()),
- 9 steps: z.array(z.string()),
- 10 }),
- 11 }),
- 12 prompt: 'Generate a lasagna recipe.',
- 13});
- 14// your custom function to record token usage:
- 15function recordUsage({
- 16 inputTokens,
- 17 outputTokens,
- 18 totalTokens,
- 19}: LanguageModelUsage) {
- 20 console.log('Prompt tokens:', inputTokens);
- 21 console.log('Completion tokens:', outputTokens);
- 22 console.log('Total tokens:', totalTokens);
- 23}
- 24// use as promise:
- 25result.usage.then(recordUsage);
- 26// use with async/await:
- 27recordUsage(await result.usage);
- 28// note: the stream needs to be consumed because of backpressure
- 29for await (const partialObject of result.partialObjectStream) {
- 30}