feat: use real openai data for token usage calculation

WiseGPT · Jan 6, 2023 · 8e99161 · 8e99161
1 parent bcf22c9
commit 8e99161
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 16 deletions.
diff --git a/src/application/open-ai/open-ai-command.handler.ts b/src/application/open-ai/open-ai-command.handler.ts
@@ -40,7 +40,8 @@ export class OpenAICommandHandler {
         message: text,
         // We rely on the fact that only 1 completion is done, this number could be wrong
         // if we used `best_of` and `n` parameters.
-        tokens: usage.completionTokens,
+        messageTokens: usage.completionTokens,
+        totalTokensSpent: usage.totalTokens,
       });
     } catch (err: any) {
       await this.conversationCommandBus.send({

diff --git a/src/config.ts b/src/config.ts
@@ -10,7 +10,9 @@ export default {
   conversation: {
     // which persona to use for conversations see `src/domain/persona/index.ts`
     personaConfigName: "slack-software-eng",
-    // ends the conversation when total conversation tokens goes over `maxConversationTokens`
-    maxConversationTokens: Number.MAX_SAFE_INTEGER,
+    // ends the conversation when sum of total tokens goes over `maximumSpentTokens`
+    // all summarization and completion requests are counted towards the total tokens
+    // persona adds overhead to each request
+    maximumSpentTokens: Number.MAX_SAFE_INTEGER,
   },
 };
diff --git a/src/domain/conversation/conversation.aggregate.ts b/src/domain/conversation/conversation.aggregate.ts
@@ -37,7 +37,7 @@ export class ConversationAggregate {
   private status: ConversationStatus;
   private readonly messages: ConversationMessage[];
   private aiStatus: ConversationAIStatus;
-  private totalTokens: number;
+  private totalTokensSpent: number;
 
   private newEvents: ConversationEvent[] = [];
 
@@ -49,7 +49,7 @@ export class ConversationAggregate {
     this.status = { status: "ONGOING" };
     this.messages = [];
     this.aiStatus = { status: "IDLE" };
-    this.totalTokens = 0;
+    this.totalTokensSpent = 0;
 
     if (createEvent) {
       this.createAndApply(createEvent);
@@ -121,8 +121,9 @@ export class ConversationAggregate {
           message: {
             id: cmd.correlationId,
             text: cmd.message,
-            tokens: cmd.tokens,
+            tokens: cmd.messageTokens,
           },
+          totalTokensSpent: cmd.totalTokensSpent,
         });
 
         this.endConversationIfWentOverLimit();
@@ -166,14 +167,14 @@ export class ConversationAggregate {
   }
 
   private endConversationIfWentOverLimit(): boolean {
-    if (this.totalTokens > config.conversation.maxConversationTokens) {
+    if (this.totalTokensSpent > config.conversation.maximumSpentTokens) {
       this.createAndApply({
         type: "CONVERSATION_ENDED",
         conversationId: this.conversationId,
         reason: {
           type: "MAXIMUM_CONVERSATION_TOKENS_REACHED",
-          maxConversationTokens: config.conversation.maxConversationTokens,
-          totalTokens: this.totalTokens,
+          maximumSpentTokens: config.conversation.maximumSpentTokens,
+          totalTokensSpent: this.totalTokensSpent,
         },
       });
 
@@ -185,13 +186,12 @@ export class ConversationAggregate {
 
   private addConversationMessage({
     message,
-    tokens,
-  }: {
+  }: // tokens,
+  {
     message: ConversationMessage;
     tokens: number;
   }): void {
     this.messages.push(message);
-    this.totalTokens += tokens;
   }
 
   private createAndApply(
@@ -223,6 +223,7 @@ export class ConversationAggregate {
     switch (event.type) {
       case "BOT_RESPONSE_ADDED": {
         this.aiStatus = { status: "IDLE" };
+        this.totalTokensSpent += event.totalTokensSpent;
 
         return this.addConversationMessage({
           message: {

diff --git a/src/domain/conversation/conversation.commands.ts b/src/domain/conversation/conversation.commands.ts
@@ -23,8 +23,8 @@ export type ProcessCompletionResponseCommand =
       botResponseType: "BOT_RESPONSE_SUCCESS";
       correlationId: string;
       message: string;
-      // how many tokens is the `message`
-      tokens: number;
+      messageTokens: number;
+      totalTokensSpent: number;
     })
   | (BaseCommand & {
       type: "PROCESS_COMPLETION_RESPONSE_COMMAND";

diff --git a/src/domain/conversation/conversation.events.ts b/src/domain/conversation/conversation.events.ts
@@ -31,15 +31,16 @@ export type BotResponseAdded = BaseEvent & {
     text: string;
     tokens: number;
   };
+  totalTokensSpent: number;
 };
 
 export type ConversationEnded = BaseEvent & {
   type: "CONVERSATION_ENDED";
   reason:
     | {
         type: "MAXIMUM_CONVERSATION_TOKENS_REACHED";
-        maxConversationTokens: number;
-        totalTokens: number;
+        maximumSpentTokens: number;
+        totalTokensSpent: number;
       }
     | {
         type: "BOT_RESPONSE_ERROR";