Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🐛 fix sum select statements

Browse Source
This commit is contained in:
Martial BE 2023-12-05 15:08:42 +08:00
parent 8bcaf182bc
commit 365744a040
No known key found for this signature in database
GPG Key ID: D06C32DF0EDB9084
1 changed files with 16 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
model/log.go
View File

@ -3,8 +3,9 @@ package model
import ( import (
"context" "context"
"fmt" "fmt"
"gorm.io/gorm"
"one-api/common" "one-api/common"
"gorm.io/gorm"
) )
type Log struct { type Log struct {
@ -134,7 +135,7 @@ func SearchUserLogs(userId int, keyword string) (logs []*Log, err error) {
} }
func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (quota int) { func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int) (quota int) {
tx := DB.Table("logs").Select("ifnull(sum(quota),0)") tx := DB.Table("logs").Select(assembleSumSelectStr("quota"))
if username != "" { if username != "" {
tx = tx.Where("username = ?", username) tx = tx.Where("username = ?", username)
} }
@ -158,7 +159,7 @@ func SumUsedQuota(logType int, startTimestamp int64, endTimestamp int64, modelNa
} }
func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) { func SumUsedToken(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string) (token int) {
tx := DB.Table("logs").Select("ifnull(sum(prompt_tokens),0) + ifnull(sum(completion_tokens),0)") tx := DB.Table("logs").Select(assembleSumSelectStr("prompt_tokens") + " + " + assembleSumSelectStr("completion_tokens"))
if username != "" { if username != "" {
tx = tx.Where("username = ?", username) tx = tx.Where("username = ?", username)
} }
@ -182,3 +183,15 @@ func DeleteOldLog(targetTimestamp int64) (int64, error) {
result := DB.Where("created_at < ?", targetTimestamp).Delete(&Log{}) result := DB.Where("created_at < ?", targetTimestamp).Delete(&Log{})
return result.RowsAffected, result.Error return result.RowsAffected, result.Error
} }
func assembleSumSelectStr(selectStr string) string {
sumSelectStr := "%s(sum(%s),0)"
nullfunc := "ifnull"
if common.UsingPostgreSQL {
nullfunc = "coalesce"
}
sumSelectStr = fmt.Sprintf(sumSelectStr, nullfunc, selectStr)
return sumSelectStr
}