feat: 记录请求的token用量

Qianxia666 · Mar 4, 2024 · 912f46f · 912f46f
1 parent 1ad1112
commit 912f46f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git a/model/log.go b/model/log.go
@@ -85,7 +85,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
 	}
 	if common.DataExportEnabled {
 		common.SafeGoroutine(func() {
-			LogQuotaData(userId, username, modelName, quota, common.GetTimestamp())
+			LogQuotaData(userId, username, modelName, quota, common.GetTimestamp(), promptTokens+completionTokens)
 		})
 	}
 }

diff --git a/model/usedata.go b/model/usedata.go
@@ -15,6 +15,7 @@ type QuotaData struct {
 	Username  string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"`
 	ModelName string `json:"model_name" gorm:"index:idx_qdt_model_user_name,priority:1;size:64;default:''"`
 	CreatedAt int64  `json:"created_at" gorm:"bigint;index:idx_qdt_created_at,priority:2"`
+	TokenUsed int    `json:"token_used" gorm:"default:0"`
 	Count     int    `json:"count" gorm:"default:0"`
 	Quota     int    `json:"quota" gorm:"default:0"`
 }
@@ -38,7 +39,7 @@ func UpdateQuotaData() {
 var CacheQuotaData = make(map[string]*QuotaData)
 var CacheQuotaDataLock = sync.Mutex{}
 
-func logQuotaDataCache(userId int, username string, modelName string, quota int, createdAt int64) {
+func logQuotaDataCache(userId int, username string, modelName string, quota int, createdAt int64, tokenUsed int) {
 	key := fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt)
 	quotaData, ok := CacheQuotaData[key]
 	if ok {
@@ -52,18 +53,19 @@ func logQuotaDataCache(userId int, username string, modelName string, quota int,
 			CreatedAt: createdAt,
 			Count:     1,
 			Quota:     quota,
+			TokenUsed: tokenUsed,
 		}
 	}
 	CacheQuotaData[key] = quotaData
 }
 
-func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64) {
+func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64, tokenUsed int) {
 	// 只精确到小时
 	createdAt = createdAt - (createdAt % 3600)
 
 	CacheQuotaDataLock.Lock()
 	defer CacheQuotaDataLock.Unlock()
-	logQuotaDataCache(userId, username, modelName, quota, createdAt)
+	logQuotaDataCache(userId, username, modelName, quota, createdAt, tokenUsed)
 }
 
 func SaveQuotaDataCache() {