Skip to content

Commit 7da36ae

Browse files
CSCITechCSCITech
authored
Merge pull request #1 from MAX-API-Next/cscitech
Cscitech
2 parents 8519b9b + 2db0d1e commit 7da36ae

15 files changed

Lines changed: 416 additions & 68 deletions

File tree

‎controller/log.go‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ import (
1515
func GetAllLogs(c *gin.Context) {
1616
pageInfo := common.GetPageQuery(c)
1717
logType, _ := strconv.Atoi(c.Query("type"))
18+
logFilter := c.Query("log_filter")
1819
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
1920
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
2021
username := c.Query("username")
@@ -24,7 +25,7 @@ func GetAllLogs(c *gin.Context) {
2425
group := c.Query("group")
2526
requestId := c.Query("request_id")
2627
upstreamRequestId := c.Query("upstream_request_id")
27-
logs, total, err := model.GetAllLogs(logType, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId)
28+
logs, total, err := model.GetAllLogs(logType, logFilter, startTimestamp, endTimestamp, modelName, username, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), channel, group, requestId, upstreamRequestId)
2829
if err != nil {
2930
common.ApiError(c, err)
3031
return
@@ -68,14 +69,15 @@ func GetUserLogs(c *gin.Context) {
6869
pageInfo := common.GetPageQuery(c)
6970
userId := c.GetInt("id")
7071
logType, _ := strconv.Atoi(c.Query("type"))
72+
logFilter := c.Query("log_filter")
7173
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
7274
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
7375
tokenName := c.Query("token_name")
7476
modelName := c.Query("model_name")
7577
group := c.Query("group")
7678
requestId := c.Query("request_id")
7779
upstreamRequestId := c.Query("upstream_request_id")
78-
logs, total, err := model.GetUserLogs(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId)
80+
logs, total, err := model.GetUserLogs(userId, logType, logFilter, startTimestamp, endTimestamp, modelName, tokenName, pageInfo.GetStartIdx(), pageInfo.GetPageSize(), group, requestId, upstreamRequestId)
7981
if err != nil {
8082
common.ApiError(c, err)
8183
return
@@ -155,14 +157,15 @@ func GetLogByKey(c *gin.Context) {
155157

156158
func GetLogsStat(c *gin.Context) {
157159
logType, _ := strconv.Atoi(c.Query("type"))
160+
logFilter := c.Query("log_filter")
158161
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
159162
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
160163
tokenName := c.Query("token_name")
161164
username := c.Query("username")
162165
modelName := c.Query("model_name")
163166
channel, _ := strconv.Atoi(c.Query("channel"))
164167
group := c.Query("group")
165-
stat, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
168+
stat, err := model.SumUsedQuota(logType, logFilter, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
166169
if err != nil {
167170
common.ApiError(c, err)
168171
return
@@ -183,13 +186,14 @@ func GetLogsStat(c *gin.Context) {
183186
func GetLogsSelfStat(c *gin.Context) {
184187
username := c.GetString("username")
185188
logType, _ := strconv.Atoi(c.Query("type"))
189+
logFilter := c.Query("log_filter")
186190
startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64)
187191
endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64)
188192
tokenName := c.Query("token_name")
189193
modelName := c.Query("model_name")
190194
channel, _ := strconv.Atoi(c.Query("channel"))
191195
group := c.Query("group")
192-
quotaNum, err := model.SumUsedQuota(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
196+
quotaNum, err := model.SumUsedQuota(logType, logFilter, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group)
193197
if err != nil {
194198
common.ApiError(c, err)
195199
return

‎controller/relay.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -387,7 +387,7 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t
387387
other["channel_name"] = c.GetString("channel_name")
388388
other["channel_type"] = c.GetInt("channel_type")
389389
adminInfo := make(map[string]interface{})
390-
adminInfo["use_channel"] = c.GetStringSlice("use_channel")
390+
service.AppendRetryLogInfo(c, adminInfo, other)
391391
isMultiKey := common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey)
392392
if isMultiKey {
393393
adminInfo["is_multi_key"] = true

‎model/log.go‎

Lines changed: 34 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,32 @@ const (
6868
LogTypeLogin = 7
6969
)
7070

71+
const LogFilterRetry = "retry"
72+
73+
func applyLogTypeFilter(tx *gorm.DB, logType int) *gorm.DB {
74+
if logType == LogTypeUnknown {
75+
return tx
76+
}
77+
return tx.Where("logs.type = ?", logType)
78+
}
79+
80+
func applyRetryLogFilter(tx *gorm.DB) *gorm.DB {
81+
retryPatterns := []string{
82+
`%"retry_log":true%`,
83+
`%"empty_retry":true%`,
84+
}
85+
return tx.Where("(logs.other LIKE ? OR logs.other LIKE ?)", retryPatterns[0], retryPatterns[1])
86+
}
87+
88+
func applyLogFilter(tx *gorm.DB, filter string) *gorm.DB {
89+
switch filter {
90+
case LogFilterRetry:
91+
return applyRetryLogFilter(tx)
92+
default:
93+
return tx
94+
}
95+
}
96+
7197
func userVisibleAuditInfo(adminInfoValue interface{}) map[string]interface{} {
7298
if !common.LogRequestContentEnabled && !common.LogResponseContentEnabled {
7399
return nil
@@ -464,13 +490,8 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) {
464490
}
465491
}
466492

467-
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) {
468-
var tx *gorm.DB
469-
if logType == LogTypeUnknown {
470-
tx = LOG_DB
471-
} else {
472-
tx = LOG_DB.Where("logs.type = ?", logType)
473-
}
493+
func GetAllLogs(logType int, logFilter string, 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) {
494+
tx := applyLogFilter(applyLogTypeFilter(LOG_DB, logType), logFilter)
474495

475496
if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil {
476497
return nil, 0, err
@@ -554,13 +575,8 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
554575

555576
const logSearchCountLimit = 10000
556577

557-
func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
558-
var tx *gorm.DB
559-
if logType == LogTypeUnknown {
560-
tx = LOG_DB.Where("logs.user_id = ?", userId)
561-
} else {
562-
tx = LOG_DB.Where("logs.user_id = ? and logs.type = ?", userId, logType)
563-
}
578+
func GetUserLogs(userId int, logType int, logFilter string, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) {
579+
tx := applyLogFilter(applyLogTypeFilter(LOG_DB.Where("logs.user_id = ?", userId), logType), logFilter)
564580

565581
if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil {
566582
return nil, 0, err
@@ -635,12 +651,15 @@ type Stat struct {
635651
Tpm int `json:"tpm"`
636652
}
637653

638-
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) {
654+
func SumUsedQuota(logType int, logFilter string, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string) (stat Stat, err error) {
639655
tx := LOG_DB.Table("logs").Select("sum(quota) quota")
640656

641657
// 为rpm和tpm创建单独的查询
642658
rpmTpmQuery := LOG_DB.Table("logs").Select("count(*) rpm, sum(prompt_tokens) + sum(completion_tokens) tpm")
643659

660+
tx = applyLogFilter(tx, logFilter)
661+
rpmTpmQuery = applyLogFilter(rpmTpmQuery, logFilter)
662+
644663
if tx, err = applyExplicitLogTextFilter(tx, "username", username); err != nil {
645664
return stat, err
646665
}

‎model/log_test.go‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package model
22

33
import (
44
"testing"
5+
"time"
56

67
"github.com/MAX-API-Next/MAX-API/common"
78
"github.com/stretchr/testify/require"
@@ -19,6 +20,116 @@ func withLogAuditSettings(t *testing.T, requestEnabled bool, responseEnabled boo
1920
})
2021
}
2122

23+
func TestGetAllLogsRetryFilter(t *testing.T) {
24+
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
25+
t.Cleanup(func() {
26+
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
27+
})
28+
29+
logs := createRetryFilterLogs(t)
30+
31+
got, total, err := GetAllLogs(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, 10, 0, "", "", "")
32+
require.NoError(t, err)
33+
require.EqualValues(t, 2, total)
34+
require.Len(t, got, 2)
35+
require.ElementsMatch(t, []int{logs[0].Id, logs[1].Id}, []int{got[0].Id, got[1].Id})
36+
}
37+
38+
func TestGetUserLogsRetryFilter(t *testing.T) {
39+
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
40+
t.Cleanup(func() {
41+
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
42+
})
43+
44+
logs := createRetryFilterLogs(t)
45+
46+
got, total, err := GetUserLogs(1, LogTypeUnknown, LogFilterRetry, 0, 0, "", "", 0, 10, "", "", "")
47+
require.NoError(t, err)
48+
require.EqualValues(t, 2, total)
49+
require.Len(t, got, 2)
50+
require.ElementsMatch(t, []int{logs[0].Id, logs[1].Id}, []int{got[0].LogId, got[1].LogId})
51+
}
52+
53+
func TestSumUsedQuotaRetryFilter(t *testing.T) {
54+
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
55+
t.Cleanup(func() {
56+
require.NoError(t, LOG_DB.Where("1 = 1").Delete(&Log{}).Error)
57+
})
58+
59+
createRetryFilterLogs(t)
60+
61+
stat, err := SumUsedQuota(LogTypeUnknown, LogFilterRetry, 0, 0, "", "", "", 0, "")
62+
require.NoError(t, err)
63+
require.Equal(t, 300, stat.Quota)
64+
require.Equal(t, 2, stat.Rpm)
65+
require.Equal(t, 30, stat.Tpm)
66+
}
67+
68+
func createRetryFilterLogs(t *testing.T) []Log {
69+
t.Helper()
70+
71+
logs := []Log{
72+
{
73+
UserId: 1,
74+
CreatedAt: time.Now().Unix() - 10,
75+
Type: LogTypeConsume,
76+
Quota: 100,
77+
PromptTokens: 5,
78+
CompletionTokens: 10,
79+
Other: common.MapToJsonStr(map[string]interface{}{
80+
"retry_log": true,
81+
"admin_info": map[string]interface{}{
82+
"use_channel": []string{"1", "2"},
83+
},
84+
}),
85+
},
86+
{
87+
UserId: 1,
88+
CreatedAt: time.Now().Unix() - 20,
89+
Type: LogTypeConsume,
90+
Quota: 200,
91+
PromptTokens: 7,
92+
CompletionTokens: 8,
93+
Other: common.MapToJsonStr(map[string]interface{}{
94+
"empty_retry": true,
95+
"admin_info": map[string]interface{}{
96+
"use_channel": []string{"3"},
97+
},
98+
}),
99+
},
100+
{
101+
UserId: 1,
102+
CreatedAt: time.Now().Unix() - 30,
103+
Type: LogTypeConsume,
104+
Quota: 400,
105+
PromptTokens: 11,
106+
CompletionTokens: 13,
107+
Other: common.MapToJsonStr(map[string]interface{}{
108+
"admin_info": map[string]interface{}{
109+
"use_channel": []string{"4"},
110+
},
111+
"empty_output_types": []string{"message", "function_call"},
112+
}),
113+
},
114+
{
115+
UserId: 2,
116+
CreatedAt: time.Now().Unix() - 40,
117+
Type: LogTypeConsume,
118+
Quota: 800,
119+
PromptTokens: 17,
120+
CompletionTokens: 19,
121+
Other: common.MapToJsonStr(map[string]interface{}{
122+
"admin_info": map[string]interface{}{
123+
"use_channel": []string{"5"},
124+
},
125+
}),
126+
},
127+
}
128+
require.NoError(t, LOG_DB.Create(&logs).Error)
129+
130+
return logs
131+
}
132+
22133
func TestFormatUserLogDetailExposesOnlyUserAuditContent(t *testing.T) {
23134
withLogAuditSettings(t, true, true)
24135

‎service/log_info_generate.go‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,17 @@ func appendRequestPath(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, other
3333
}
3434
}
3535

36+
func AppendRetryLogInfo(ctx *gin.Context, adminInfo map[string]interface{}, other map[string]interface{}) {
37+
if ctx == nil || adminInfo == nil || other == nil {
38+
return
39+
}
40+
useChannel := ctx.GetStringSlice("use_channel")
41+
adminInfo["use_channel"] = useChannel
42+
if len(useChannel) > 1 {
43+
other["retry_log"] = true
44+
}
45+
}
46+
3647
func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelRatio, groupRatio, completionRatio float64,
3748
cacheTokens int, cacheRatio float64, modelPrice float64, userGroupRatio float64) map[string]interface{} {
3849
other := make(map[string]interface{})
@@ -58,7 +69,7 @@ func GenerateTextOtherInfo(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, m
5869
}
5970

6071
adminInfo := make(map[string]interface{})
61-
adminInfo["use_channel"] = ctx.GetStringSlice("use_channel")
72+
AppendRetryLogInfo(ctx, adminInfo, other)
6273
isMultiKey := common.GetContextKeyBool(ctx, constant.ContextKeyChannelIsMultiKey)
6374
if isMultiKey {
6475
adminInfo["is_multi_key"] = true

‎service/log_info_generate_test.go‎

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
package service
2+
3+
import (
4+
"net/http/httptest"
5+
"testing"
6+
"time"
7+
8+
relaycommon "github.com/MAX-API-Next/MAX-API/relay/common"
9+
"github.com/gin-gonic/gin"
10+
"github.com/stretchr/testify/require"
11+
)
12+
13+
func TestGenerateTextOtherInfoMarksRetryLog(t *testing.T) {
14+
gin.SetMode(gin.TestMode)
15+
rec := httptest.NewRecorder()
16+
ctx, _ := gin.CreateTestContext(rec)
17+
ctx.Set("use_channel", []string{"1", "2"})
18+
19+
now := time.Now()
20+
info := &relaycommon.RelayInfo{
21+
StartTime: now,
22+
FirstResponseTime: now,
23+
ChannelMeta: &relaycommon.ChannelMeta{},
24+
}
25+
26+
other := GenerateTextOtherInfo(ctx, info, 1, 1, 1, 0, 0, 0, 1)
27+
28+
require.Equal(t, true, other["retry_log"])
29+
adminInfo, ok := other["admin_info"].(map[string]interface{})
30+
require.True(t, ok)
31+
require.Equal(t, []string{"1", "2"}, adminInfo["use_channel"])
32+
}
33+
34+
func TestGenerateTextOtherInfoDoesNotMarkSingleChannelAsRetry(t *testing.T) {
35+
gin.SetMode(gin.TestMode)
36+
rec := httptest.NewRecorder()
37+
ctx, _ := gin.CreateTestContext(rec)
38+
ctx.Set("use_channel", []string{"1"})
39+
40+
now := time.Now()
41+
info := &relaycommon.RelayInfo{
42+
StartTime: now,
43+
FirstResponseTime: now,
44+
ChannelMeta: &relaycommon.ChannelMeta{},
45+
}
46+
47+
other := GenerateTextOtherInfo(ctx, info, 1, 1, 1, 0, 0, 0, 1)
48+
49+
require.NotContains(t, other, "retry_log")
50+
}

0 commit comments

Comments
 (0)