初始化导入 new-api 源码
This commit is contained in:
OpenClaw Task Bot
225
model/usedata.go
Normal file
225
model/usedata.go
Normal file
@@ -0,0 +1,225 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"sort"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
"gorm.io/gorm/clause"
|
||||
)
|
||||
|
||||
// QuotaData 柱状图数据
|
||||
type QuotaData struct {
|
||||
Id int `json:"id"`
|
||||
UserID int `json:"user_id" gorm:"index"`
|
||||
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"`
|
||||
}
|
||||
|
||||
func UpdateQuotaData() {
|
||||
for {
|
||||
if common.DataExportEnabled {
|
||||
common.SysLog("正在更新数据看板数据...")
|
||||
SaveQuotaDataCache()
|
||||
}
|
||||
sleepMinutes := common.DataExportInterval
|
||||
if sleepMinutes <= 0 {
|
||||
common.SysLog(fmt.Sprintf("DataExportInterval=%d 非法，回退为 5 分钟", sleepMinutes))
|
||||
sleepMinutes = 5
|
||||
}
|
||||
time.Sleep(time.Duration(sleepMinutes) * time.Minute)
|
||||
}
|
||||
}
|
||||
|
||||
var CacheQuotaData = make(map[string]*QuotaData)
|
||||
var CacheQuotaDataLock = sync.Mutex{}
|
||||
|
||||
func quotaDataKey(userId int, username string, modelName string, createdAt int64) string {
|
||||
return fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt)
|
||||
}
|
||||
|
||||
func logQuotaDataCache(userId int, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) {
|
||||
key := quotaDataKey(userId, username, modelName, createdAt)
|
||||
quotaData, ok := CacheQuotaData[key]
|
||||
if ok {
|
||||
quotaData.Count += count
|
||||
quotaData.Quota += quota
|
||||
quotaData.TokenUsed += tokenUsed
|
||||
} else {
|
||||
quotaData = &QuotaData{
|
||||
UserID: userId,
|
||||
Username: username,
|
||||
ModelName: modelName,
|
||||
CreatedAt: createdAt,
|
||||
Count: count,
|
||||
Quota: quota,
|
||||
TokenUsed: tokenUsed,
|
||||
}
|
||||
}
|
||||
CacheQuotaData[key] = quotaData
|
||||
}
|
||||
|
||||
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, 1, quota, createdAt, tokenUsed)
|
||||
}
|
||||
|
||||
func SaveQuotaDataCache() {
|
||||
CacheQuotaDataLock.Lock()
|
||||
if len(CacheQuotaData) == 0 {
|
||||
CacheQuotaDataLock.Unlock()
|
||||
return
|
||||
}
|
||||
cacheToFlush := CacheQuotaData
|
||||
CacheQuotaData = make(map[string]*QuotaData)
|
||||
CacheQuotaDataLock.Unlock()
|
||||
|
||||
size := len(cacheToFlush)
|
||||
for _, quotaData := range cacheToFlush {
|
||||
if err := upsertQuotaData(quotaData); err != nil {
|
||||
common.SysLog(fmt.Sprintf("save quota data cache error: %s", err.Error()))
|
||||
CacheQuotaDataLock.Lock()
|
||||
logQuotaDataCache(quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.Count, quotaData.Quota, quotaData.CreatedAt, quotaData.TokenUsed)
|
||||
CacheQuotaDataLock.Unlock()
|
||||
}
|
||||
}
|
||||
common.SysLog(fmt.Sprintf("保存数据看板数据成功，共保存%d条数据", size))
|
||||
}
|
||||
|
||||
func upsertQuotaData(quotaData *QuotaData) error {
|
||||
return DB.Clauses(clause.OnConflict{
|
||||
Columns: []clause.Column{
|
||||
{Name: "user_id"},
|
||||
{Name: "username"},
|
||||
{Name: "model_name"},
|
||||
{Name: "created_at"},
|
||||
},
|
||||
DoUpdates: clause.Assignments(map[string]interface{}{
|
||||
"count": gorm.Expr("quota_data.count + VALUES(count)"),
|
||||
"quota": gorm.Expr("quota_data.quota + VALUES(quota)"),
|
||||
"token_used": gorm.Expr("quota_data.token_used + VALUES(token_used)"),
|
||||
}),
|
||||
}).Table("quota_data").Create(quotaData).Error
|
||||
}
|
||||
|
||||
func increaseQuotaData(userId int, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) {
|
||||
err := upsertQuotaData(&QuotaData{
|
||||
UserID: userId,
|
||||
Username: username,
|
||||
ModelName: modelName,
|
||||
CreatedAt: createdAt,
|
||||
Count: count,
|
||||
Quota: quota,
|
||||
TokenUsed: tokenUsed,
|
||||
})
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("increaseQuotaData error: %s", err))
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeQuotaDataResult(quotaDatas []*QuotaData, err error) ([]*QuotaData, error) {
|
||||
if err != nil {
|
||||
if errors.Is(err, gorm.ErrRecordNotFound) {
|
||||
return []*QuotaData{}, nil
|
||||
}
|
||||
return nil, err
|
||||
}
|
||||
if quotaDatas == nil {
|
||||
return []*QuotaData{}, nil
|
||||
}
|
||||
return quotaDatas, nil
|
||||
}
|
||||
|
||||
func mergeQuotaDataWithCache(quotaDatas []*QuotaData, matcher func(*QuotaData) bool) []*QuotaData {
|
||||
merged := make(map[string]*QuotaData, len(quotaDatas))
|
||||
for _, item := range quotaDatas {
|
||||
if item == nil {
|
||||
continue
|
||||
}
|
||||
copied := *item
|
||||
merged[quotaDataKey(copied.UserID, copied.Username, copied.ModelName, copied.CreatedAt)] = &copied
|
||||
}
|
||||
|
||||
CacheQuotaDataLock.Lock()
|
||||
defer CacheQuotaDataLock.Unlock()
|
||||
for _, cached := range CacheQuotaData {
|
||||
if cached == nil || !matcher(cached) {
|
||||
continue
|
||||
}
|
||||
key := quotaDataKey(cached.UserID, cached.Username, cached.ModelName, cached.CreatedAt)
|
||||
if existing, ok := merged[key]; ok {
|
||||
existing.Count += cached.Count
|
||||
existing.Quota += cached.Quota
|
||||
existing.TokenUsed += cached.TokenUsed
|
||||
} else {
|
||||
copied := *cached
|
||||
merged[key] = &copied
|
||||
}
|
||||
}
|
||||
|
||||
result := make([]*QuotaData, 0, len(merged))
|
||||
for _, item := range merged {
|
||||
result = append(result, item)
|
||||
}
|
||||
sort.Slice(result, func(i, j int) bool {
|
||||
if result[i].CreatedAt == result[j].CreatedAt {
|
||||
if result[i].ModelName == result[j].ModelName {
|
||||
return result[i].UserID < result[j].UserID
|
||||
}
|
||||
return result[i].ModelName < result[j].ModelName
|
||||
}
|
||||
return result[i].CreatedAt < result[j].CreatedAt
|
||||
})
|
||||
return result
|
||||
}
|
||||
|
||||
func GetQuotaDataByUsername(username string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) {
|
||||
var quotaDatas []*QuotaData
|
||||
err = DB.Table("quota_data").Where("username = ? and created_at >= ? and created_at <= ?", username, startTime, endTime).Find("aDatas).Error
|
||||
quotaDatas, err = normalizeQuotaDataResult(quotaDatas, err)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return mergeQuotaDataWithCache(quotaDatas, func(item *QuotaData) bool {
|
||||
return item.Username == username && item.CreatedAt >= startTime && item.CreatedAt <= endTime
|
||||
}), nil
|
||||
}
|
||||
|
||||
func GetQuotaDataByUserId(userId int, startTime int64, endTime int64) (quotaData []*QuotaData, err error) {
|
||||
var quotaDatas []*QuotaData
|
||||
err = DB.Table("quota_data").Where("user_id = ? and created_at >= ? and created_at <= ?", userId, startTime, endTime).Find("aDatas).Error
|
||||
quotaDatas, err = normalizeQuotaDataResult(quotaDatas, err)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return mergeQuotaDataWithCache(quotaDatas, func(item *QuotaData) bool {
|
||||
return item.UserID == userId && item.CreatedAt >= startTime && item.CreatedAt <= endTime
|
||||
}), nil
|
||||
}
|
||||
|
||||
func GetAllQuotaDates(startTime int64, endTime int64, username string) (quotaData []*QuotaData, err error) {
|
||||
if username != "" {
|
||||
return GetQuotaDataByUsername(username, startTime, endTime)
|
||||
}
|
||||
var quotaDatas []*QuotaData
|
||||
err = DB.Table("quota_data").Select("model_name, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used, created_at").Where("created_at >= ? and created_at <= ?", startTime, endTime).Group("model_name, created_at").Find("aDatas).Error
|
||||
quotaDatas, err = normalizeQuotaDataResult(quotaDatas, err)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return mergeQuotaDataWithCache(quotaDatas, func(item *QuotaData) bool {
|
||||
return item.CreatedAt >= startTime && item.CreatedAt <= endTime
|
||||
}), nil
|
||||
}
|
||||
Reference in New Issue
Block a user