feat: add tokenlens integration for language model usage tracking (#1161)

This commit is contained in:
Nicklas Scharpff 2025-09-08 16:56:10 +02:00 committed by GitHub
parent 7327c9a18a
commit 1e539b65b3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 363 additions and 6 deletions

View file

@ -2,6 +2,7 @@ import {
convertToModelMessages,
createUIMessageStream,
JsonToSseTransformStream,
LanguageModelUsage,
smoothStream,
stepCountIs,
streamText,
@ -149,6 +150,8 @@ export async function POST(request: Request) {
const streamId = generateUUID();
await createStreamId({ streamId, chatId: id });
let finalUsage: LanguageModelUsage | undefined;
const stream = createUIMessageStream({
execute: ({ writer: dataStream }) => {
const result = streamText({
@ -179,6 +182,10 @@ export async function POST(request: Request) {
isEnabled: isProductionEnvironment,
functionId: 'stream-text',
},
onFinish: ({ usage }) => {
finalUsage = usage;
dataStream.write({ type: 'data-usage', data: usage });
},
});
result.consumeStream();