xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-16 21:23:44 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

🔖 chore: migration logger package

Browse Source
This commit is contained in:
MartialBE
2024-05-29 01:04:23 +08:00
parent 79524108a3
commit ce12558ad6
44 changed files with 207 additions and 174 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
model/balancer.go
View File

@@ -4,6 +4,7 @@ import (
"errors"
"math/rand"
"one-api/common"
"one-api/common/logger"
"one-api/common/utils"
"strings"
"sync"
@@ -162,7 +163,7 @@ func (cc *ChannelsChooser) Load() {
abilities, err := GetAbilityChannelGroup()
if err != nil {
common.SysLog("get enabled abilities failed: " + err.Error())
logger.SysLog("get enabled abilities failed: " + err.Error())
return
}
@@ -216,5 +217,5 @@ func (cc *ChannelsChooser) Load() {
cc.Channels = newChannels
cc.Match = newMatchList
cc.Unlock()
common.SysLog("channels Load success")
logger.SysLog("channels Load success")
}

9
model/cache.go
View File

@@ -4,6 +4,7 @@ import (
"encoding/json"
"fmt"
"one-api/common"
"one-api/common/logger"
"strconv"
"time"
)
@@ -34,7 +35,7 @@ func CacheGetTokenByKey(key string) (*Token, error) {
}
err = common.RedisSet(fmt.Sprintf("token:%s", key), string(jsonBytes), time.Duration(TokenCacheSeconds)*time.Second)
if err != nil {
common.SysError("Redis set token error: " + err.Error())
logger.SysError("Redis set token error: " + err.Error())
}
return &token, nil
}
@@ -54,7 +55,7 @@ func CacheGetUserGroup(id int) (group string, err error) {
}
err = common.RedisSet(fmt.Sprintf("user_group:%d", id), group, time.Duration(TokenCacheSeconds)*time.Second)
if err != nil {
common.SysError("Redis set user group error: " + err.Error())
logger.SysError("Redis set user group error: " + err.Error())
}
}
return group, err
@@ -72,7 +73,7 @@ func CacheGetUserQuota(id int) (quota int, err error) {
}
err = common.RedisSet(fmt.Sprintf("user_quota:%d", id), fmt.Sprintf("%d", quota), time.Duration(TokenCacheSeconds)*time.Second)
if err != nil {
common.SysError("Redis set user quota error: " + err.Error())
logger.SysError("Redis set user quota error: " + err.Error())
}
return quota, err
}
@@ -119,7 +120,7 @@ func CacheIsUserEnabled(userId int) (bool, error) {
}
err = common.RedisSet(fmt.Sprintf("user_enabled:%d", userId), enabled, time.Duration(TokenCacheSeconds)*time.Second)
if err != nil {
common.SysError("Redis set user enabled error: " + err.Error())
logger.SysError("Redis set user enabled error: " + err.Error())
}
return userEnabled, err
}

11
model/channel.go
View File

