From 6f5dccd595ef7f4d90e6315281d4ee66670a5eab Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Tue, 17 Sep 2024 14:56:10 +0800
Subject: [PATCH] feat: estimate completion usage calculation added to chat.

- Estimated usage calculation has been added to chat completion fetching to track message usage.
- Added functionality to estimate completion usage tokens based on input and prompt data.
---
 src/renderer/src/services/api.ts      |  3 ++-
 src/renderer/src/services/messages.ts | 21 ++++++++++++++++++++-
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/src/renderer/src/services/api.ts b/src/renderer/src/services/api.ts
index 450b9032..39e00aec 100644
--- a/src/renderer/src/services/api.ts
+++ b/src/renderer/src/services/api.ts
@@ -15,7 +15,7 @@ import {
   getTranslateModel
 } from './assistant'
 import { EVENT_NAMES, EventEmitter } from './event'
-import { filterMessages } from './messages'
+import { estimateMessagesToken, filterMessages } from './messages'
 
 export async function fetchChatCompletion({
   messages,
@@ -67,6 +67,7 @@ export async function fetchChatCompletion({
       onResponse({ ...message, status: 'pending' })
     })
     message.status = 'success'
+    message.usage = message.usage || (await estimateMessagesToken({ assistant, messages: [...messages, message] }))
   } catch (error: any) {
     message.content = `Error: ${error.message}`
     message.status = 'error'
diff --git a/src/renderer/src/services/messages.ts b/src/renderer/src/services/messages.ts
index bc1a2351..86418687 100644
--- a/src/renderer/src/services/messages.ts
+++ b/src/renderer/src/services/messages.ts
@@ -1,7 +1,8 @@
 import { DEFAULT_CONEXTCOUNT } from '@renderer/config/constant'
 import { Assistant, Message } from '@renderer/types'
 import { GPTTokens } from 'gpt-tokens'
-import { isEmpty, takeRight } from 'lodash'
+import { isEmpty, last, takeRight } from 'lodash'
+import { CompletionUsage } from 'openai/resources'
 
 import { getAssistantSettings } from './assistant'
 import FileManager from './file'
@@ -44,6 +45,24 @@ export function estimateInputTokenCount(text: string) {
   return input.usedTokens - 7
 }
 
+export async function estimateMessagesToken({
+  assistant,
+  messages
+}: {
+  assistant: Assistant
+  messages: Message[]
+}): Promise<CompletionUsage> {
+  const responseMessageContent = last(messages)?.content
+  const inputMessageContent = messages[messages.length - 2]?.content
+  const completion_tokens = await estimateInputTokenCount(responseMessageContent ?? '')
+  const prompt_tokens = await estimateInputTokenCount(assistant.prompt + inputMessageContent ?? '')
+  return {
+    completion_tokens,
+    prompt_tokens: prompt_tokens,
+    total_tokens: prompt_tokens + completion_tokens
+  } as CompletionUsage
+}
+
 export function estimateHistoryTokenCount(assistant: Assistant, msgs: Message[]) {
   const { contextCount } = getAssistantSettings(assistant)