diff --git a/model/log.go b/model/log.go index 44bc5b7..2740c5a 100644 --- 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 index 9ca623d..b2f3025 100644 --- 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() {