feat: 新增数据看板

This commit is contained in:
CaIon
2024-01-07 18:31:14 +08:00
parent c09df83f34
commit bf8794d257
16 changed files with 455 additions and 6 deletions

View File

@@ -59,9 +59,10 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
if !common.LogConsumeEnabled {
return
}
username := GetUsernameById(userId)
log := &Log{
UserId: userId,
Username: GetUsernameById(userId),
Username: username,
CreatedAt: common.GetTimestamp(),
Type: LogTypeConsume,
Content: content,
@@ -77,6 +78,9 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
if err != nil {
common.LogError(ctx, "failed to record log: "+err.Error())
}
if common.DataExportEnabled {
LogQuotaData(userId, username, modelName, quota, common.GetTimestamp())
}
}
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int) (logs []*Log, err error) {

View File

@@ -127,6 +127,10 @@ func InitDB() (err error) {
if err != nil {
return err
}
err = db.AutoMigrate(&QuotaData{})
if err != nil {
return err
}
common.SysLog("database migrated")
err = createRootAccountIfNeed()
return err

View File

@@ -37,6 +37,8 @@ func InitOptionMap() {
common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled)
common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled)
common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled)
common.OptionMap["DrawingEnabled"] = strconv.FormatBool(common.DrawingEnabled)
common.OptionMap["DataExportEnabled"] = strconv.FormatBool(common.DataExportEnabled)
common.OptionMap["ChannelDisableThreshold"] = strconv.FormatFloat(common.ChannelDisableThreshold, 'f', -1, 64)
common.OptionMap["EmailDomainRestrictionEnabled"] = strconv.FormatBool(common.EmailDomainRestrictionEnabled)
common.OptionMap["EmailDomainWhitelist"] = strings.Join(common.EmailDomainWhitelist, ",")
@@ -76,6 +78,7 @@ func InitOptionMap() {
common.OptionMap["ChatLink"] = common.ChatLink
common.OptionMap["QuotaPerUnit"] = strconv.FormatFloat(common.QuotaPerUnit, 'f', -1, 64)
common.OptionMap["RetryTimes"] = strconv.Itoa(common.RetryTimes)
common.OptionMap["DataExportInterval"] = strconv.Itoa(common.DataExportInterval)
common.OptionMapRWMutex.Unlock()
loadOptionsFromDatabase()
@@ -157,6 +160,12 @@ func updateOptionMap(key string, value string) (err error) {
common.LogConsumeEnabled = boolValue
case "DisplayInCurrencyEnabled":
common.DisplayInCurrencyEnabled = boolValue
case "DisplayTokenStatEnabled":
common.DisplayTokenStatEnabled = boolValue
case "DrawingEnabled":
common.DrawingEnabled = boolValue
case "DataExportEnabled":
common.DataExportEnabled = boolValue
}
}
switch key {
@@ -217,6 +226,8 @@ func updateOptionMap(key string, value string) (err error) {
common.PreConsumedQuota, _ = strconv.Atoi(value)
case "RetryTimes":
common.RetryTimes, _ = strconv.Atoi(value)
case "DataExportInterval":
common.DataExportInterval, _ = strconv.Atoi(value)
case "ModelRatio":
err = common.UpdateModelRatioByJSONString(value)
case "GroupRatio":

87
model/usedata.go Normal file
View File

@@ -0,0 +1,87 @@
package model
import (
"fmt"
"one-api/common"
"time"
)
// QuotaData 柱状图数据
type QuotaData struct {
Id int `json:"id"`
UserID int `json:"user_id" gorm:"index"`
Username string `json:"username" gorm:"index:index_quota_data_model_user_name,priority:2;default:''"`
ModelName string `json:"model_name" gorm:"index;index:index_quota_data_model_user_name,priority:1;default:''"`
CreatedAt int64 `json:"created_at" gorm:"bigint;index:index_quota_data_created_at,priority:2"`
Count int `json:"count" gorm:"default:0"`
Quota int `json:"quota" gorm:"default:0"`
}
func UpdateQuotaData(frequency int) {
for {
common.SysLog("正在更新数据看板数据...")
SaveQuotaDataCache()
time.Sleep(time.Duration(frequency) * time.Minute)
}
}
var CacheQuotaData = make(map[string]*QuotaData)
func LogQuotaDataCache(userId int, username string, modelName string, quota int, createdAt int64) {
// 只精确到小时
createdAt = createdAt - (createdAt % 3600)
key := fmt.Sprintf("%d-%s-%s-%d", userId, username, modelName, createdAt)
quotaData, ok := CacheQuotaData[key]
if ok {
quotaData.Count += 1
quotaData.Quota += quota
} else {
quotaData = &QuotaData{
UserID: userId,
Username: username,
ModelName: modelName,
CreatedAt: createdAt,
Count: 1,
Quota: quota,
}
}
CacheQuotaData[key] = quotaData
}
func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64) {
LogQuotaDataCache(userId, username, modelName, quota, createdAt)
}
func SaveQuotaDataCache() {
// 如果缓存中有数据,就保存到数据库中
// 1. 先查询数据库中是否有数据
// 2. 如果有数据,就更新数据
// 3. 如果没有数据,就插入数据
for _, quotaData := range CacheQuotaData {
quotaDataDB := &QuotaData{}
DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?",
quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt).First(quotaDataDB)
if quotaDataDB.Id > 0 {
quotaDataDB.Count += quotaData.Count
quotaDataDB.Quota += quotaData.Quota
DB.Table("quota_data").Save(quotaDataDB)
} else {
DB.Table("quota_data").Create(quotaData)
}
}
CacheQuotaData = make(map[string]*QuotaData)
}
func GetQuotaDataByUsername(username string) (quotaData []*QuotaData, err error) {
var quotaDatas []*QuotaData
// 从quota_data表中查询数据
err = DB.Table("quota_data").Where("username = ?", username).Find(&quotaDatas).Error
return quotaDatas, err
}
func GetAllQuotaDates() (quotaData []*QuotaData, err error) {
var quotaDatas []*QuotaData
// 从quota_data表中查询数据
err = DB.Table("quota_data").Find(&quotaDatas).Error
return quotaDatas, err
}