Skip to content

Commit

Permalink
feat: ip log
Browse files Browse the repository at this point in the history
  • Loading branch information
@zijiren233
zijiren233 committed Jan 23, 2025
1 parent 6e8b5e9 commit b52dd06
Show file tree
Hide file tree
Showing 6 changed files with 45 additions and 6 deletions.
8 changes: 6 additions & 2 deletions service/aiproxy/common/consume/consume.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ func AsyncConsume(
inputPrice,
outputPrice float64,
content string,
ip string,
requestDetail *model.RequestDetail,
) {
if meta.IsChannelTest {
Expand All @@ -49,6 +50,7 @@ func AsyncConsume(
inputPrice,
outputPrice,
content,
ip,
requestDetail,
)
}
Expand All @@ -62,6 +64,7 @@ func Consume(
inputPrice,
outputPrice float64,
content string,
ip string,
requestDetail *model.RequestDetail,
) {
if meta.IsChannelTest {
Expand All @@ -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())
}
Expand Down Expand Up @@ -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 {
Expand Down Expand Up @@ -166,6 +169,7 @@ func recordConsume(meta *meta.Meta, code int, usage *relaymodel.Usage, inputPric
meta.Endpoint,
content,
meta.Mode,
ip,
requestDetail,
)
}
8 changes: 8 additions & 0 deletions service/aiproxy/controller/log.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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())
Expand Down Expand Up @@ -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,
Expand All @@ -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())
Expand Down Expand Up @@ -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,
Expand All @@ -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())
Expand Down Expand Up @@ -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,
Expand All @@ -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())
Expand Down
1 change: 1 addition & 0 deletions service/aiproxy/middleware/distributor.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,7 @@ func distribute(c *gin.Context, mode int) {
0,
0,
errMsg,
c.ClientIP(),
nil,
)
abortLogWithMessage(c, http.StatusTooManyRequests, errMsg)
Expand Down
28 changes: 24 additions & 4 deletions service/aiproxy/model/log.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -118,6 +119,7 @@ func RecordConsumeLog(
endpoint string,
content string,
mode int,
ip string,
requestDetail *RequestDetail,
) error {
defer func() {
Expand All @@ -141,6 +143,7 @@ func RecordConsumeLog(
TokenName: tokenName,
Model: modelName,
Mode: mode,
IP: ip,
UsedAmount: amount,
Price: price,
CompletionPrice: completionPrice,
Expand Down Expand Up @@ -203,6 +206,7 @@ func getLogs(
mode int,
codeType CodeType,
withBody bool,
ip string,
) (int64, []*Log, error) {
tx := LogDB.Model(&Log{})
if group != "" {
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -366,6 +375,7 @@ func searchLogs(
mode int,
codeType CodeType,
withBody bool,
ip string,
) (int64, []*Log, error) {
tx := LogDB.Model(&Log{})
if group != "" {
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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...)
}
Expand Down Expand Up @@ -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
}
Expand Down Expand Up @@ -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
}
Expand Down
2 changes: 2 additions & 0 deletions service/aiproxy/model/utils.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,6 +59,7 @@ func BatchRecordConsume(
endpoint string,
content string,
mode int,
ip string,
requestDetail *RequestDetail,
) error {
errs := []error{}
Expand All @@ -79,6 +80,7 @@ func BatchRecordConsume(
endpoint,
content,
mode,
ip,
requestDetail,
)
if err != nil {
Expand Down
4 changes: 4 additions & 0 deletions service/aiproxy/relay/controller/handle.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa
0,
0,
errMsg,
c.ClientIP(),
nil,
)
return openai.ErrorWrapperWithMessage(
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -125,6 +128,7 @@ func Handle(meta *meta.Meta, c *gin.Context, preProcess func() (*PreCheckGroupBa
preCheckReq.InputPrice,
preCheckReq.OutputPrice,
"",
c.ClientIP(),
nil,
)

Expand Down

0 comments on commit b52dd06

Please sign in to comment.