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.
This commit is contained in:
kangfenmao 2024-09-17 14:56:10 +08:00
parent 0af35b9f10
commit 6f5dccd595
2 changed files with 22 additions and 2 deletions

View File

@ -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'

View File

@ -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)