@@ -2,6 +2,7 @@ package model
import (
"one-api/common"
"one-api/common/logger"
"one-api/common/utils"
"strings"
@@ -240,7 +241,7 @@ func (channel *Channel) UpdateResponseTime(responseTime int64) {
ResponseTime: int(responseTime),
}).Error
if err != nil {
common.SysError("failed to update response time: " + err.Error())
logger.SysError("failed to update response time: " + err.Error())
}
}
@@ -250,7 +251,7 @@ func (channel *Channel) UpdateBalance(balance float64) {
Balance: balance,
}).Error
if err != nil {
common.SysError("failed to update balance: " + err.Error())
logger.SysError("failed to update balance: " + err.Error())
}
}
@@ -283,11 +284,11 @@ func (channel *Channel) StatusToStr() string {
func UpdateChannelStatusById(id int, status int) {
err := UpdateAbilityStatus(id, status == common.ChannelStatusEnabled)
if err != nil {
common.SysError("failed to update ability status: " + err.Error())
logger.SysError("failed to update ability status: " + err.Error())
}
err = DB.Model(&Channel{}).Where("id = ?", id).Update("status", status).Error
if err != nil {
common.SysError("failed to update channel status: " + err.Error())
logger.SysError("failed to update channel status: " + err.Error())
}
if err == nil {
@@ -307,7 +308,7 @@ func UpdateChannelUsedQuota(id int, quota int) {
func updateChannelUsedQuota(id int, quota int) {
err := DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", gorm.Expr("used_quota + ?", quota)).Error
if err != nil {
common.SysError("failed to update channel used quota: " + err.Error())
logger.SysError("failed to update channel used quota: " + err.Error())
}
}

7
model/log.go
View File

@@ -4,6 +4,7 @@ import (
"context"
"fmt"
"one-api/common"
"one-api/common/logger"
"one-api/common/utils"
"gorm.io/gorm"
@@ -48,12 +49,12 @@ func RecordLog(userId int, logType int, content string) {
}
err := DB.Create(log).Error
if err != nil {
common.SysError("failed to record log: " + err.Error())
logger.SysError("failed to record log: " + err.Error())
}
}
func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptTokens int, completionTokens int, modelName string, tokenName string, quota int, content string, requestTime int) {
common.LogInfo(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
logger.LogInfo(ctx, fmt.Sprintf("record consume log: userId=%d, channelId=%d, promptTokens=%d, completionTokens=%d, modelName=%s, tokenName=%s, quota=%d, content=%s", userId, channelId, promptTokens, completionTokens, modelName, tokenName, quota, content))
if !common.LogConsumeEnabled {
return
}
@@ -73,7 +74,7 @@ func RecordConsumeLog(ctx context.Context, userId int, channelId int, promptToke
}
err := DB.Create(log).Error
if err != nil {
common.LogError(ctx, "failed to record log: "+err.Error())
logger.LogError(ctx, "failed to record log: "+err.Error())
}
}

19
model/main.go
View File

@@ -3,6 +3,7 @@ package model
import (
"fmt"
"one-api/common"
"one-api/common/logger"
"one-api/common/utils"
"strconv"
"strings"
@@ -20,7 +21,7 @@ var DB *gorm.DB
func SetupDB() {
err := InitDB()
if err != nil {
common.FatalLog("failed to initialize database: " + err.Error())
logger.FatalLog("failed to initialize database: " + err.Error())
}
ChannelGroup.Load()
common.RootUserEmail = GetRootUserEmail()
@@ -28,7 +29,7 @@ func SetupDB() {
if viper.GetBool("batch_update_enabled") {
common.BatchUpdateEnabled = true
common.BatchUpdateInterval = utils.GetOrDefault("batch_update_interval", 5)
common.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s")
logger.SysLog("batch update enabled with interval " + strconv.Itoa(common.BatchUpdateInterval) + "s")
InitBatchUpdater()
}
}
@@ -37,7 +38,7 @@ func createRootAccountIfNeed() error {
var user User
//if user.Status != common.UserStatusEnabled {
if err := DB.First(&user).Error; err != nil {
common.SysLog("no user exists, create a root user for you: username is root, password is 123456")
logger.SysLog("no user exists, create a root user for you: username is root, password is 123456")
hashedPassword, err := common.Password2Hash("123456")
if err != nil {
return err
@@ -61,7 +62,7 @@ func chooseDB() (*gorm.DB, error) {
dsn := viper.GetString("sql_dsn")
if strings.HasPrefix(dsn, "postgres://") {
// Use PostgreSQL
common.SysLog("using PostgreSQL as database")
logger.SysLog("using PostgreSQL as database")
common.UsingPostgreSQL = true
return gorm.Open(postgres.New(postgres.Config{
DSN: dsn,
@@ -71,13 +72,13 @@ func chooseDB() (*gorm.DB, error) {
})
}
// Use MySQL
common.SysLog("using MySQL as database")
logger.SysLog("using MySQL as database")
return gorm.Open(mysql.Open(dsn), &gorm.Config{
PrepareStmt: true, // precompile SQL
})
}
// Use SQLite
common.SysLog("SQL_DSN not set, using SQLite as database")
logger.SysLog("SQL_DSN not set, using SQLite as database")
common.UsingSQLite = true
config := fmt.Sprintf("?_busy_timeout=%d", utils.GetOrDefault("sqlite_busy_timeout", 3000))
return gorm.Open(sqlite.Open(viper.GetString("sqlite_path")+config), &gorm.Config{
@@ -104,7 +105,7 @@ func InitDB() (err error) {
if !common.IsMasterNode {
return nil
}
common.SysLog("database migration started")
logger.SysLog("database migration started")
migration(DB)
@@ -152,11 +153,11 @@ func InitDB() (err error) {
if err != nil {
return err
}
common.SysLog("database migrated")
logger.SysLog("database migrated")
err = createRootAccountIfNeed()
return err
} else {
common.FatalLog(err)
logger.FatalLog(err)
}
return err
}

4
model/migrate.go
View File

@@ -1,7 +1,7 @@
package model
import (
"one-api/common"
"one-api/common/logger"
"github.com/go-gormigrate/gormigrate/v2"
"gorm.io/gorm"
@@ -22,7 +22,7 @@ func removeKeyIndexMigration() *gormigrate.Migration {
err := tx.Migrator().DropIndex(&Channel{}, "idx_channels_key")
if err != nil {
common.SysLog("remove idx_channels_key Failure: " + err.Error())
logger.SysLog("remove idx_channels_key Failure: " + err.Error())
}
return nil
},

5
model/option.go
View File

@@ -2,6 +2,7 @@ package model
import (
"one-api/common"
"one-api/common/logger"
"strconv"
"strings"
"time"
@@ -90,7 +91,7 @@ func loadOptionsFromDatabase() {
for _, option := range options {
err := updateOptionMap(option.Key, option.Value)
if err != nil {
common.SysError("failed to update option map: " + err.Error())
logger.SysError("failed to update option map: " + err.Error())
}
}
}
@@ -98,7 +99,7 @@ func loadOptionsFromDatabase() {
func SyncOptions(frequency int) {
for {
time.Sleep(time.Duration(frequency) * time.Second)
common.SysLog("syncing options from database")
logger.SysLog("syncing options from database")
loadOptionsFromDatabase()
}
}

13
model/token.go
View File

@@ -4,6 +4,7 @@ import (
"errors"
"fmt"
"one-api/common"
"one-api/common/logger"
"one-api/common/stmp"
"one-api/common/utils"
@@ -58,7 +59,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
}
token, err = CacheGetTokenByKey(key)
if err != nil {
common.SysError("CacheGetTokenByKey failed: " + err.Error())
logger.SysError("CacheGetTokenByKey failed: " + err.Error())
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, errors.New("无效的令牌")
}
@@ -77,7 +78,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
token.Status = common.TokenStatusExpired
err := token.SelectUpdate()
if err != nil {
common.SysError("failed to update token status" + err.Error())
logger.SysError("failed to update token status" + err.Error())
}
}
return nil, errors.New("该令牌已过期")
@@ -88,7 +89,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
token.Status = common.TokenStatusExhausted
err := token.SelectUpdate()
if err != nil {
common.SysError("failed to update token status" + err.Error())
logger.SysError("failed to update token status" + err.Error())
}
}
return nil, errors.New("该令牌额度已用尽")
@@ -254,12 +255,12 @@ func sendQuotaWarningEmail(userId int, userQuota int, noMoreQuota bool) {
user := User{Id: userId}
if err := user.FillUserById(); err != nil {
common.SysError("failed to fetch user email: " + err.Error())
logger.SysError("failed to fetch user email: " + err.Error())
return
}
if user.Email == "" {
common.SysError("user email is empty")
logger.SysError("user email is empty")
return
}
@@ -271,7 +272,7 @@ func sendQuotaWarningEmail(userId int, userQuota int, noMoreQuota bool) {
err := stmp.SendQuotaWarningCodeEmail(userName, user.Email, userQuota, noMoreQuota)
if err != nil {
common.SysError("failed to send email" + err.Error())
logger.SysError("failed to send email" + err.Error())
}
}

9
model/user.go
View File

@@ -4,6 +4,7 @@ import (
"errors"
"fmt"
"one-api/common"
"one-api/common/logger"
"one-api/common/utils"
"strings"
@@ -306,7 +307,7 @@ func IsAdmin(userId int) bool {
var user User
err := DB.Where("id = ?", userId).Select("role").Find(&user).Error
if err != nil {
common.SysError("no such user " + err.Error())
logger.SysError("no such user " + err.Error())
return false
}
return user.Role >= common.RoleAdminUser
@@ -415,7 +416,7 @@ func updateUserUsedQuotaAndRequestCount(id int, quota int, count int) {
},
).Error
if err != nil {
common.SysError("failed to update user used quota and request count: " + err.Error())
logger.SysError("failed to update user used quota and request count: " + err.Error())
}
}
@@ -426,14 +427,14 @@ func updateUserUsedQuota(id int, quota int) {
},
).Error
if err != nil {
common.SysError("failed to update user used quota: " + err.Error())
logger.SysError("failed to update user used quota: " + err.Error())
}
}
func updateUserRequestCount(id int, count int) {
err := DB.Model(&User{}).Where("id = ?", id).Update("request_count", gorm.Expr("request_count + ?", count)).Error
if err != nil {
common.SysError("failed to update user request count: " + err.Error())
logger.SysError("failed to update user request count: " + err.Error())
}
}

9
model/utils.go
View File

@@ -2,6 +2,7 @@ package model
import (
"one-api/common"
"one-api/common/logger"
"sync"
"time"
)
@@ -45,7 +46,7 @@ func addNewRecord(type_ int, id int, value int) {
}
func batchUpdate() {
common.SysLog("batch update started")
logger.SysLog("batch update started")
for i := 0; i < BatchUpdateTypeCount; i++ {
batchUpdateLocks[i].Lock()
store := batchUpdateStores[i]
@@ -57,12 +58,12 @@ func batchUpdate() {
case BatchUpdateTypeUserQuota:
err := increaseUserQuota(key, value)
if err != nil {
common.SysError("failed to batch update user quota: " + err.Error())
logger.SysError("failed to batch update user quota: " + err.Error())
}
case BatchUpdateTypeTokenQuota:
err := increaseTokenQuota(key, value)
if err != nil {
common.SysError("failed to batch update token quota: " + err.Error())
logger.SysError("failed to batch update token quota: " + err.Error())
}
case BatchUpdateTypeUsedQuota:
updateUserUsedQuota(key, value)
@@ -73,5 +74,5 @@ func batchUpdate() {
}
}
}
common.SysLog("batch update finished")
logger.SysLog("batch update finished")
}