From cf265515cb7a59b651ba81c38b6edfda53083191 Mon Sep 17 00:00:00 2001
From: Connor Fogarty <40175891+connorff@users.noreply.github.com>
Date: Thu, 15 Aug 2024 00:19:12 -0500
Subject: [PATCH] feat: send analytics event for token usage (#37)

---
 .../bolt/app/lib/.server/llm/stream-text.ts     |  3 ---
 .../app/lib/.server/llm/switchable-stream.ts    |  2 --
 packages/bolt/app/lib/analytics.ts              |  1 +
 packages/bolt/app/routes/api.chat.ts            | 17 ++++++++++++++++-
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/packages/bolt/app/lib/.server/llm/stream-text.ts b/packages/bolt/app/lib/.server/llm/stream-text.ts
index b1149e4..cf937fd 100644
--- a/packages/bolt/app/lib/.server/llm/stream-text.ts
+++ b/packages/bolt/app/lib/.server/llm/stream-text.ts
@@ -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,
   });
 }
diff --git a/packages/bolt/app/lib/.server/llm/switchable-stream.ts b/packages/bolt/app/lib/.server/llm/switchable-stream.ts
index 1f351a2..638db7a 100644
--- a/packages/bolt/app/lib/.server/llm/switchable-stream.ts
+++ b/packages/bolt/app/lib/.server/llm/switchable-stream.ts
@@ -24,8 +24,6 @@ export default class SwitchableStream extends TransformStream {
       await this._currentReader.cancel();
     }
 
-    console.log('Switching stream');
-
     this._currentReader = newStream.getReader();
 
     this._pumpStream();
diff --git a/packages/bolt/app/lib/analytics.ts b/packages/bolt/app/lib/analytics.ts
index df41c3e..bf1bcdc 100644
--- a/packages/bolt/app/lib/analytics.ts
+++ b/packages/bolt/app/lib/analytics.ts
@@ -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`,
 }
 
diff --git a/packages/bolt/app/routes/api.chat.ts b/packages/bolt/app/routes/api.chat.ts
index dbf1a7e..f20f50c 100644
--- a/packages/bolt/app/routes/api.chat.ts
+++ b/packages/bolt/app/routes/api.chat.ts
@@ -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();
         }