From dd6b4579326064258f6eb56f0a3ab2cd1c71f53b Mon Sep 17 00:00:00 2001
From: Oleksii Myronenko <mironenkoalexey@gmail.com>
Date: Wed, 29 Mar 2023 23:34:06 +0300
Subject: [PATCH] #477 Estimate tokens usage with streaming completions

---
 src/chatgpt-api.ts | 16 +++++++++++++---
 src/types.ts       | 14 +++++++++++++-
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/src/chatgpt-api.ts b/src/chatgpt-api.ts
index b44cce2d..3ae23a06 100644
--- a/src/chatgpt-api.ts
+++ b/src/chatgpt-api.ts
@@ -211,16 +211,17 @@ export class ChatGPTAPI {
                     result.id = response.id
                   }
 
-                  if (response?.choices?.length) {
+                  if (response.choices?.length) {
                     const delta = response.choices[0].delta
                     result.delta = delta.content
                     if (delta?.content) result.text += delta.content
-                    result.detail = response
 
                     if (delta.role) {
                       result.role = delta.role
                     }
 
+                    result.detail = response
+
                     onProgress?.(result)
                   }
                 } catch (err) {
@@ -286,7 +287,16 @@ export class ChatGPTAPI {
           }
         }
       }
-    ).then((message) => {
+    ).then(async (message) => {
+      if (message.detail && !message.detail.usage) {
+        const promptTokens = numTokens
+        const completionTokens = await this._getTokenCount(message.text)
+        message.detail.usage = {
+          prompt_tokens: promptTokens,
+          completion_tokens: completionTokens,
+          total_tokens: promptTokens + completionTokens
+        }
+      }
       return this._upsertMessage(message).then(() => message)
     })
 
diff --git a/src/types.ts b/src/types.ts
index fd23cf7b..cb6bbfa3 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -59,13 +59,25 @@ export type SendMessageBrowserOptions = {
   abortSignal?: AbortSignal
 }
 
+interface CreateChatCompletionStreamResponse
+  extends openai.CreateChatCompletionDeltaResponse {
+  usage: CreateCompletionStreamResponseUsage
+}
+
+interface CreateCompletionStreamResponseUsage
+  extends openai.CreateCompletionResponseUsage {
+  estimated: true
+}
+
 export interface ChatMessage {
   id: string
   text: string
   role: Role
   name?: string
   delta?: string
-  detail?: any
+  detail?:
+    | openai.CreateChatCompletionResponse
+    | CreateChatCompletionStreamResponse
 
   // relevant for both ChatGPTAPI and ChatGPTUnofficialProxyAPI
   parentMessageId?: string