feat: add tokenlens integration for language model usage tracking (#1161)
This commit is contained in:
parent
7327c9a18a
commit
1e539b65b3
8 changed files with 363 additions and 6 deletions
|
|
@ -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();
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue