feat: add model name in chat usage records (#493)

* chore: optimize statistics chart render

* feat: add model name in chat usage records

* lint fix

Author: zhujunsan

service/src/index.ts

@@ -2,6 +2,7 @@ import express from 'express'
 import jwt from 'jsonwebtoken'
 import * as dotenv from 'dotenv'
 import { ObjectId } from 'mongodb'
+import type { TiktokenModel } from 'gpt-token'
 import { textTokens } from 'gpt-token'
 import speakeasy from 'speakeasy'
 import { TwoFAConfig } from './types'
@@ -383,6 +384,8 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
     globalThis.console.error(`Unable to get chat room \t ${userId}\t ${roomId}`)
   if (room != null && isNotEmptyString(room.prompt))
     systemMessage = room.prompt
+  const model = room.chatModel
+
   let lastResponse
   let result
   let message: ChatInfo
@@ -451,9 +454,9 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
       if (!result.data.detail)
         result.data.detail = {}
       result.data.detail.usage = new UsageResponse()
-      // 因为 token 本身不计算, 所以这里默认以 gpt 3.5 的算做一个伪统计
-      result.data.detail.usage.prompt_tokens = textTokens(prompt, 'gpt-3.5-turbo')
-      result.data.detail.usage.completion_tokens = textTokens(result.data.text, 'gpt-3.5-turbo')
+      // if no usage data, calculate using Tiktoken library
+      result.data.detail.usage.prompt_tokens = textTokens(prompt, model as TiktokenModel)
+      result.data.detail.usage.completion_tokens = textTokens(result.data.text, model as TiktokenModel)
       result.data.detail.usage.total_tokens = result.data.detail.usage.prompt_tokens + result.data.detail.usage.completion_tokens
       result.data.detail.usage.estimated = true
     }
@@ -498,6 +501,7 @@ router.post('/chat-process', [auth, limiter], async (req, res) => {
           roomId,
           message._id,
           result.data.id,
+          model,
           result.data.detail?.usage as UsageResponse)
       }
       // update personal useAmount moved here

service/src/storage/model.ts

@@ -149,16 +149,18 @@ export class ChatUsage {
   roomId: number
   chatId: ObjectId
   messageId: string
+  model: string
   promptTokens: number
   completionTokens: number
   totalTokens: number
   estimated: boolean
   dateTime: number
-  constructor(userId: ObjectId, roomId: number, chatId: ObjectId, messageId: string, usage: UsageResponse) {
+  constructor(userId: ObjectId, roomId: number, chatId: ObjectId, messageId: string, model: string, usage?: UsageResponse) {
     this.userId = userId
     this.roomId = roomId
     this.chatId = chatId
     this.messageId = messageId
+    this.model = model
     if (usage) {
       this.promptTokens = usage.prompt_tokens
       this.completionTokens = usage.completion_tokens

service/src/storage/mongo.ts

@@ -102,8 +102,13 @@ export async function updateChat(chatId: string, response: string, messageId: st
   await chatCol.updateOne(query, update)
 }
 
-export async function insertChatUsage(userId: ObjectId, roomId: number, chatId: ObjectId, messageId: string, usage: UsageResponse) {
-  const chatUsage = new ChatUsage(userId, roomId, chatId, messageId, usage)
+export async function insertChatUsage(userId: ObjectId,
+  roomId: number,
+  chatId: ObjectId,
+  messageId: string,
+  model: string,
+  usage: UsageResponse) {
+  const chatUsage = new ChatUsage(userId, roomId, chatId, messageId, model, usage)
   await usageCol.insertOne(chatUsage)
   return chatUsage
 }

service/src/utils/image.ts

@@ -13,10 +13,10 @@ fs.mkdir('uploads').then(() => {
 
 export async function convertImageUrl(uploadFileKey: string): Promise<string> {
   const imageData = await fs.readFile(`uploads/${uploadFileKey}`)
-  //判断文件格式
+  // 判断文件格式
   const imageType = await fileType.fileTypeFromBuffer(imageData)
-  const mimeType = imageType.mime;
+  const mimeType = imageType.mime
   // 将图片数据转换为 Base64 编码的字符串
   const base64Image = imageData.toString('base64')
   return `data:${mimeType};base64,${base64Image}`
-}
+}