Skip to content

Commit e44bc53

Browse files
zijiren233zijiren233
authored
refactor: new ranking page (#501)
1 parent 0cb3437 commit e44bc53

24 files changed

Lines changed: 973 additions & 414 deletions

File tree

‎core/controller/consumption_ranking.go‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
package controller
2+
3+
import (
4+
"net/http"
5+
"time"
6+
7+
"github.com/gin-gonic/gin"
8+
"github.com/labring/aiproxy/core/controller/utils"
9+
"github.com/labring/aiproxy/core/middleware"
10+
"github.com/labring/aiproxy/core/model"
11+
)
12+
13+
type ConsumptionRankingResponseItem struct {
14+
Rank int `json:"rank"`
15+
model.ConsumptionRankingItem
16+
}
17+
18+
func normalizeConsumptionRankingType(rankingType string) model.ConsumptionRankingType {
19+
switch model.ConsumptionRankingType(rankingType) {
20+
case model.ConsumptionRankingTypeChannel,
21+
model.ConsumptionRankingTypeModel,
22+
model.ConsumptionRankingTypeGroup:
23+
return model.ConsumptionRankingType(rankingType)
24+
default:
25+
return model.ConsumptionRankingTypeGroup
26+
}
27+
}
28+
29+
// GetConsumptionRanking godoc
30+
//
31+
// @Summary Get consumption ranking
32+
// @Description Returns channel, model, or group consumption ranking aggregated from summary data
33+
// @Tags groups
34+
// @Produce json
35+
// @Security ApiKeyAuth
36+
// @Param type query string false "Ranking type: channel, model, group" default(group)
37+
// @Param start_timestamp query int64 false "Start timestamp"
38+
// @Param end_timestamp query int64 false "End timestamp"
39+
// @Param timezone query string false "Timezone, default is Local"
40+
// @Param page query int false "Page number"
41+
// @Param per_page query int false "Items per page"
42+
// @Param order query string false "Order: used_amount_desc, used_amount_asc, request_count_desc, request_count_asc, total_tokens_desc, total_tokens_asc, channel_id_asc, channel_id_desc, model_asc, model_desc, group_id_asc, group_id_desc"
43+
// @Success 200 {object} middleware.APIResponse{data=map[string]any}
44+
// @Router /api/groups/consumption_ranking [get]
45+
// @Router /api/groups/ranking [get]
46+
func GetConsumptionRanking(c *gin.Context) {
47+
page, perPage := utils.ParsePageParams(c)
48+
49+
timezone := c.DefaultQuery("timezone", "Local")
50+
if _, err := time.LoadLocation(timezone); err != nil {
51+
timezone = "Local"
52+
}
53+
54+
startTime, endTime := utils.ParseTimeRange(c, 0)
55+
56+
rankingType := normalizeConsumptionRankingType(
57+
c.DefaultQuery("type", string(model.ConsumptionRankingTypeGroup)),
58+
)
59+
60+
items, total, _, err := model.GetConsumptionRanking(
61+
rankingType,
62+
startTime,
63+
endTime,
64+
page,
65+
perPage,
66+
c.Query("order"),
67+
)
68+
if err != nil {
69+
middleware.ErrorResponse(c, http.StatusInternalServerError, err.Error())
70+
return
71+
}
72+
73+
page, perPage = model.NormalizePageParams(page, perPage)
74+
75+
responseItems := make([]ConsumptionRankingResponseItem, len(items))
76+
for i, item := range items {
77+
responseItems[i] = ConsumptionRankingResponseItem{
78+
Rank: (page-1)*perPage + i + 1,
79+
ConsumptionRankingItem: item,
80+
}
81+
}
82+
83+
middleware.SuccessResponse(c, gin.H{
84+
"items": responseItems,
85+
"total": total,
86+
"type": rankingType,
87+
"timezone": timezone,
88+
})
89+
}

‎core/controller/group_ranking.go‎

Lines changed: 0 additions & 62 deletions
This file was deleted.

‎core/model/consumption_ranking.go‎

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
package model
2+
3+
import (
4+
"strings"
5+
"time"
6+
7+
"gorm.io/gorm"
8+
)
9+
10+
type ConsumptionRankingType string
11+
12+
const (
13+
ConsumptionRankingTypeGroup ConsumptionRankingType = "group"
14+
ConsumptionRankingTypeChannel ConsumptionRankingType = "channel"
15+
ConsumptionRankingTypeModel ConsumptionRankingType = "model"
16+
)
17+
18+
type ConsumptionRankingItem struct {
19+
GroupID string `json:"group_id,omitempty" gorm:"column:group_id"`
20+
ChannelID int `json:"channel_id,omitempty" gorm:"column:channel_id"`
21+
Model string `json:"model,omitempty" gorm:"column:model"`
22+
RequestCount int64 `json:"request_count" gorm:"column:request_count"`
23+
UsedAmount float64 `json:"used_amount" gorm:"column:used_amount"`
24+
InputTokens int64 `json:"input_tokens" gorm:"column:input_tokens"`
25+
OutputTokens int64 `json:"output_tokens" gorm:"column:output_tokens"`
26+
TotalTokens int64 `json:"total_tokens" gorm:"column:total_tokens"`
27+
}
28+
29+
func normalizeConsumptionRankingOrder(order, dimension string) (normalized, clause string) {
30+
order = strings.ToLower(strings.TrimSpace(order))
31+
dimensionOrderDesc := dimension + "_desc"
32+
dimensionOrderAsc := dimension + "_asc"
33+
34+
switch order {
35+
case "used_amount_asc":
36+
return "used_amount_asc", "used_amount ASC, request_count DESC, " + dimension + " ASC"
37+
case "request_count_desc":
38+
return "request_count_desc", "request_count DESC, used_amount DESC, " + dimension + " ASC"
39+
case "request_count_asc":
40+
return "request_count_asc", "request_count ASC, used_amount DESC, " + dimension + " ASC"
41+
case "total_tokens_desc":
42+
return "total_tokens_desc", "total_tokens DESC, used_amount DESC, " + dimension + " ASC"
43+
case "total_tokens_asc":
44+
return "total_tokens_asc", "total_tokens ASC, used_amount DESC, " + dimension + " ASC"
45+
case dimensionOrderDesc:
46+
return dimensionOrderDesc, dimension + " DESC"
47+
case dimensionOrderAsc:
48+
return dimensionOrderAsc, dimension + " ASC"
49+
default:
50+
return "used_amount_desc", "used_amount DESC, request_count DESC, " + dimension + " ASC"
51+
}
52+
}
53+
54+
//nolint:unparam
55+
func buildConsumptionRankingTimeQuery(
56+
tx *gorm.DB,
57+
timestampColumn string,
58+
start, end time.Time,
59+
) *gorm.DB {
60+
switch {
61+
case !start.IsZero() && !end.IsZero():
62+
tx = tx.Where(timestampColumn+" BETWEEN ? AND ?", start.Unix(), end.Unix())
63+
case !start.IsZero():
64+
tx = tx.Where(timestampColumn+" >= ?", start.Unix())
65+
case !end.IsZero():
66+
tx = tx.Where(timestampColumn+" <= ?", end.Unix())
67+
}
68+
69+
return tx
70+
}
71+
72+
func GetConsumptionRanking(
73+
rankingType ConsumptionRankingType,
74+
start, end time.Time,
75+
page, perPage int,
76+
order string,
77+
) ([]ConsumptionRankingItem, int64, string, error) {
78+
page, perPage = NormalizePageParams(page, perPage)
79+
limit, offset := toLimitOffset(page, perPage)
80+
81+
var (
82+
baseQuery *gorm.DB
83+
groupField string
84+
selectField string
85+
)
86+
87+
switch rankingType {
88+
case ConsumptionRankingTypeChannel:
89+
groupField = "channel_id"
90+
selectField = "channel_id"
91+
baseQuery = buildConsumptionRankingTimeQuery(
92+
LogDB.Model(&Summary{}),
93+
"hour_timestamp",
94+
start,
95+
end,
96+
)
97+
case ConsumptionRankingTypeModel:
98+
groupField = "model"
99+
selectField = "model"
100+
baseQuery = buildConsumptionRankingTimeQuery(
101+
LogDB.Model(&Summary{}),
102+
"hour_timestamp",
103+
start,
104+
end,
105+
)
106+
case ConsumptionRankingTypeGroup:
107+
groupField = "group_id"
108+
selectField = "group_id"
109+
baseQuery = buildConsumptionRankingTimeQuery(
110+
LogDB.Model(&GroupSummary{}),
111+
"hour_timestamp",
112+
start,
113+
end,
114+
)
115+
default:
116+
groupField = "group_id"
117+
selectField = "group_id"
118+
baseQuery = buildConsumptionRankingTimeQuery(
119+
LogDB.Model(&GroupSummary{}),
120+
"hour_timestamp",
121+
start,
122+
end,
123+
)
124+
}
125+
126+
normalizedOrder, orderClause := normalizeConsumptionRankingOrder(order, groupField)
127+
128+
var total int64
129+
if err := baseQuery.Session(&gorm.Session{}).
130+
Distinct(selectField).
131+
Count(&total).
132+
Error; err != nil {
133+
return nil, 0, normalizedOrder, err
134+
}
135+
136+
items := make([]ConsumptionRankingItem, 0, perPage)
137+
if err := baseQuery.
138+
Session(&gorm.Session{}).
139+
Select(
140+
selectField + ", " +
141+
"SUM(request_count) as request_count, " +
142+
"SUM(used_amount) as used_amount, " +
143+
"SUM(input_tokens) as input_tokens, " +
144+
"SUM(output_tokens) as output_tokens, " +
145+
"SUM(total_tokens) as total_tokens",
146+
).
147+
Group(selectField).
148+
Order(orderClause).
149+
Offset(offset).
150+
Limit(limit).
151+
Find(&items).Error; err != nil {
152+
return nil, 0, normalizedOrder, err
153+
}
154+
155+
return items, total, normalizedOrder, nil
156+
}

‎core/model/modelconfig.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -295,7 +295,7 @@ func GetModelConfig(model string) (ModelConfig, error) {
295295
err := DB.Model(&ModelConfig{}).
296296
Where("model = ?", model).
297297
Omit("created_at", "updated_at").
298-
First(config).
298+
First(&config).
299299
Error
300300

301301
return config, HandleNotFound(err, ErrModelConfigNotFound)

0 commit comments

Comments
 (0)