xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-17 16:06:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: 记录请求的token用量

Browse Source
This commit is contained in:
1808837298@qq.com 2024-03-04 14:58:30 +08:00
parent 1ad1112351
commit 912f46fcd2
2 changed files with 6 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
model/log.go
View File

@ -85,7 +85,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
} }
if common.DataExportEnabled { if common.DataExportEnabled {
common.SafeGoroutine(func() { common.SafeGoroutine(func() {
LogQuotaData(userId, username, modelName, quota, common.GetTimestamp()) LogQuotaData(userId, username, modelName, quota, common.GetTimestamp(), promptTokens+completionTokens)
}) })
} }
} }

8
model/usedata.go
View File

@ -15,6 +15,7 @@ type QuotaData struct {
Username string `json:"username" gorm:"index:idx_qdt_model_user_name,priority:2;size:64;default:''"` 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:''"` 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"` 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"` Count int `json:"count" gorm:"default:0"`
Quota int `json:"quota" gorm:"default:0"` Quota int `json:"quota" gorm:"default:0"`
} }
@ -38,7 +39,7 @@ func UpdateQuotaData() {
var CacheQuotaData = make(map[string]*QuotaData) var CacheQuotaData = make(map[string]*QuotaData)
var CacheQuotaDataLock = sync.Mutex{} 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) key := fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt)
quotaData, ok := CacheQuotaData[key] quotaData, ok := CacheQuotaData[key]
if ok { if ok {
@ -52,18 +53,19 @@ func logQuotaDataCache(userId int, username string, modelName string, quota int,
CreatedAt: createdAt, CreatedAt: createdAt,
Count: 1, Count: 1,
Quota: quota, Quota: quota,
TokenUsed: tokenUsed,
} }
} }
CacheQuotaData[key] = quotaData 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) createdAt = createdAt - (createdAt % 3600)
CacheQuotaDataLock.Lock() CacheQuotaDataLock.Lock()
defer CacheQuotaDataLock.Unlock() defer CacheQuotaDataLock.Unlock()
logQuotaDataCache(userId, username, modelName, quota, createdAt) logQuotaDataCache(userId, username, modelName, quota, createdAt, tokenUsed)
} }
func SaveQuotaDataCache() { func SaveQuotaDataCache() {