feat: send analytics event for token usage (#37)

This commit is contained in:
Connor Fogarty 2024-08-15 00:19:12 -05:00 committed by GitHub
parent fb2d957351
commit cf265515cb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 17 additions and 6 deletions

View File

@ -30,9 +30,6 @@ export function streamText(messages: Messages, env: Env, options?: StreamingOpti
'anthropic-beta': 'max-tokens-3-5-sonnet-2024-07-15',
},
messages: convertToCoreMessages(messages),
onFinish: ({ finishReason, usage, warnings }) => {
console.log({ finishReason, usage, warnings });
},
...options,
});
}

View File

@ -24,8 +24,6 @@ export default class SwitchableStream extends TransformStream {
await this._currentReader.cancel();
}
console.log('Switching stream');
this._currentReader = newStream.getReader();
this._pumpStream();

View File

@ -13,6 +13,7 @@ const MESSAGE_PREFIX = 'Bolt';
export enum AnalyticsTrackEvent {
MessageSent = `${MESSAGE_PREFIX} Message Sent`,
MessageComplete = `${MESSAGE_PREFIX} Message Complete`,
ChatCreated = `${MESSAGE_PREFIX} Chat Created`,
}

View File

@ -4,6 +4,8 @@ import { CONTINUE_PROMPT } from '~/lib/.server/llm/prompts';
import { streamText, type Messages, type StreamingOptions } from '~/lib/.server/llm/stream-text';
import SwitchableStream from '~/lib/.server/llm/switchable-stream';
import { handleWithAuth } from '~/lib/.server/login';
import { getSession } from '~/lib/.server/sessions';
import { AnalyticsAction, AnalyticsTrackEvent, sendEventInternal } from '~/lib/analytics';
export async function action(args: ActionFunctionArgs) {
return handleWithAuth(args, chatAction);
@ -17,8 +19,21 @@ async function chatAction({ context, request }: ActionFunctionArgs) {
try {
const options: StreamingOptions = {
toolChoice: 'none',
onFinish: async ({ text: content, finishReason }) => {
onFinish: async ({ text: content, finishReason, usage }) => {
if (finishReason !== 'length') {
const { session } = await getSession(request, context.cloudflare.env);
await sendEventInternal(session.data, {
action: AnalyticsAction.Track,
payload: {
event: AnalyticsTrackEvent.MessageComplete,
properties: {
usage,
finishReason,
},
},
});
return stream.close();
}