#477 Estimate tokens usage with streaming completions

transitive-bullshit · Mar 29, 2023 · dd6b457 · dd6b457
1 parent 4613b95
commit dd6b457
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 4 deletions.
diff --git 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
@@ -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