feat(dashboard): add traffic flow sankey chart (#5465)
* feat(dashboard): add traffic flow sankey chart Add dashboard flow APIs and a Sankey-based flow view with user, optional API key, model, and channel layers.\n\nReuse the dashboard VChart palette, add precise link/node tooltips and interactions, and cover filtering, layer ordering, color stability, and error states with tests. * feat: build flow chart from quota data --------- Co-authored-by: CaIon <i@caion.me>
This commit is contained in:
Quaternijkon
+31
-3
@@ -298,6 +298,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
|
||||
username := c.GetString("username")
|
||||
requestId := c.GetString(common.RequestIdKey)
|
||||
upstreamRequestId := c.GetString(common.UpstreamRequestIdKey)
|
||||
createdAt := common.GetTimestamp()
|
||||
otherStr := common.MapToJsonStr(params.Other)
|
||||
// 判断是否需要记录 IP
|
||||
needRecordIp := false
|
||||
@@ -309,7 +310,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
|
||||
log := &Log{
|
||||
UserId: userId,
|
||||
Username: username,
|
||||
CreatedAt: common.GetTimestamp(),
|
||||
CreatedAt: createdAt,
|
||||
Type: LogTypeConsume,
|
||||
Content: params.Content,
|
||||
PromptTokens: params.PromptTokens,
|
||||
@@ -338,7 +339,18 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
|
||||
}
|
||||
if common.DataExportEnabled {
|
||||
gopool.Go(func() {
|
||||
LogQuotaData(userId, username, params.ModelName, params.Quota, common.GetTimestamp(), params.PromptTokens+params.CompletionTokens)
|
||||
LogQuotaData(QuotaDataLogParams{
|
||||
UserID: userId,
|
||||
Username: username,
|
||||
ModelName: params.ModelName,
|
||||
Quota: params.Quota,
|
||||
CreatedAt: createdAt,
|
||||
TokenUsed: params.PromptTokens + params.CompletionTokens,
|
||||
UseGroup: params.Group,
|
||||
TokenID: params.TokenId,
|
||||
ChannelID: params.ChannelId,
|
||||
NodeName: common.NodeName,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -366,10 +378,11 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
|
||||
tokenName = token.Name
|
||||
}
|
||||
}
|
||||
createdAt := common.GetTimestamp()
|
||||
log := &Log{
|
||||
UserId: params.UserId,
|
||||
Username: username,
|
||||
CreatedAt: common.GetTimestamp(),
|
||||
CreatedAt: createdAt,
|
||||
Type: params.LogType,
|
||||
Content: params.Content,
|
||||
TokenName: tokenName,
|
||||
@@ -384,6 +397,21 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
|
||||
if err != nil {
|
||||
common.SysLog("failed to record task billing log: " + err.Error())
|
||||
}
|
||||
if params.LogType == LogTypeConsume && common.DataExportEnabled {
|
||||
gopool.Go(func() {
|
||||
LogQuotaData(QuotaDataLogParams{
|
||||
UserID: params.UserId,
|
||||
Username: username,
|
||||
ModelName: params.ModelName,
|
||||
Quota: params.Quota,
|
||||
CreatedAt: createdAt,
|
||||
UseGroup: params.Group,
|
||||
TokenID: params.TokenId,
|
||||
ChannelID: params.ChannelId,
|
||||
NodeName: common.NodeName,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
|
||||
|
||||
@@ -40,6 +40,7 @@ func TestMain(m *testing.M) {
|
||||
&Token{},
|
||||
&Log{},
|
||||
&Channel{},
|
||||
&QuotaData{},
|
||||
&Ability{},
|
||||
&TopUp{},
|
||||
&SubscriptionPlan{},
|
||||
@@ -62,6 +63,7 @@ func truncateTables(t *testing.T) {
|
||||
DB.Exec("DELETE FROM tokens")
|
||||
DB.Exec("DELETE FROM logs")
|
||||
DB.Exec("DELETE FROM channels")
|
||||
DB.Exec("DELETE FROM quota_data")
|
||||
DB.Exec("DELETE FROM abilities")
|
||||
DB.Exec("DELETE FROM top_ups")
|
||||
DB.Exec("DELETE FROM subscription_orders")
|
||||
|
||||
+76
-31
@@ -16,11 +16,28 @@ 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"`
|
||||
UseGroup string `json:"use_group" gorm:"index;size:64;default:''"`
|
||||
TokenID int `json:"token_id" gorm:"index;default:0"`
|
||||
ChannelID int `json:"channel_id" gorm:"index;default:0"`
|
||||
NodeName string `json:"node_name" gorm:"index;size:64;default:''"`
|
||||
TokenUsed int `json:"token_used" gorm:"default:0"`
|
||||
Count int `json:"count" gorm:"default:0"`
|
||||
Quota int `json:"quota" gorm:"default:0"`
|
||||
}
|
||||
|
||||
type QuotaDataLogParams struct {
|
||||
UserID int
|
||||
Username string
|
||||
ModelName string
|
||||
Quota int
|
||||
CreatedAt int64
|
||||
TokenUsed int
|
||||
UseGroup string
|
||||
TokenID int
|
||||
ChannelID int
|
||||
NodeName string
|
||||
}
|
||||
|
||||
func UpdateQuotaData() {
|
||||
for {
|
||||
if common.DataExportEnabled {
|
||||
@@ -34,34 +51,50 @@ func UpdateQuotaData() {
|
||||
var CacheQuotaData = make(map[string]*QuotaData)
|
||||
var CacheQuotaDataLock = sync.Mutex{}
|
||||
|
||||
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]
|
||||
func logQuotaDataCache(quotaData *QuotaData) {
|
||||
key := fmt.Sprintf("%d\x00%s\x00%s\x00%d\x00%s\x00%d\x00%d\x00%s",
|
||||
quotaData.UserID,
|
||||
quotaData.Username,
|
||||
quotaData.ModelName,
|
||||
quotaData.CreatedAt,
|
||||
quotaData.UseGroup,
|
||||
quotaData.TokenID,
|
||||
quotaData.ChannelID,
|
||||
quotaData.NodeName,
|
||||
)
|
||||
count := quotaData.Count
|
||||
quota := quotaData.Quota
|
||||
tokenUsed := quotaData.TokenUsed
|
||||
cachedQuotaData, ok := CacheQuotaData[key]
|
||||
if ok {
|
||||
quotaData.Count += 1
|
||||
quotaData.Quota += quota
|
||||
quotaData.TokenUsed += tokenUsed
|
||||
} else {
|
||||
quotaData = &QuotaData{
|
||||
UserID: userId,
|
||||
Username: username,
|
||||
ModelName: modelName,
|
||||
CreatedAt: createdAt,
|
||||
Count: 1,
|
||||
Quota: quota,
|
||||
TokenUsed: tokenUsed,
|
||||
}
|
||||
cachedQuotaData.Count += count
|
||||
cachedQuotaData.Quota += quota
|
||||
cachedQuotaData.TokenUsed += tokenUsed
|
||||
quotaData = cachedQuotaData
|
||||
}
|
||||
CacheQuotaData[key] = quotaData
|
||||
}
|
||||
|
||||
func LogQuotaData(userId int, username string, modelName string, quota int, createdAt int64, tokenUsed int) {
|
||||
func LogQuotaData(params QuotaDataLogParams) {
|
||||
// 只精确到小时
|
||||
createdAt = createdAt - (createdAt % 3600)
|
||||
createdAt := params.CreatedAt - (params.CreatedAt % 3600)
|
||||
quotaData := &QuotaData{
|
||||
UserID: params.UserID,
|
||||
Username: params.Username,
|
||||
ModelName: params.ModelName,
|
||||
CreatedAt: createdAt,
|
||||
UseGroup: params.UseGroup,
|
||||
TokenID: params.TokenID,
|
||||
ChannelID: params.ChannelID,
|
||||
NodeName: params.NodeName,
|
||||
Count: 1,
|
||||
Quota: params.Quota,
|
||||
TokenUsed: params.TokenUsed,
|
||||
}
|
||||
|
||||
CacheQuotaDataLock.Lock()
|
||||
defer CacheQuotaDataLock.Unlock()
|
||||
logQuotaDataCache(userId, username, modelName, quota, createdAt, tokenUsed)
|
||||
logQuotaDataCache(quotaData)
|
||||
}
|
||||
|
||||
func SaveQuotaDataCache() {
|
||||
@@ -74,13 +107,15 @@ func SaveQuotaDataCache() {
|
||||
// 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)
|
||||
DB.Table("quota_data").
|
||||
Where("user_id = ? and username = ? and model_name = ? and created_at = ? and use_group = ? and token_id = ? and channel_id = ? and node_name = ?",
|
||||
quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt, quotaData.UseGroup, quotaData.TokenID, quotaData.ChannelID, quotaData.NodeName).
|
||||
First(quotaDataDB)
|
||||
if quotaDataDB.Id > 0 {
|
||||
//quotaDataDB.Count += quotaData.Count
|
||||
//quotaDataDB.Quota += quotaData.Quota
|
||||
//DB.Table("quota_data").Save(quotaDataDB)
|
||||
increaseQuotaData(quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.Count, quotaData.Quota, quotaData.CreatedAt, quotaData.TokenUsed)
|
||||
increaseQuotaData(quotaData)
|
||||
} else {
|
||||
DB.Table("quota_data").Create(quotaData)
|
||||
}
|
||||
@@ -89,13 +124,15 @@ func SaveQuotaDataCache() {
|
||||
common.SysLog(fmt.Sprintf("保存数据看板数据成功，共保存%d条数据", size))
|
||||
}
|
||||
|
||||
func increaseQuotaData(userId int, username string, modelName string, count int, quota int, createdAt int64, tokenUsed int) {
|
||||
err := DB.Table("quota_data").Where("user_id = ? and username = ? and model_name = ? and created_at = ?",
|
||||
userId, username, modelName, createdAt).Updates(map[string]interface{}{
|
||||
"count": gorm.Expr("count + ?", count),
|
||||
"quota": gorm.Expr("quota + ?", quota),
|
||||
"token_used": gorm.Expr("token_used + ?", tokenUsed),
|
||||
}).Error
|
||||
func increaseQuotaData(quotaData *QuotaData) {
|
||||
err := DB.Table("quota_data").
|
||||
Where("user_id = ? and username = ? and model_name = ? and created_at = ? and use_group = ? and token_id = ? and channel_id = ? and node_name = ?",
|
||||
quotaData.UserID, quotaData.Username, quotaData.ModelName, quotaData.CreatedAt, quotaData.UseGroup, quotaData.TokenID, quotaData.ChannelID, quotaData.NodeName).
|
||||
Updates(map[string]interface{}{
|
||||
"count": gorm.Expr("count + ?", quotaData.Count),
|
||||
"quota": gorm.Expr("quota + ?", quotaData.Quota),
|
||||
"token_used": gorm.Expr("token_used + ?", quotaData.TokenUsed),
|
||||
}).Error
|
||||
if err != nil {
|
||||
common.SysLog(fmt.Sprintf("increaseQuotaData error: %s", err))
|
||||
}
|
||||
@@ -104,14 +141,22 @@ func increaseQuotaData(userId int, username string, modelName string, count int,
|
||||
func GetQuotaDataByUsername(username string, startTime int64, endTime int64) (quotaData []*QuotaData, err error) {
|
||||
var quotaDatas []*QuotaData
|
||||
// 从quota_data表中查询数据
|
||||
err = DB.Table("quota_data").Where("username = ? and created_at >= ? and created_at <= ?", username, startTime, endTime).Find("aDatas).Error
|
||||
err = DB.Table("quota_data").
|
||||
Select("user_id, username, model_name, created_at, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used").
|
||||
Where("username = ? and created_at >= ? and created_at <= ?", username, startTime, endTime).
|
||||
Group("user_id, username, model_name, created_at").
|
||||
Find("aDatas).Error
|
||||
return quotaDatas, err
|
||||
}
|
||||
|
||||
func GetQuotaDataByUserId(userId int, startTime int64, endTime int64) (quotaData []*QuotaData, err error) {
|
||||
var quotaDatas []*QuotaData
|
||||
// 从quota_data表中查询数据
|
||||
err = DB.Table("quota_data").Where("user_id = ? and created_at >= ? and created_at <= ?", userId, startTime, endTime).Find("aDatas).Error
|
||||
err = DB.Table("quota_data").
|
||||
Select("user_id, username, model_name, created_at, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used").
|
||||
Where("user_id = ? and created_at >= ? and created_at <= ?", userId, startTime, endTime).
|
||||
Group("user_id, username, model_name, created_at").
|
||||
Find("aDatas).Error
|
||||
return quotaDatas, err
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
type FlowQuotaData struct {
|
||||
UserID int `json:"user_id,omitempty" gorm:"column:user_id"`
|
||||
Username string `json:"username,omitempty" gorm:"column:username"`
|
||||
NodeName string `json:"node_name,omitempty" gorm:"column:node_name"`
|
||||
TokenID int `json:"token_id,omitempty" gorm:"column:token_id"`
|
||||
TokenName string `json:"token_name,omitempty" gorm:"-"`
|
||||
UseGroup string `json:"use_group" gorm:"column:use_group"`
|
||||
ChannelID int `json:"channel_id,omitempty" gorm:"column:channel_id"`
|
||||
ChannelName string `json:"channel_name,omitempty" gorm:"-"`
|
||||
ModelName string `json:"model_name" gorm:"column:model_name"`
|
||||
TokenUsed int `json:"token_used" gorm:"column:token_used"`
|
||||
Count int `json:"count" gorm:"column:count"`
|
||||
Quota int `json:"quota" gorm:"column:quota"`
|
||||
}
|
||||
|
||||
func GetFlowQuotaData(startTime int64, endTime int64, username string, userID int, role int) ([]*FlowQuotaData, error) {
|
||||
switch {
|
||||
case role >= common.RoleRootUser:
|
||||
return getRootFlowQuotaData(startTime, endTime, username)
|
||||
case role >= common.RoleAdminUser:
|
||||
return getAdminFlowQuotaData(startTime, endTime, username)
|
||||
default:
|
||||
return getSelfFlowQuotaData(startTime, endTime, userID)
|
||||
}
|
||||
}
|
||||
|
||||
func flowQuotaBaseQuery(startTime int64, endTime int64) *gorm.DB {
|
||||
query := DB.Table("quota_data").
|
||||
Where("use_group <> ''").
|
||||
Where("created_at >= ? and created_at <= ?", startTime, endTime)
|
||||
return query
|
||||
}
|
||||
|
||||
func getSelfFlowQuotaData(startTime int64, endTime int64, userID int) ([]*FlowQuotaData, error) {
|
||||
rows := make([]*FlowQuotaData, 0)
|
||||
err := flowQuotaBaseQuery(startTime, endTime).
|
||||
Select("token_id, use_group, model_name, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used").
|
||||
Where("user_id = ?", userID).
|
||||
Group("token_id, use_group, model_name").
|
||||
Order("quota DESC").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rows, fillFlowTokenNames(rows)
|
||||
}
|
||||
|
||||
func getAdminFlowQuotaData(startTime int64, endTime int64, username string) ([]*FlowQuotaData, error) {
|
||||
rows := make([]*FlowQuotaData, 0)
|
||||
query := flowQuotaBaseQuery(startTime, endTime).
|
||||
Select("user_id, username, use_group, model_name, channel_id, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used")
|
||||
if username != "" {
|
||||
query = query.Where("username = ?", username)
|
||||
}
|
||||
err := query.
|
||||
Group("user_id, username, use_group, model_name, channel_id").
|
||||
Order("quota DESC").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return rows, fillFlowChannelNames(rows)
|
||||
}
|
||||
|
||||
func getRootFlowQuotaData(startTime int64, endTime int64, username string) ([]*FlowQuotaData, error) {
|
||||
rows := make([]*FlowQuotaData, 0)
|
||||
query := flowQuotaBaseQuery(startTime, endTime).
|
||||
Select("user_id, username, node_name, token_id, use_group, model_name, channel_id, sum(count) as count, sum(quota) as quota, sum(token_used) as token_used")
|
||||
if username != "" {
|
||||
query = query.Where("username = ?", username)
|
||||
}
|
||||
err := query.
|
||||
Group("user_id, username, node_name, token_id, use_group, model_name, channel_id").
|
||||
Order("quota DESC").
|
||||
Find(&rows).Error
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := fillFlowTokenNames(rows); err != nil {
|
||||
return rows, err
|
||||
}
|
||||
return rows, fillFlowChannelNames(rows)
|
||||
}
|
||||
|
||||
func fillFlowTokenNames(rows []*FlowQuotaData) error {
|
||||
tokenIDSet := make(map[int]struct{})
|
||||
tokenIDs := make([]int, 0)
|
||||
for _, row := range rows {
|
||||
if row.TokenID == 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := tokenIDSet[row.TokenID]; ok {
|
||||
continue
|
||||
}
|
||||
tokenIDSet[row.TokenID] = struct{}{}
|
||||
tokenIDs = append(tokenIDs, row.TokenID)
|
||||
}
|
||||
if len(tokenIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var tokens []struct {
|
||||
Id int `gorm:"column:id"`
|
||||
Name string `gorm:"column:name"`
|
||||
}
|
||||
if err := DB.Model(&Token{}).Select("id, name").Where("id IN ?", tokenIDs).Find(&tokens).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
tokenNameByID := make(map[int]string, len(tokens))
|
||||
for _, token := range tokens {
|
||||
tokenNameByID[token.Id] = token.Name
|
||||
}
|
||||
// Deleted tokens are intentionally not resolved here: leave TokenName empty
|
||||
// so the frontend can render a localized "deleted (id)" label instead.
|
||||
for _, row := range rows {
|
||||
if name := tokenNameByID[row.TokenID]; name != "" {
|
||||
row.TokenName = name
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func fillFlowChannelNames(rows []*FlowQuotaData) error {
|
||||
channelIDSet := make(map[int]struct{})
|
||||
channelIDs := make([]int, 0)
|
||||
for _, row := range rows {
|
||||
if row.ChannelID == 0 {
|
||||
continue
|
||||
}
|
||||
if _, ok := channelIDSet[row.ChannelID]; ok {
|
||||
continue
|
||||
}
|
||||
channelIDSet[row.ChannelID] = struct{}{}
|
||||
channelIDs = append(channelIDs, row.ChannelID)
|
||||
}
|
||||
if len(channelIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
channelNameByID := make(map[int]string, len(channelIDs))
|
||||
if common.MemoryCacheEnabled {
|
||||
for _, channelID := range channelIDs {
|
||||
if channel, err := CacheGetChannel(channelID); err == nil {
|
||||
channelNameByID[channelID] = channel.Name
|
||||
}
|
||||
}
|
||||
} else {
|
||||
var channels []struct {
|
||||
Id int `gorm:"column:id"`
|
||||
Name string `gorm:"column:name"`
|
||||
}
|
||||
if err := DB.Table("channels").Select("id, name").Where("id IN ?", channelIDs).Find(&channels).Error; err != nil {
|
||||
return err
|
||||
}
|
||||
for _, channel := range channels {
|
||||
channelNameByID[channel.Id] = channel.Name
|
||||
}
|
||||
}
|
||||
for _, row := range rows {
|
||||
if name := channelNameByID[row.ChannelID]; name != "" {
|
||||
row.ChannelName = name
|
||||
continue
|
||||
}
|
||||
if row.ChannelID > 0 {
|
||||
row.ChannelName = fmt.Sprintf("channel-%d", row.ChannelID)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,193 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/QuantumNous/new-api/common"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func seedFlowQuotaData(t *testing.T, quotaData QuotaData) {
|
||||
t.Helper()
|
||||
require.NoError(t, DB.Create("aData).Error)
|
||||
}
|
||||
|
||||
func seedFlowLookupData(t *testing.T) {
|
||||
t.Helper()
|
||||
require.NoError(t, DB.Create(&Channel{Id: 1, Name: "east"}).Error)
|
||||
require.NoError(t, DB.Create(&Channel{Id: 2, Name: "west"}).Error)
|
||||
require.NoError(t, DB.Create(&Token{Id: 11, UserId: 1, Key: "sk-primary", Name: "primary"}).Error)
|
||||
require.NoError(t, DB.Create(&Token{Id: 22, UserId: 2, Key: "sk-backup", Name: "backup"}).Error)
|
||||
require.NoError(t, DB.Delete(&Token{Id: 11}).Error)
|
||||
}
|
||||
|
||||
func TestGetFlowQuotaDataUsesQuotaDataRoleSpecificDimensions(t *testing.T) {
|
||||
truncateTables(t)
|
||||
seedFlowLookupData(t)
|
||||
|
||||
seedFlowQuotaData(t, QuotaData{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
NodeName: "node-a",
|
||||
TokenID: 11,
|
||||
UseGroup: "vip",
|
||||
ModelName: "gpt-a",
|
||||
ChannelID: 1,
|
||||
CreatedAt: 1000,
|
||||
Count: 2,
|
||||
Quota: 100,
|
||||
TokenUsed: 40,
|
||||
})
|
||||
seedFlowQuotaData(t, QuotaData{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
NodeName: "node-a",
|
||||
TokenID: 11,
|
||||
UseGroup: "vip",
|
||||
ModelName: "gpt-a",
|
||||
ChannelID: 1,
|
||||
CreatedAt: 1100,
|
||||
Count: 1,
|
||||
Quota: 50,
|
||||
TokenUsed: 20,
|
||||
})
|
||||
seedFlowQuotaData(t, QuotaData{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
NodeName: "node-a",
|
||||
TokenID: 11,
|
||||
UseGroup: "vip",
|
||||
ModelName: "gpt-a",
|
||||
ChannelID: 2,
|
||||
CreatedAt: 1200,
|
||||
Count: 1,
|
||||
Quota: 25,
|
||||
TokenUsed: 10,
|
||||
})
|
||||
seedFlowQuotaData(t, QuotaData{
|
||||
UserID: 2,
|
||||
Username: "bob",
|
||||
NodeName: "node-b",
|
||||
TokenID: 22,
|
||||
UseGroup: "default",
|
||||
ModelName: "gpt-b",
|
||||
ChannelID: 1,
|
||||
CreatedAt: 1300,
|
||||
Count: 3,
|
||||
Quota: 70,
|
||||
TokenUsed: 30,
|
||||
})
|
||||
seedFlowQuotaData(t, QuotaData{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
ModelName: "legacy",
|
||||
CreatedAt: 1400,
|
||||
Count: 99,
|
||||
Quota: 999,
|
||||
TokenUsed: 999,
|
||||
})
|
||||
|
||||
rootRows, err := GetFlowQuotaData(900, 2000, "", 0, common.RoleRootUser)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, rootRows, 3)
|
||||
// Token 11 was soft-deleted, so its name is intentionally left empty for the
|
||||
// frontend to render a localized "deleted (id)" label instead.
|
||||
require.Equal(t, FlowQuotaData{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
NodeName: "node-a",
|
||||
TokenID: 11,
|
||||
TokenName: "",
|
||||
UseGroup: "vip",
|
||||
ChannelID: 1,
|
||||
ChannelName: "east",
|
||||
ModelName: "gpt-a",
|
||||
TokenUsed: 60,
|
||||
Count: 3,
|
||||
Quota: 150,
|
||||
}, *rootRows[0])
|
||||
// A token that still exists resolves to its current name.
|
||||
require.Equal(t, 22, rootRows[1].TokenID)
|
||||
require.Equal(t, "backup", rootRows[1].TokenName)
|
||||
|
||||
adminRows, err := GetFlowQuotaData(900, 2000, "alice", 0, common.RoleAdminUser)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, adminRows, 2)
|
||||
require.Equal(t, 0, adminRows[0].TokenID)
|
||||
require.Empty(t, adminRows[0].TokenName)
|
||||
require.Empty(t, adminRows[0].NodeName)
|
||||
require.Equal(t, "alice", adminRows[0].Username)
|
||||
require.Equal(t, "vip", adminRows[0].UseGroup)
|
||||
require.Equal(t, "east", adminRows[0].ChannelName)
|
||||
require.Equal(t, 150, adminRows[0].Quota)
|
||||
|
||||
selfRows, err := GetFlowQuotaData(900, 2000, "", 1, common.RoleCommonUser)
|
||||
require.NoError(t, err)
|
||||
require.Len(t, selfRows, 1)
|
||||
require.Empty(t, selfRows[0].Username)
|
||||
require.Equal(t, 0, selfRows[0].ChannelID)
|
||||
require.Empty(t, selfRows[0].ChannelName)
|
||||
require.Empty(t, selfRows[0].TokenName)
|
||||
require.Equal(t, "vip", selfRows[0].UseGroup)
|
||||
require.Equal(t, 175, selfRows[0].Quota)
|
||||
}
|
||||
|
||||
func TestLogQuotaDataSplitsRowsByUseGroupTokenChannelAndNode(t *testing.T) {
|
||||
truncateTables(t)
|
||||
CacheQuotaDataLock.Lock()
|
||||
CacheQuotaData = make(map[string]*QuotaData)
|
||||
CacheQuotaDataLock.Unlock()
|
||||
|
||||
LogQuotaData(QuotaDataLogParams{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
ModelName: "gpt-a",
|
||||
CreatedAt: 3661,
|
||||
UseGroup: "vip",
|
||||
TokenID: 11,
|
||||
ChannelID: 1,
|
||||
NodeName: "node-a",
|
||||
Quota: 100,
|
||||
TokenUsed: 40,
|
||||
})
|
||||
LogQuotaData(QuotaDataLogParams{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
ModelName: "gpt-a",
|
||||
CreatedAt: 3700,
|
||||
UseGroup: "vip",
|
||||
TokenID: 11,
|
||||
ChannelID: 1,
|
||||
NodeName: "node-a",
|
||||
Quota: 50,
|
||||
TokenUsed: 20,
|
||||
})
|
||||
LogQuotaData(QuotaDataLogParams{
|
||||
UserID: 1,
|
||||
Username: "alice",
|
||||
ModelName: "gpt-a",
|
||||
CreatedAt: 3700,
|
||||
UseGroup: "default",
|
||||
TokenID: 11,
|
||||
ChannelID: 1,
|
||||
NodeName: "node-a",
|
||||
Quota: 25,
|
||||
TokenUsed: 10,
|
||||
})
|
||||
|
||||
SaveQuotaDataCache()
|
||||
|
||||
var rows []QuotaData
|
||||
require.NoError(t, DB.Order("quota DESC").Find(&rows).Error)
|
||||
require.Len(t, rows, 2)
|
||||
require.Equal(t, int64(3600), rows[0].CreatedAt)
|
||||
require.Equal(t, "vip", rows[0].UseGroup)
|
||||
require.Equal(t, 11, rows[0].TokenID)
|
||||
require.Equal(t, 1, rows[0].ChannelID)
|
||||
require.Equal(t, "node-a", rows[0].NodeName)
|
||||
require.Equal(t, 2, rows[0].Count)
|
||||
require.Equal(t, 150, rows[0].Quota)
|
||||
require.Equal(t, 60, rows[0].TokenUsed)
|
||||
require.Equal(t, "default", rows[1].UseGroup)
|
||||
require.Equal(t, 25, rows[1].Quota)
|
||||
}
|
||||
Reference in New Issue
Block a user