From b52dd06c6f5dda0625e3e441e3a577fadad6f1f7 Mon Sep 17 00:00:00 2001 From: zijiren233 Date: Thu, 23 Jan 2025 11:59:01 +0800 Subject: [PATCH] feat: ip log --- service/aiproxy/common/consume/consume.go | 8 +++++-- service/aiproxy/controller/log.go | 8 +++++++ service/aiproxy/middleware/distributor.go | 1 + service/aiproxy/model/log.go | 28 ++++++++++++++++++---- service/aiproxy/model/utils.go | 2 ++ service/aiproxy/relay/controller/handle.go | 4 ++++ 6 files changed, 45 insertions(+), 6 deletions(-) diff --git a/service/aiproxy/common/consume/consume.go b/service/aiproxy/common/consume/consume.go index 139de8675b1..05e1e0fb327 100644 --- a/service/aiproxy/common/consume/consume.go +++ b/service/aiproxy/common/consume/consume.go @@ -26,6 +26,7 @@ func AsyncConsume( inputPrice, outputPrice float64, content string, + ip string, requestDetail *model.RequestDetail, ) { if meta.IsChannelTest { @@ -49,6 +50,7 @@ func AsyncConsume( inputPrice, outputPrice, content, + ip, requestDetail, ) } @@ -62,6 +64,7 @@ func Consume( inputPrice, outputPrice float64, content string, + ip string, requestDetail *model.RequestDetail, ) { if meta.IsChannelTest { @@ -70,7 +73,7 @@ func Consume( amount := calculateAmount(ctx, usage, inputPrice, outputPrice, postGroupConsumer, meta) - err := recordConsume(meta, code, usage, inputPrice, outputPrice, content, requestDetail, amount) + err := recordConsume(meta, code, usage, inputPrice, outputPrice, content, ip, requestDetail, amount) if err != nil { log.Error("error batch record consume: " + err.Error()) } @@ -136,7 +139,7 @@ func processGroupConsume( return consumedAmount } -func recordConsume(meta *meta.Meta, code int, usage *relaymodel.Usage, inputPrice, outputPrice float64, content string, requestDetail *model.RequestDetail, amount float64) error { +func recordConsume(meta *meta.Meta, code int, usage *relaymodel.Usage, inputPrice, outputPrice float64, content string, ip string, requestDetail *model.RequestDetail, amount float64) error { promptTokens := 0 completionTokens := 0 if usage != nil { @@ -166,6 +169,7 @@ func recordConsume(meta *meta.Meta, code int, usage *relaymodel.Usage, inputPric meta.Endpoint, content, meta.Mode, + ip, requestDetail, ) } diff --git a/service/aiproxy/controller/log.go b/service/aiproxy/controller/log.go index c4b8e209f39..081201d3b66 100644 --- a/service/aiproxy/controller/log.go +++ b/service/aiproxy/controller/log.go @@ -43,6 +43,7 @@ func GetLogs(c *gin.Context) { mode, _ := strconv.Atoi(c.Query("mode")) codeType := c.Query("code_type") withBody, _ := strconv.ParseBool(c.Query("with_body")) + ip := c.Query("ip") result, err := model.GetLogs( group, startTimestampTime, @@ -59,6 +60,7 @@ func GetLogs(c *gin.Context) { mode, model.CodeType(codeType), withBody, + ip, ) if err != nil { middleware.ErrorResponse(c, http.StatusOK, err.Error()) @@ -104,6 +106,7 @@ func GetGroupLogs(c *gin.Context) { mode, _ := strconv.Atoi(c.Query("mode")) codeType := c.Query("code_type") withBody, _ := strconv.ParseBool(c.Query("with_body")) + ip := c.Query("ip") result, err := model.GetGroupLogs( group, startTimestampTime, @@ -120,6 +123,7 @@ func GetGroupLogs(c *gin.Context) { mode, model.CodeType(codeType), withBody, + ip, ) if err != nil { middleware.ErrorResponse(c, http.StatusOK, err.Error()) @@ -158,6 +162,7 @@ func SearchLogs(c *gin.Context) { mode, _ := strconv.Atoi(c.Query("mode")) codeType := c.Query("code_type") withBody, _ := strconv.ParseBool(c.Query("with_body")) + ip := c.Query("ip") result, err := model.SearchLogs( group, keyword, @@ -175,6 +180,7 @@ func SearchLogs(c *gin.Context) { mode, model.CodeType(codeType), withBody, + ip, ) if err != nil { middleware.ErrorResponse(c, http.StatusOK, err.Error()) @@ -217,6 +223,7 @@ func SearchGroupLogs(c *gin.Context) { mode, _ := strconv.Atoi(c.Query("mode")) codeType := c.Query("code_type") withBody, _ := strconv.ParseBool(c.Query("with_body")) + ip := c.Query("ip") result, err := model.SearchGroupLogs( group, keyword, @@ -234,6 +241,7 @@ func SearchGroupLogs(c *gin.Context) { mode, model.CodeType(codeType), withBody, + ip, ) if err != nil { middleware.ErrorResponse(c, http.StatusOK, err.Error()) diff --git a/service/aiproxy/middleware/distributor.go b/service/aiproxy/middleware/distributor.go index 1ffed4711af..2befdd35e34 100644 --- a/service/aiproxy/middleware/distributor.go +++ b/service/aiproxy/middleware/distributor.go @@ -200,6 +200,7 @@ func distribute(c *gin.Context, mode int) { 0, 0, errMsg, + c.ClientIP(), nil, ) abortLogWithMessage(c, http.StatusTooManyRequests, errMsg) diff --git a/service/aiproxy/model/log.go b/service/aiproxy/model/log.go index ab24eafe3ab..328ae14fb93 100644 --- a/service/aiproxy/model/log.go +++ b/service/aiproxy/model/log.go @@ -44,6 +44,7 @@ type Log struct { ChannelID int `gorm:"index" json:"channel"` Code int `gorm:"index" json:"code"` Mode int `json:"mode"` + IP string `json:"ip"` } func (l *Log) MarshalJSON() ([]byte, error) { @@ -118,6 +119,7 @@ func RecordConsumeLog( endpoint string, content string, mode int, + ip string, requestDetail *RequestDetail, ) error { defer func() { @@ -141,6 +143,7 @@ func RecordConsumeLog( TokenName: tokenName, Model: modelName, Mode: mode, + IP: ip, UsedAmount: amount, Price: price, CompletionPrice: completionPrice, @@ -203,6 +206,7 @@ func getLogs( mode int, codeType CodeType, withBody bool, + ip string, ) (int64, []*Log, error) { tx := LogDB.Model(&Log{}) if group != "" { @@ -235,6 +239,9 @@ func getLogs( if endpoint != "" { tx = tx.Where("endpoint = ?", endpoint) } + if ip != "" { + tx = tx.Where("ip = ?", ip) + } switch codeType { case CodeTypeSuccess: tx = tx.Where("code = 200") @@ -287,8 +294,9 @@ func GetLogs( mode int, codeType CodeType, withBody bool, + ip string, ) (*GetLogsResult, error) { - total, logs, err := getLogs(group, startTimestamp, endTimestamp, modelName, requestID, tokenID, tokenName, startIdx, num, channelID, endpoint, order, mode, codeType, withBody) + total, logs, err := getLogs(group, startTimestamp, endTimestamp, modelName, requestID, tokenID, tokenName, startIdx, num, channelID, endpoint, order, mode, codeType, withBody, ip) if err != nil { return nil, err } @@ -323,11 +331,12 @@ func GetGroupLogs( mode int, codeType CodeType, withBody bool, + ip string, ) (*GetGroupLogsResult, error) { if group == "" { return nil, errors.New("group is required") } - total, logs, err := getLogs(group, startTimestamp, endTimestamp, modelName, requestID, tokenID, tokenName, startIdx, num, channelID, endpoint, order, mode, codeType, withBody) + total, logs, err := getLogs(group, startTimestamp, endTimestamp, modelName, requestID, tokenID, tokenName, startIdx, num, channelID, endpoint, order, mode, codeType, withBody, ip) if err != nil { return nil, err } @@ -366,6 +375,7 @@ func searchLogs( mode int, codeType CodeType, withBody bool, + ip string, ) (int64, []*Log, error) { tx := LogDB.Model(&Log{}) if group != "" { @@ -400,6 +410,9 @@ func searchLogs( if channelID != 0 { tx = tx.Where("channel_id = ?", channelID) } + if ip != "" { + tx = tx.Where("ip = ?", ip) + } switch codeType { case CodeTypeSuccess: tx = tx.Where("code = 200") @@ -470,6 +483,11 @@ func searchLogs( } values = append(values, "%"+keyword+"%") + if ip != "" { + conditions = append(conditions, "ip = ?") + values = append(values, ip) + } + if len(conditions) > 0 { tx = tx.Where(fmt.Sprintf("(%s)", strings.Join(conditions, " OR ")), values...) } @@ -526,8 +544,9 @@ func SearchLogs( mode int, codeType CodeType, withBody bool, + ip string, ) (*GetLogsResult, error) { - total, logs, err := searchLogs(group, keyword, page, perPage, endpoint, requestID, tokenID, tokenName, modelName, startTimestamp, endTimestamp, channelID, order, mode, codeType, withBody) + total, logs, err := searchLogs(group, keyword, page, perPage, endpoint, requestID, tokenID, tokenName, modelName, startTimestamp, endTimestamp, channelID, order, mode, codeType, withBody, ip) if err != nil { return nil, err } @@ -563,11 +582,12 @@ func SearchGroupLogs( mode int, codeType CodeType, withBody bool, + ip string, ) (*GetGroupLogsResult, error) { if group == "" { return nil, errors.New("group is required") } - total, logs, err := searchLogs(group, keyword, page, perPage, endpoint, requestID, tokenID, tokenName, modelName, startTimestamp, endTimestamp, channelID, order, mode, codeType, withBody) + total, logs, err := searchLogs(group, keyword, page, perPage, endpoint, requestID, tokenID, tokenName, modelName, startTimestamp, endTimestamp, channelID, order, mode, codeType, withBody, ip) if err != nil { return nil, err } diff --git a/service/aiproxy/model/utils.go b/service/aiproxy/model/utils.go index 249554235e0..148454496f5 100644 --- a/service/aiproxy/model/utils.go +++ b/service/aiproxy/model/utils.go @@ -59,6 +59,7 @@ func BatchRecordConsume( endpoint string, content string, mode int, + ip string, requestDetail *RequestDetail, ) error { errs := []error{} @@ -79,6 +80,7 @@ func BatchRecordConsume( endpoint, content, mode, + ip, requestDetail, ) if err != nil { diff --git a/service/aiproxy/relay/controller/handle.go b/service/aiproxy/relay/controller/handle.go index db985ec1bdb..7ad2767f1ea 100644 --- a/service/aiproxy/relay/controller/handle.go +++ b/service/aiproxy/relay/controller/handle.go @@ -42,6 +42,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa 0, 0, errMsg, + c.ClientIP(), nil, ) return openai.ErrorWrapperWithMessage( @@ -76,6 +77,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa 0, 0, err.Error(), + c.ClientIP(), detail, ) return openai.ErrorWrapper(err, "invalid_request", http.StatusBadRequest) @@ -111,6 +113,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa preCheckReq.InputPrice, preCheckReq.OutputPrice, respErr.Error.JSONOrEmpty(), + c.ClientIP(), detail, ) return respErr @@ -125,6 +128,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa preCheckReq.InputPrice, preCheckReq.OutputPrice, "", + c.ClientIP(), nil, )