xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-17 07:56: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 {
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:''"`
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() {