chaos/chaos-api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into feat_subscribe_sp1

Browse Source
This commit is contained in:
Little Write
2025-09-27 22:54:52 +08:00
parent 764aaef8d9 1d4d243f66
commit 6a2071cb38
617 changed files with 64415 additions and 27156 deletions
Download Patch File Download Diff File Expand all files Collapse all files
model/ability.go
+26 -6
View File
@@ -136,13 +136,13 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
}
}
} else {
return nil, errors.New("channel not found")
return nil, nil
}
err = DB.First(&channel, "id = ?", channel.Id).Error
return &channel, err
}
func (channel *Channel) AddAbilities() error {
func (channel *Channel) AddAbilities(tx *gorm.DB) error {
models_ := strings.Split(channel.Models, ",")
groups_ := strings.Split(channel.Group, ",")
abilitySet := make(map[string]struct{})
@@ -169,8 +169,13 @@ func (channel *Channel) AddAbilities() error {
if len(abilities) == 0 {
return nil
}
// choose DB or provided tx
useDB := DB
if tx != nil {
useDB = tx
}
for _, chunk := range lo.Chunk(abilities, 50) {
err := DB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
err := useDB.Clauses(clause.OnConflict{DoNothing: true}).Create(&chunk).Error
if err != nil {
return err
}
@@ -284,6 +289,21 @@ func FixAbility() (int, int, error) {
return 0, 0, errors.New("已经有一个修复任务在运行中，请稍后再试")
}
defer fixLock.Unlock()
// truncate abilities table
if common.UsingSQLite {
err := DB.Exec("DELETE FROM abilities").Error
if err != nil {
common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
return 0, 0, err
}
} else {
err := DB.Exec("TRUNCATE TABLE abilities").Error
if err != nil {
common.SysLog(fmt.Sprintf("Truncate abilities failed: %s", err.Error()))
return 0, 0, err
}
}
var channels []*Channel
// Find all channels
err := DB.Model(&Channel{}).Find(&channels).Error
@@ -300,15 +320,15 @@ func FixAbility() (int, int, error) {
// Delete all abilities of this channel
err = DB.Where("channel_id IN ?", ids).Delete(&Ability{}).Error
if err != nil {
common.SysError(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error()))
failCount += len(chunk)
continue
}
// Then add new abilities
for _, channel := range chunk {
err = channel.AddAbilities()
err = channel.AddAbilities(nil)
if err != nil {
common.SysError(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error()))
common.SysLog(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error()))
failCount++
} else {
successCount++
model/channel.go
+147 -64
View File
@@ -13,6 +13,7 @@ import (
"strings"
"sync"
"github.com/samber/lo"
"gorm.io/gorm"
)
@@ -44,16 +45,25 @@ type Channel struct {
Tag *string `json:"tag" gorm:"index"`
Setting *string `json:"setting" gorm:"type:text"` // 渠道额外设置
ParamOverride *string `json:"param_override" gorm:"type:text"`
HeaderOverride *string `json:"header_override" gorm:"type:text"`
Remark string `json:"remark,omitempty" gorm:"type:varchar(255)" validate:"max=255"`
// add after v0.8.5
ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"`
OtherSettings string `json:"settings" gorm:"column:settings"` // 其他设置，存储azure版本等不需要检索的信息，详见dto.ChannelOtherSettings
// cache info
Keys []string `json:"-" gorm:"-"`
}
type ChannelInfo struct {
IsMultiKey bool `json:"is_multi_key"` // 是否多Key模式
MultiKeySize int `json:"multi_key_size"` // 多Key模式下的Key数量
MultiKeyStatusList map[int]int `json:"multi_key_status_list"` // key状态列表，key index -> status
MultiKeyPollingIndex int `json:"multi_key_polling_index"` // 多Key模式下轮询的key索引
MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
IsMultiKey bool `json:"is_multi_key"` // 是否多Key模式
MultiKeySize int `json:"multi_key_size"` // 多Key模式下的Key数量
MultiKeyStatusList map[int]int `json:"multi_key_status_list"` // key状态列表，key index -> status
MultiKeyDisabledReason map[int]string `json:"multi_key_disabled_reason,omitempty"` // key禁用原因列表，key index -> reason
MultiKeyDisabledTime map[int]int64 `json:"multi_key_disabled_time,omitempty"` // key禁用时间列表，key index -> time
MultiKeyPollingIndex int `json:"multi_key_polling_index"` // 多Key模式下轮询的key索引
MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"`
}
// Value implements driver.Valuer interface
@@ -67,15 +77,18 @@ func (c *ChannelInfo) Scan(value interface{}) error {
return common.Unmarshal(bytesValue, c)
}
func (channel *Channel) getKeys() []string {
func (channel *Channel) GetKeys() []string {
if channel.Key == "" {
return []string{}
}
if len(channel.Keys) > 0 {
return channel.Keys
}
trimmed := strings.TrimSpace(channel.Key)
// If the key starts with '[', try to parse it as a JSON array (e.g., for Vertex AI scenarios)
if strings.HasPrefix(trimmed, "[") {
var arr []json.RawMessage
if err := json.Unmarshal([]byte(trimmed), &arr); err == nil {
if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
res := make([]string, len(arr))
for i, v := range arr {
res[i] = string(v)
@@ -95,12 +108,16 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
}
// Obtain all keys (split by \n)
keys := channel.getKeys()
keys := channel.GetKeys()
if len(keys) == 0 {
// No keys available, return error, should disable the channel
return "", 0, types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey)
}
lock := GetChannelPollingLock(channel.Id)
lock.Lock()
defer lock.Unlock()
statusList := channel.ChannelInfo.MultiKeyStatusList
// helper to get key status, default to enabled when missing
getStatus := func(idx int) int {
@@ -132,13 +149,10 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
return keys[selectedIdx], selectedIdx, nil
case constant.MultiKeyModePolling:
// Use channel-specific lock to ensure thread-safe polling
lock := getChannelPollingLock(channel.Id)
lock.Lock()
defer lock.Unlock()
channelInfo, err := CacheGetChannelInfo(channel.Id)
if err != nil {
return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed)
return "", 0, types.NewError(err, types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
}
//println("before polling index:", channel.ChannelInfo.MultiKeyPollingIndex)
defer func() {
@@ -197,9 +211,9 @@ func (channel *Channel) GetGroups() []string {
func (channel *Channel) GetOtherInfo() map[string]interface{} {
otherInfo := make(map[string]interface{})
if channel.OtherInfo != "" {
err := json.Unmarshal([]byte(channel.OtherInfo), &otherInfo)
err := common.Unmarshal([]byte(channel.OtherInfo), &otherInfo)
if err != nil {
common.SysError("failed to unmarshal other info: " + err.Error())
common.SysLog(fmt.Sprintf("failed to unmarshal other info: channel_id=%d, tag=%s, name=%s, error=%v", channel.Id, channel.GetTag(), channel.Name, err))
}
}
return otherInfo
@@ -208,7 +222,7 @@ func (channel *Channel) GetOtherInfo() map[string]interface{} {
func (channel *Channel) SetOtherInfo(otherInfo map[string]interface{}) {
otherInfoBytes, err := json.Marshal(otherInfo)
if err != nil {
common.SysError("failed to marshal other info: " + err.Error())
common.SysLog(fmt.Sprintf("failed to marshal other info: channel_id=%d, tag=%s, name=%s, error=%v", channel.Id, channel.GetTag(), channel.Name, err))
return
}
channel.OtherInfo = string(otherInfoBytes)
@@ -236,6 +250,10 @@ func (channel *Channel) Save() error {
return DB.Save(channel).Error
}
func (channel *Channel) SaveWithoutKey() error {
return DB.Omit("key").Save(channel).Error
}
func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool) ([]*Channel, error) {
var channels []*Channel
var err error
@@ -328,38 +346,54 @@ func GetChannelById(id int, selectAll bool) (*Channel, error) {
}
func BatchInsertChannels(channels []Channel) error {
var err error
err = DB.Create(&channels).Error
if err != nil {
return err
if len(channels) == 0 {
return nil
}
for _, channel_ := range channels {
err = channel_.AddAbilities()
if err != nil {
tx := DB.Begin()
if tx.Error != nil {
return tx.Error
}
defer func() {
if r := recover(); r != nil {
tx.Rollback()
}
}()
for _, chunk := range lo.Chunk(channels, 50) {
if err := tx.Create(&chunk).Error; err != nil {
tx.Rollback()
return err
}
for _, channel_ := range chunk {
if err := channel_.AddAbilities(tx); err != nil {
tx.Rollback()
return err
}
}
}
return nil
return tx.Commit().Error
}
func BatchDeleteChannels(ids []int) error {
//使用事务 删除channel表和channel_ability表
if len(ids) == 0 {
return nil
}
// 使用事务 分批删除channel表和abilities表
tx := DB.Begin()
err := tx.Where("id in (?)", ids).Delete(&Channel{}).Error
if err != nil {
// 回滚事务
tx.Rollback()
return err
if tx.Error != nil {
return tx.Error
}
err = tx.Where("channel_id in (?)", ids).Delete(&Ability{}).Error
if err != nil {
// 回滚事务
tx.Rollback()
return err
for _, chunk := range lo.Chunk(ids, 200) {
if err := tx.Where("id in (?)", chunk).Delete(&Channel{}).Error; err != nil {
tx.Rollback()
return err
}
if err := tx.Where("channel_id in (?)", chunk).Delete(&Ability{}).Error; err != nil {
tx.Rollback()
return err
}
}
// 提交事务
tx.Commit()
return err
return tx.Commit().Error
}
func (channel *Channel) GetPriority() int64 {
@@ -380,7 +414,11 @@ func (channel *Channel) GetBaseURL() string {
if channel.BaseURL == nil {
return ""
}
return *channel.BaseURL
url := *channel.BaseURL
if url == "" {
url = constant.ChannelBaseURLs[channel.Type]
}
return url
}
func (channel *Channel) GetModelMapping() string {
@@ -403,7 +441,7 @@ func (channel *Channel) Insert() error {
if err != nil {
return err
}
err = channel.AddAbilities()
err = channel.AddAbilities(nil)
return err
}
@@ -425,7 +463,7 @@ func (channel *Channel) Update() error {
trimmed := strings.TrimSpace(keyStr)
if strings.HasPrefix(trimmed, "[") {
var arr []json.RawMessage
if err := json.Unmarshal([]byte(trimmed), &arr); err == nil {
if err := common.Unmarshal([]byte(trimmed), &arr); err == nil {
keys = make([]string, len(arr))
for i, v := range arr {
keys[i] = string(v)
@@ -462,7 +500,7 @@ func (channel *Channel) UpdateResponseTime(responseTime int64) {
ResponseTime: int(responseTime),
}).Error
if err != nil {
common.SysError("failed to update response time: " + err.Error())
common.SysLog(fmt.Sprintf("failed to update response time: channel_id=%d, error=%v", channel.Id, err))
}
}
@@ -472,7 +510,7 @@ func (channel *Channel) UpdateBalance(balance float64) {
Balance: balance,
}).Error
if err != nil {
common.SysError("failed to update balance: " + err.Error())
common.SysLog(fmt.Sprintf("failed to update balance: channel_id=%d, error=%v", channel.Id, err))
}
}
@@ -491,8 +529,8 @@ var channelStatusLock sync.Mutex
// channelPollingLocks stores locks for each channel.id to ensure thread-safe polling
var channelPollingLocks sync.Map
// getChannelPollingLock returns or creates a mutex for the given channel ID
func getChannelPollingLock(channelId int) *sync.Mutex {
// GetChannelPollingLock returns or creates a mutex for the given channel ID
func GetChannelPollingLock(channelId int) *sync.Mutex {
if lock, exists := channelPollingLocks.Load(channelId); exists {
return lock.(*sync.Mutex)
}
@@ -522,8 +560,8 @@ func CleanupChannelPollingLocks() {
})
}
func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int) {
keys := channel.getKeys()
func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason string) {
keys := channel.GetKeys()
if len(keys) == 0 {
channel.Status = status
} else {
@@ -541,6 +579,14 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int) {
delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex)
} else {
channel.ChannelInfo.MultiKeyStatusList[keyIndex] = status
if channel.ChannelInfo.MultiKeyDisabledReason == nil {
channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string)
}
if channel.ChannelInfo.MultiKeyDisabledTime == nil {
channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64)
}
channel.ChannelInfo.MultiKeyDisabledReason[keyIndex] = reason
channel.ChannelInfo.MultiKeyDisabledTime[keyIndex] = common.GetTimestamp()
}
if len(channel.ChannelInfo.MultiKeyStatusList) >= channel.ChannelInfo.MultiKeySize {
channel.Status = common.ChannelStatusAutoDisabled
@@ -562,8 +608,12 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
return false
}
if channelCache.ChannelInfo.IsMultiKey {
// Use per-channel lock to prevent concurrent map read/write with GetNextEnabledKey
pollingLock := GetChannelPollingLock(channelId)
pollingLock.Lock()
// 如果是多Key模式，更新缓存中的状态
handlerMultiKeyUpdate(channelCache, usingKey, status)
handlerMultiKeyUpdate(channelCache, usingKey, status, reason)
pollingLock.Unlock()
//CacheUpdateChannel(channelCache)
//return true
} else {
@@ -571,10 +621,6 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
if channelCache.Status == status {
return false
}
// 如果缓存渠道不存在(说明已经被禁用)，且要设置的状态不为启用，直接返回
if status != common.ChannelStatusEnabled {
return false
}
CacheUpdateChannelStatus(channelId, status)
}
}
@@ -584,7 +630,7 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
if shouldUpdateAbilities {
err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled)
if err != nil {
common.SysError("failed to update ability status: " + err.Error())
common.SysLog(fmt.Sprintf("failed to update ability status: channel_id=%d, error=%v", channelId, err))
}
}
}()
@@ -598,7 +644,11 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
if channel.ChannelInfo.IsMultiKey {
beforeStatus := channel.Status
handlerMultiKeyUpdate(channel, usingKey, status)
// Protect map writes with the same per-channel lock used by readers
pollingLock := GetChannelPollingLock(channelId)
pollingLock.Lock()
handlerMultiKeyUpdate(channel, usingKey, status, reason)
pollingLock.Unlock()
if beforeStatus != channel.Status {
shouldUpdateAbilities = true
}
@@ -610,9 +660,9 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
channel.Status = status
shouldUpdateAbilities = true
}
err = channel.Save()
err = channel.SaveWithoutKey()
if err != nil {
common.SysError("failed to update channel status: " + err.Error())
common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channel.Id, status, err))
return false
}
}
@@ -674,7 +724,7 @@ func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *
for _, channel := range channels {
err = channel.UpdateAbilities(nil)
if err != nil {
common.SysError("failed to update abilities: " + err.Error())
common.SysLog(fmt.Sprintf("failed to update abilities: channel_id=%d, tag=%s, error=%v", channel.Id, channel.GetTag(), err))
}
}
}
@@ -698,7 +748,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())
common.SysLog(fmt.Sprintf("failed to update channel used quota: channel_id=%d, delta_quota=%d, error=%v", id, quota, err))
}
}
@@ -778,7 +828,7 @@ func SearchTags(keyword string, group string, model string, idSort bool) ([]*str
func (channel *Channel) ValidateSettings() error {
channelParams := &dto.ChannelSettings{}
if channel.Setting != nil && *channel.Setting != "" {
err := json.Unmarshal([]byte(*channel.Setting), channelParams)
err := common.Unmarshal([]byte(*channel.Setting), channelParams)
if err != nil {
return err
}
@@ -789,9 +839,9 @@ func (channel *Channel) ValidateSettings() error {
func (channel *Channel) GetSetting() dto.ChannelSettings {
setting := dto.ChannelSettings{}
if channel.Setting != nil && *channel.Setting != "" {
err := json.Unmarshal([]byte(*channel.Setting), &setting)
err := common.Unmarshal([]byte(*channel.Setting), &setting)
if err != nil {
common.SysError("failed to unmarshal setting: " + err.Error())
common.SysLog(fmt.Sprintf("failed to unmarshal setting: channel_id=%d, error=%v", channel.Id, err))
channel.Setting = nil // 清空设置以避免后续错误
_ = channel.Save() // 保存修改
}
@@ -800,25 +850,58 @@ func (channel *Channel) GetSetting() dto.ChannelSettings {
}
func (channel *Channel) SetSetting(setting dto.ChannelSettings) {
settingBytes, err := json.Marshal(setting)
settingBytes, err := common.Marshal(setting)
if err != nil {
common.SysError("failed to marshal setting: " + err.Error())
common.SysLog(fmt.Sprintf("failed to marshal setting: channel_id=%d, error=%v", channel.Id, err))
return
}
channel.Setting = common.GetPointer[string](string(settingBytes))
}
func (channel *Channel) GetOtherSettings() dto.ChannelOtherSettings {
setting := dto.ChannelOtherSettings{}
if channel.OtherSettings != "" {
err := common.UnmarshalJsonStr(channel.OtherSettings, &setting)
if err != nil {
common.SysLog(fmt.Sprintf("failed to unmarshal setting: channel_id=%d, error=%v", channel.Id, err))
channel.OtherSettings = "{}" // 清空设置以避免后续错误
_ = channel.Save() // 保存修改
}
}
return setting
}
func (channel *Channel) SetOtherSettings(setting dto.ChannelOtherSettings) {
settingBytes, err := common.Marshal(setting)
if err != nil {
common.SysLog(fmt.Sprintf("failed to marshal setting: channel_id=%d, error=%v", channel.Id, err))
return
}
channel.OtherSettings = string(settingBytes)
}
func (channel *Channel) GetParamOverride() map[string]interface{} {
paramOverride := make(map[string]interface{})
if channel.ParamOverride != nil && *channel.ParamOverride != "" {
err := json.Unmarshal([]byte(*channel.ParamOverride), &paramOverride)
err := common.Unmarshal([]byte(*channel.ParamOverride), &paramOverride)
if err != nil {
common.SysError("failed to unmarshal param override: " + err.Error())
common.SysLog(fmt.Sprintf("failed to unmarshal param override: channel_id=%d, error=%v", channel.Id, err))
}
}
return paramOverride
}
func (channel *Channel) GetHeaderOverride() map[string]interface{} {
headerOverride := make(map[string]interface{})
if channel.HeaderOverride != nil && *channel.HeaderOverride != "" {
err := common.Unmarshal([]byte(*channel.HeaderOverride), &headerOverride)
if err != nil {
common.SysLog(fmt.Sprintf("failed to unmarshal header override: channel_id=%d, error=%v", channel.Id, err))
}
}
return headerOverride
}
func GetChannelsByIds(ids []int) ([]*Channel, error) {
var channels []*Channel
err := DB.Where("id in (?)", ids).Find(&channels).Error
model/channel_cache.go
+39 -17
View File
@@ -5,7 +5,9 @@ import (
"fmt"
"math/rand"
"one-api/common"
"one-api/constant"
"one-api/setting"
"one-api/setting/ratio_setting"
"sort"
"strings"
"sync"
@@ -66,6 +68,20 @@ func InitChannelCache() {
channelSyncLock.Lock()
group2model2channels = newGroup2model2channels
//channelsIDM = newChannelId2channel
for i, channel := range newChannelId2channel {
if channel.ChannelInfo.IsMultiKey {
channel.Keys = channel.GetKeys()
if channel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling {
if oldChannel, ok := channelsIDM[i]; ok {
// 存在旧的渠道，如果是多key且轮询，保留轮询索引信息
if oldChannel.ChannelInfo.IsMultiKey && oldChannel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling {
channel.ChannelInfo.MultiKeyPollingIndex = oldChannel.ChannelInfo.MultiKeyPollingIndex
}
}
}
}
}
channelsIDM = newChannelId2channel
channelSyncLock.Unlock()
common.SysLog("channels synced from database")
@@ -109,20 +125,10 @@ func CacheGetRandomSatisfiedChannel(c *gin.Context, group string, model string,
return nil, group, err
}
}
if channel == nil {
return nil, group, errors.New("channel not found")
}
return channel, selectGroup, nil
}
func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) {
if strings.HasPrefix(model, "gpt-4-gizmo") {
model = "gpt-4-gizmo-*"
}
if strings.HasPrefix(model, "gpt-4o-gizmo") {
model = "gpt-4o-gizmo-*"
}
// if memory cache is disabled, get channel directly from database
if !common.MemoryCacheEnabled {
return GetRandomSatisfiedChannel(group, model, retry)
@@ -130,10 +136,18 @@ func getRandomSatisfiedChannel(group string, model string, retry int) (*Channel,
channelSyncLock.RLock()
defer channelSyncLock.RUnlock()
// First, try to find channels with the exact model name.
channels := group2model2channels[group][model]
// If no channels found, try to find channels with the normalized model name.
if len(channels) == 0 {
return nil, errors.New("channel not found")
normalizedModel := ratio_setting.FormatMatchingModelName(model)
channels = group2model2channels[group][normalizedModel]
}
if len(channels) == 0 {
return nil, nil
}
if len(channels) == 1 {
@@ -206,9 +220,6 @@ func CacheGetChannel(id int) (*Channel, error) {
if !ok {
return nil, fmt.Errorf("渠道# %d，已不存在", id)
}
if c.Status != common.ChannelStatusEnabled {
return nil, fmt.Errorf("渠道# %d，已被禁用", id)
}
return c, nil
}
@@ -227,9 +238,6 @@ func CacheGetChannelInfo(id int) (*ChannelInfo, error) {
if !ok {
return nil, fmt.Errorf("渠道# %d，已不存在", id)
}
if c.Status != common.ChannelStatusEnabled {
return nil, fmt.Errorf("渠道# %d，已被禁用", id)
}
return &c.ChannelInfo, nil
}
@@ -242,6 +250,20 @@ func CacheUpdateChannelStatus(id int, status int) {
if channel, ok := channelsIDM[id]; ok {
channel.Status = status
}
if status != common.ChannelStatusEnabled {
// delete the channel from group2model2channels
for group, model2channels := range group2model2channels {
for model, channels := range model2channels {
for i, channelId := range channels {
if channelId == id {
// remove the channel from the slice
group2model2channels[group][model] = append(channels[:i], channels[i+1:]...)
break
}
}
}
}
}
}
func CacheUpdateChannel(channel *Channel) {
model/log.go
+12 -15
View File
@@ -4,6 +4,8 @@ import (
"context"
"fmt"
"one-api/common"
"one-api/logger"
"one-api/types"
"os"
"strings"
"time"
@@ -87,13 +89,13 @@ func RecordLog(userId int, logType int, content string) {
}
err := LOG_DB.Create(log).Error
if err != nil {
common.SysError("failed to record log: " + err.Error())
common.SysLog("failed to record log: " + err.Error())
}
}
func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string, tokenName string, content string, tokenId int, useTimeSeconds int,
isStream bool, group string, other map[string]interface{}) {
common.LogInfo(c, fmt.Sprintf("record error log: userId=%d, channelId=%d, modelName=%s, tokenName=%s, content=%s", userId, channelId, modelName, tokenName, content))
logger.LogInfo(c, fmt.Sprintf("record error log: userId=%d, channelId=%d, modelName=%s, tokenName=%s, content=%s", userId, channelId, modelName, tokenName, content))
username := c.GetString("username")
otherStr := common.MapToJsonStr(other)
// 判断是否需要记录 IP
@@ -129,7 +131,7 @@ func RecordErrorLog(c *gin.Context, userId int, channelId int, modelName string,
}
err := LOG_DB.Create(log).Error
if err != nil {
common.LogError(c, "failed to record log: "+err.Error())
logger.LogError(c, "failed to record log: "+err.Error())
}
}
@@ -142,7 +144,6 @@ type RecordConsumeLogParams struct {
Quota int `json:"quota"`
Content string `json:"content"`
TokenId int `json:"token_id"`
UserQuota int `json:"user_quota"`
UseTimeSeconds int `json:"use_time_seconds"`
IsStream bool `json:"is_stream"`
Group string `json:"group"`
@@ -150,10 +151,10 @@ type RecordConsumeLogParams struct {
}
func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams) {
common.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params)))
if !common.LogConsumeEnabled {
return
}
logger.LogInfo(c, fmt.Sprintf("record consume log: userId=%d, params=%s", userId, common.GetJsonString(params)))
username := c.GetString("username")
otherStr := common.MapToJsonStr(params.Other)
// 判断是否需要记录 IP
@@ -189,7 +190,7 @@ func RecordConsumeLog(c *gin.Context, userId int, params RecordConsumeLogParams)
}
err := LOG_DB.Create(log).Error
if err != nil {
common.LogError(c, "failed to record log: "+err.Error())
logger.LogError(c, "failed to record log: "+err.Error())
}
if common.DataExportEnabled {
gopool.Go(func() {
@@ -236,26 +237,22 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName
return nil, 0, err
}
channelIdsMap := make(map[int]struct{})
channelMap := make(map[int]string)
channelIds := types.NewSet[int]()
for _, log := range logs {
if log.ChannelId != 0 {
channelIdsMap[log.ChannelId] = struct{}{}
channelIds.Add(log.ChannelId)
}
}
channelIds := make([]int, 0, len(channelIdsMap))
for channelId := range channelIdsMap {
channelIds = append(channelIds, channelId)
}
if len(channelIds) > 0 {
if channelIds.Len() > 0 {
var channels []struct {
Id int `gorm:"column:id"`
Name string `gorm:"column:name"`
}
if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds).Find(&channels).Error; err != nil {
if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil {
return logs, total, err
}
channelMap := make(map[int]string, len(channels))
for _, channel := range channels {
channelMap[channel.Id] = channel.Name
}
model/main.go
+115 -3
View File
@@ -180,6 +180,12 @@ func InitDB() (err error) {
db = db.Debug()
}
DB = db
// MySQL charset/collation startup check: ensure Chinese-capable charset
if common.UsingMySQL {
if err := checkMySQLChineseSupport(DB); err != nil {
panic(err)
}
}
sqlDB, err := DB.DB()
if err != nil {
return err
@@ -214,6 +220,12 @@ func InitLogDB() (err error) {
db = db.Debug()
}
LOG_DB = db
// If log DB is MySQL, also ensure Chinese-capable charset
if common.LogSqlType == common.DatabaseTypeMySQL {
if err := checkMySQLChineseSupport(LOG_DB); err != nil {
panic(err)
}
}
sqlDB, err := LOG_DB.DB()
if err != nil {
return err
@@ -235,9 +247,6 @@ func InitLogDB() (err error) {
}
func migrateDB() error {
if !common.UsingPostgreSQL {
return migrateDBFast()
}
err := DB.AutoMigrate(
&Channel{},
&Token{},
@@ -250,7 +259,12 @@ func migrateDB() error {
&TopUp{},
&QuotaData{},
&Task{},
&Model{},
&Vendor{},
&PrefillGroup{},
&Setup{},
&TwoFA{},
&TwoFABackupCode{},
)
if err != nil {
return err
@@ -259,6 +273,7 @@ func migrateDB() error {
}
func migrateDBFast() error {
var wg sync.WaitGroup
migrations := []struct {
@@ -276,7 +291,12 @@ func migrateDBFast() error {
{&TopUp{}, "TopUp"},
{&QuotaData{}, "QuotaData"},
{&Task{}, "Task"},
{&Model{}, "Model"},
{&Vendor{}, "Vendor"},
{&PrefillGroup{}, "PrefillGroup"},
{&Setup{}, "Setup"},
{&TwoFA{}, "TwoFA"},
{&TwoFABackupCode{}, "TwoFABackupCode"},
}
// 动态计算migration数量，确保errChan缓冲区足够大
errChan := make(chan error, len(migrations))
@@ -332,6 +352,98 @@ func CloseDB() error {
return closeDB(DB)
}
// checkMySQLChineseSupport ensures the MySQL connection and current schema
// default charset/collation can store Chinese characters. It allows common
// Chinese-capable charsets (utf8mb4, utf8, gbk, big5, gb18030) and panics otherwise.
func checkMySQLChineseSupport(db *gorm.DB) error {
// 仅检测：当前库默认字符集/排序规则 + 各表的排序规则（隐含字符集）
// Read current schema defaults
var schemaCharset, schemaCollation string
err := db.Raw("SELECT DEFAULT_CHARACTER_SET_NAME, DEFAULT_COLLATION_NAME FROM information_schema.SCHEMATA WHERE SCHEMA_NAME = DATABASE()").Row().Scan(&schemaCharset, &schemaCollation)
if err != nil {
return fmt.Errorf("读取当前库默认字符集/排序规则失败 / Failed to read schema default charset/collation: %v", err)
}
toLower := func(s string) string { return strings.ToLower(s) }
// Allowed charsets that can store Chinese text
allowedCharsets := map[string]string{
"utf8mb4": "utf8mb4_",
"utf8": "utf8_",
"gbk": "gbk_",
"big5": "big5_",
"gb18030": "gb18030_",
}
isChineseCapable := func(cs, cl string) bool {
csLower := toLower(cs)
clLower := toLower(cl)
if prefix, ok := allowedCharsets[csLower]; ok {
if clLower == "" {
return true
}
return strings.HasPrefix(clLower, prefix)
}
// 如果仅提供了排序规则，尝试按排序规则前缀判断
for _, prefix := range allowedCharsets {
if strings.HasPrefix(clLower, prefix) {
return true
}
}
return false
}
// 1) 当前库默认值必须支持中文
if !isChineseCapable(schemaCharset, schemaCollation) {
return fmt.Errorf("当前库默认字符集/排序规则不支持中文：schema(%s/%s)。请将库设置为 utf8mb4/utf8/gbk/big5/gb18030 / Schema default charset/collation is not Chinese-capable: schema(%s/%s). Please set to utf8mb4/utf8/gbk/big5/gb18030",
schemaCharset, schemaCollation, schemaCharset, schemaCollation)
}
// 2) 所有物理表的排序规则（隐含字符集）必须支持中文
type tableInfo struct {
Name string
Collation *string
}
var tables []tableInfo
if err := db.Raw("SELECT TABLE_NAME, TABLE_COLLATION FROM information_schema.TABLES WHERE TABLE_SCHEMA = DATABASE() AND TABLE_TYPE = 'BASE TABLE'").Scan(&tables).Error; err != nil {
return fmt.Errorf("读取表排序规则失败 / Failed to read table collations: %v", err)
}
var badTables []string
for _, t := range tables {
// NULL 或空表示继承库默认设置，已在上面校验库默认，视为通过
if t.Collation == nil || *t.Collation == "" {
continue
}
cl := *t.Collation
// 仅凭排序规则判断是否中文可用
ok := false
lower := strings.ToLower(cl)
for _, prefix := range allowedCharsets {
if strings.HasPrefix(lower, prefix) {
ok = true
break
}
}
if !ok {
badTables = append(badTables, fmt.Sprintf("%s(%s)", t.Name, cl))
}
}
if len(badTables) > 0 {
// 限制输出数量以避免日志过长
maxShow := 20
shown := badTables
if len(shown) > maxShow {
shown = shown[:maxShow]
}
return fmt.Errorf(
"存在不支持中文的表，请修复其排序规则/字符集。示例（最多展示 %d 项）：%v / Found tables not Chinese-capable. Please fix their collation/charset. Examples (showing up to %d): %v",
maxShow, shown, maxShow, shown,
)
}
return nil
}
var (
lastPingTime time.Time
pingMutex sync.Mutex
model/missing_models.go
+30
View File
@@ -0,0 +1,30 @@
package model
// GetMissingModels returns model names that are referenced in the system
func GetMissingModels() ([]string, error) {
// 1. 获取所有已启用模型（去重）
models := GetEnabledModels()
if len(models) == 0 {
return []string{}, nil
}
// 2. 查询已有的元数据模型名
var existing []string
if err := DB.Model(&Model{}).Where("model_name IN ?", models).Pluck("model_name", &existing).Error; err != nil {
return nil, err
}
existingSet := make(map[string]struct{}, len(existing))
for _, e := range existing {
existingSet[e] = struct{}{}
}
// 3. 收集缺失模型
var missing []string
for _, name := range models {
if _, ok := existingSet[name]; !ok {
missing = append(missing, name)
}
}
return missing, nil
}
model/model_extra.go
+31
View File
@@ -0,0 +1,31 @@
package model
func GetModelEnableGroups(modelName string) []string {
// 确保缓存最新
GetPricing()
if modelName == "" {
return make([]string, 0)
}
modelEnableGroupsLock.RLock()
groups, ok := modelEnableGroups[modelName]
modelEnableGroupsLock.RUnlock()
if !ok {
return make([]string, 0)
}
return groups
}
// GetModelQuotaTypes 返回指定模型的计费类型集合（来自缓存）
func GetModelQuotaTypes(modelName string) []int {
GetPricing()
modelEnableGroupsLock.RLock()
quota, ok := modelQuotaTypeMap[modelName]
modelEnableGroupsLock.RUnlock()
if !ok {
return []int{}
}
return []int{quota}
}
model/model_meta.go
+147
View File
@@ -0,0 +1,147 @@
package model
import (
"one-api/common"
"strconv"
"gorm.io/gorm"
)
const (
NameRuleExact = iota
NameRulePrefix
NameRuleContains
NameRuleSuffix
)
type BoundChannel struct {
Name string `json:"name"`
Type int `json:"type"`
}
type Model struct {
Id int `json:"id"`
ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name_delete_at,priority:1"`
Description string `json:"description,omitempty" gorm:"type:text"`
Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"`
Tags string `json:"tags,omitempty" gorm:"type:varchar(255)"`
VendorID int `json:"vendor_id,omitempty" gorm:"index"`
Endpoints string `json:"endpoints,omitempty" gorm:"type:text"`
Status int `json:"status" gorm:"default:1"`
SyncOfficial int `json:"sync_official" gorm:"default:1"`
CreatedTime int64 `json:"created_time" gorm:"bigint"`
UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_model_name_delete_at,priority:2"`
BoundChannels []BoundChannel `json:"bound_channels,omitempty" gorm:"-"`
EnableGroups []string `json:"enable_groups,omitempty" gorm:"-"`
QuotaTypes []int `json:"quota_types,omitempty" gorm:"-"`
NameRule int `json:"name_rule" gorm:"default:0"`
MatchedModels []string `json:"matched_models,omitempty" gorm:"-"`
MatchedCount int `json:"matched_count,omitempty" gorm:"-"`
}
func (mi *Model) Insert() error {
now := common.GetTimestamp()
mi.CreatedTime = now
mi.UpdatedTime = now
return DB.Create(mi).Error
}
func IsModelNameDuplicated(id int, name string) (bool, error) {
if name == "" {
return false, nil
}
var cnt int64
err := DB.Model(&Model{}).Where("model_name = ? AND id <> ?", name, id).Count(&cnt).Error
return cnt > 0, err
}
func (mi *Model) Update() error {
mi.UpdatedTime = common.GetTimestamp()
return DB.Session(&gorm.Session{AllowGlobalUpdate: false, FullSaveAssociations: false}).
Model(&Model{}).
Where("id = ?", mi.Id).
Omit("created_time").
Select("*").
Updates(mi).Error
}
func (mi *Model) Delete() error {
return DB.Delete(mi).Error
}
func GetVendorModelCounts() (map[int64]int64, error) {
var stats []struct {
VendorID int64
Count int64
}
if err := DB.Model(&Model{}).
Select("vendor_id as vendor_id, count(*) as count").
Group("vendor_id").
Scan(&stats).Error; err != nil {
return nil, err
}
m := make(map[int64]int64, len(stats))
for _, s := range stats {
m[s.VendorID] = s.Count
}
return m, nil
}
func GetAllModels(offset int, limit int) ([]*Model, error) {
var models []*Model
err := DB.Order("id DESC").Offset(offset).Limit(limit).Find(&models).Error
return models, err
}
func GetBoundChannelsByModelsMap(modelNames []string) (map[string][]BoundChannel, error) {
result := make(map[string][]BoundChannel)
if len(modelNames) == 0 {
return result, nil
}
type row struct {
Model string
Name string
Type int
}
var rows []row
err := DB.Table("channels").
Select("abilities.model as model, channels.name as name, channels.type as type").
Joins("JOIN abilities ON abilities.channel_id = channels.id").
Where("abilities.model IN ? AND abilities.enabled = ?", modelNames, true).
Distinct().
Scan(&rows).Error
if err != nil {
return nil, err
}
for _, r := range rows {
result[r.Model] = append(result[r.Model], BoundChannel{Name: r.Name, Type: r.Type})
}
return result, nil
}
func SearchModels(keyword string, vendor string, offset int, limit int) ([]*Model, int64, error) {
var models []*Model
db := DB.Model(&Model{})
if keyword != "" {
like := "%" + keyword + "%"
db = db.Where("model_name LIKE ? OR description LIKE ? OR tags LIKE ?", like, like, like)
}
if vendor != "" {
if vid, err := strconv.Atoi(vendor); err == nil {
db = db.Where("models.vendor_id = ?", vid)
} else {
db = db.Joins("JOIN vendors ON vendors.id = models.vendor_id").Where("vendors.name LIKE ?", "%"+vendor+"%")
}
}
var total int64
if err := db.Count(&total).Error; err != nil {
return nil, 0, err
}
if err := db.Order("models.id DESC").Offset(offset).Limit(limit).Find(&models).Error; err != nil {
return nil, 0, err
}
return models, total, nil
}
model/option.go
+32 -20
View File
@@ -6,6 +6,7 @@ import (
"one-api/setting/config"
"one-api/setting/operation_setting"
"one-api/setting/ratio_setting"
"one-api/setting/system_setting"
"strconv"
"strings"
"time"
@@ -66,16 +67,16 @@ func InitOptionMap() {
common.OptionMap["SystemName"] = common.SystemName
common.OptionMap["Logo"] = common.Logo
common.OptionMap["ServerAddress"] = ""
common.OptionMap["WorkerUrl"] = setting.WorkerUrl
common.OptionMap["WorkerValidKey"] = setting.WorkerValidKey
common.OptionMap["WorkerAllowHttpImageRequestEnabled"] = strconv.FormatBool(setting.WorkerAllowHttpImageRequestEnabled)
common.OptionMap["WorkerUrl"] = system_setting.WorkerUrl
common.OptionMap["WorkerValidKey"] = system_setting.WorkerValidKey
common.OptionMap["WorkerAllowHttpImageRequestEnabled"] = strconv.FormatBool(system_setting.WorkerAllowHttpImageRequestEnabled)
common.OptionMap["PayAddress"] = ""
common.OptionMap["CustomCallbackAddress"] = ""
common.OptionMap["EpayId"] = ""
common.OptionMap["EpayKey"] = ""
common.OptionMap["Price"] = strconv.FormatFloat(setting.Price, 'f', -1, 64)
common.OptionMap["USDExchangeRate"] = strconv.FormatFloat(setting.USDExchangeRate, 'f', -1, 64)
common.OptionMap["MinTopUp"] = strconv.Itoa(setting.MinTopUp)
common.OptionMap["Price"] = strconv.FormatFloat(operation_setting.Price, 'f', -1, 64)
common.OptionMap["USDExchangeRate"] = strconv.FormatFloat(operation_setting.USDExchangeRate, 'f', -1, 64)
common.OptionMap["MinTopUp"] = strconv.Itoa(operation_setting.MinTopUp)
common.OptionMap["StripeMinTopUp"] = strconv.Itoa(setting.StripeMinTopUp)
common.OptionMap["StripeApiSecret"] = setting.StripeApiSecret
common.OptionMap["StripeWebhookSecret"] = setting.StripeWebhookSecret
@@ -89,7 +90,7 @@ func InitOptionMap() {
common.OptionMap["Chats"] = setting.Chats2JsonString()
common.OptionMap["AutoGroups"] = setting.AutoGroups2JsonString()
common.OptionMap["DefaultUseAutoGroup"] = strconv.FormatBool(setting.DefaultUseAutoGroup)
common.OptionMap["PayMethods"] = setting.PayMethods2JsonString()
common.OptionMap["PayMethods"] = operation_setting.PayMethods2JsonString()
common.OptionMap["GitHubClientId"] = ""
common.OptionMap["GitHubClientSecret"] = ""
common.OptionMap["TelegramBotToken"] = ""
@@ -115,6 +116,9 @@ func InitOptionMap() {
common.OptionMap["GroupGroupRatio"] = ratio_setting.GroupGroupRatio2JSONString()
common.OptionMap["UserUsableGroups"] = setting.UserUsableGroups2JSONString()
common.OptionMap["CompletionRatio"] = ratio_setting.CompletionRatio2JSONString()
common.OptionMap["ImageRatio"] = ratio_setting.ImageRatio2JSONString()
common.OptionMap["AudioRatio"] = ratio_setting.AudioRatio2JSONString()
common.OptionMap["AudioCompletionRatio"] = ratio_setting.AudioCompletionRatio2JSONString()
common.OptionMap["TopUpLink"] = common.TopUpLink
//common.OptionMap["ChatLink"] = common.ChatLink
//common.OptionMap["ChatLink2"] = common.ChatLink2
@@ -154,7 +158,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())
common.SysLog("failed to update option map: " + err.Error())
}
}
}
@@ -275,7 +279,7 @@ func updateOptionMap(key string, value string) (err error) {
case "SMTPSSLEnabled":
common.SMTPSSLEnabled = boolValue
case "WorkerAllowHttpImageRequestEnabled":
setting.WorkerAllowHttpImageRequestEnabled = boolValue
system_setting.WorkerAllowHttpImageRequestEnabled = boolValue
case "DefaultUseAutoGroup":
setting.DefaultUseAutoGroup = boolValue
case "ExposeRatioEnabled":
@@ -297,29 +301,29 @@ func updateOptionMap(key string, value string) (err error) {
case "SMTPToken":
common.SMTPToken = value
case "ServerAddress":
setting.ServerAddress = value
system_setting.ServerAddress = value
case "WorkerUrl":
setting.WorkerUrl = value
system_setting.WorkerUrl = value
case "WorkerValidKey":
setting.WorkerValidKey = value
system_setting.WorkerValidKey = value
case "PayAddress":
setting.PayAddress = value
operation_setting.PayAddress = value
case "Chats":
err = setting.UpdateChatsByJsonString(value)
case "AutoGroups":
err = setting.UpdateAutoGroupsByJsonString(value)
case "CustomCallbackAddress":
setting.CustomCallbackAddress = value
operation_setting.CustomCallbackAddress = value
case "EpayId":
setting.EpayId = value
operation_setting.EpayId = value
case "EpayKey":
setting.EpayKey = value
operation_setting.EpayKey = value
case "Price":
setting.Price, _ = strconv.ParseFloat(value, 64)
operation_setting.Price, _ = strconv.ParseFloat(value, 64)
case "USDExchangeRate":
setting.USDExchangeRate, _ = strconv.ParseFloat(value, 64)
operation_setting.USDExchangeRate, _ = strconv.ParseFloat(value, 64)
case "MinTopUp":
setting.MinTopUp, _ = strconv.Atoi(value)
operation_setting.MinTopUp, _ = strconv.Atoi(value)
case "StripeApiSecret":
setting.StripeApiSecret = value
case "StripeWebhookSecret":
@@ -348,6 +352,8 @@ func updateOptionMap(key string, value string) (err error) {
common.LinuxDOClientId = value
case "LinuxDOClientSecret":
common.LinuxDOClientSecret = value
case "LinuxDOMinimumTrustLevel":
common.LinuxDOMinimumTrustLevel, _ = strconv.Atoi(value)
case "Footer":
common.Footer = value
case "SystemName":
@@ -406,6 +412,12 @@ func updateOptionMap(key string, value string) (err error) {
err = ratio_setting.UpdateModelPriceByJSONString(value)
case "CacheRatio":
err = ratio_setting.UpdateCacheRatioByJSONString(value)
case "ImageRatio":
err = ratio_setting.UpdateImageRatioByJSONString(value)
case "AudioRatio":
err = ratio_setting.UpdateAudioRatioByJSONString(value)
case "AudioCompletionRatio":
err = ratio_setting.UpdateAudioCompletionRatioByJSONString(value)
case "TopUpLink":
common.TopUpLink = value
//case "ChatLink":
@@ -423,7 +435,7 @@ func updateOptionMap(key string, value string) (err error) {
case "StreamCacheQueueLength":
setting.StreamCacheQueueLength, _ = strconv.Atoi(value)
case "PayMethods":
err = setting.UpdatePayMethodsByJsonString(value)
err = operation_setting.UpdatePayMethodsByJsonString(value)
}
return err
}
model/prefill_group.go
+126
View File
@@ -0,0 +1,126 @@
package model
import (
"database/sql/driver"
"encoding/json"
"one-api/common"
"gorm.io/gorm"
)
// PrefillGroup 用于存储可复用的“组”信息，例如模型组、标签组、端点组等。
// Name 字段保持唯一，用于在前端下拉框中展示。
// Type 字段用于区分组的类别，可选值如：model、tag、endpoint。
// Items 字段使用 JSON 数组保存对应类型的字符串集合，示例：
// ["gpt-4o", "gpt-3.5-turbo"]
// 设计遵循 3NF，避免冗余，提供灵活扩展能力。
// JSONValue 基于 json.RawMessage 实现，支持从数据库的 []byte 和 string 两种类型读取
type JSONValue json.RawMessage
// Value 实现 driver.Valuer 接口，用于数据库写入
func (j JSONValue) Value() (driver.Value, error) {
if j == nil {
return nil, nil
}
return []byte(j), nil
}
// Scan 实现 sql.Scanner 接口，兼容不同驱动返回的类型
func (j *JSONValue) Scan(value interface{}) error {
switch v := value.(type) {
case nil:
*j = nil
return nil
case []byte:
// 拷贝底层字节，避免保留底层缓冲区
b := make([]byte, len(v))
copy(b, v)
*j = JSONValue(b)
return nil
case string:
*j = JSONValue([]byte(v))
return nil
default:
// 其他类型尝试序列化为 JSON
b, err := json.Marshal(v)
if err != nil {
return err
}
*j = JSONValue(b)
return nil
}
}
// MarshalJSON 确保在对外编码时与 json.RawMessage 行为一致
func (j JSONValue) MarshalJSON() ([]byte, error) {
if j == nil {
return []byte("null"), nil
}
return j, nil
}
// UnmarshalJSON 确保在对外解码时与 json.RawMessage 行为一致
func (j *JSONValue) UnmarshalJSON(data []byte) error {
if data == nil {
*j = nil
return nil
}
b := make([]byte, len(data))
copy(b, data)
*j = JSONValue(b)
return nil
}
type PrefillGroup struct {
Id int `json:"id"`
Name string `json:"name" gorm:"size:64;not null;uniqueIndex:uk_prefill_name,where:deleted_at IS NULL"`
Type string `json:"type" gorm:"size:32;index;not null"`
Items JSONValue `json:"items" gorm:"type:json"`
Description string `json:"description,omitempty" gorm:"type:varchar(255)"`
CreatedTime int64 `json:"created_time" gorm:"bigint"`
UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
}
// Insert 新建组
func (g *PrefillGroup) Insert() error {
now := common.GetTimestamp()
g.CreatedTime = now
g.UpdatedTime = now
return DB.Create(g).Error
}
// IsPrefillGroupNameDuplicated 检查组名称是否重复（排除自身 ID）
func IsPrefillGroupNameDuplicated(id int, name string) (bool, error) {
if name == "" {
return false, nil
}
var cnt int64
err := DB.Model(&PrefillGroup{}).Where("name = ? AND id <> ?", name, id).Count(&cnt).Error
return cnt > 0, err
}
// Update 更新组
func (g *PrefillGroup) Update() error {
g.UpdatedTime = common.GetTimestamp()
return DB.Save(g).Error
}
// DeleteByID 根据 ID 删除组
func DeletePrefillGroupByID(id int) error {
return DB.Delete(&PrefillGroup{}, id).Error
}
// GetAllPrefillGroups 获取全部组，可按类型过滤（为空则返回全部）
func GetAllPrefillGroups(groupType string) ([]*PrefillGroup, error) {
var groups []*PrefillGroup
query := DB.Model(&PrefillGroup{})
if groupType != "" {
query = query.Where("type = ?", groupType)
}
if err := query.Order("updated_time DESC").Find(&groups).Error; err != nil {
return nil, err
}
return groups, nil
}
model/pricing.go
+194 -9
View File
@@ -1,7 +1,10 @@
package model
import (
"encoding/json"
"fmt"
"strings"
"one-api/common"
"one-api/constant"
"one-api/setting/ratio_setting"
@@ -12,6 +15,10 @@ import (
type Pricing struct {
ModelName string `json:"model_name"`
Description string `json:"description,omitempty"`
Icon string `json:"icon,omitempty"`
Tags string `json:"tags,omitempty"`
VendorID int `json:"vendor_id,omitempty"`
QuotaType int `json:"quota_type"`
ModelRatio float64 `json:"model_ratio"`
ModelPrice float64 `json:"model_price"`
@@ -21,10 +28,24 @@ type Pricing struct {
SupportedEndpointTypes []constant.EndpointType `json:"supported_endpoint_types"`
}
type PricingVendor struct {
ID int `json:"id"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
Icon string `json:"icon,omitempty"`
}
var (
pricingMap []Pricing
lastGetPricingTime time.Time
updatePricingLock sync.Mutex
pricingMap []Pricing
vendorsList []PricingVendor
supportedEndpointMap map[string]common.EndpointInfo
lastGetPricingTime time.Time
updatePricingLock sync.Mutex
// 缓存映射：模型名 -> 启用分组 / 计费类型
modelEnableGroups = make(map[string][]string)
modelQuotaTypeMap = make(map[string]int)
modelEnableGroupsLock = sync.RWMutex{}
)
var (
@@ -46,6 +67,15 @@ func GetPricing() []Pricing {
return pricingMap
}
// GetVendors 返回当前定价接口使用到的供应商信息
func GetVendors() []PricingVendor {
if time.Since(lastGetPricingTime) > time.Minute*1 || len(pricingMap) == 0 {
// 保证先刷新一次
GetPricing()
}
return vendorsList
}
func GetModelSupportEndpointTypes(model string) []constant.EndpointType {
if model == "" {
return make([]constant.EndpointType, 0)
@@ -62,9 +92,83 @@ func updatePricing() {
//modelRatios := common.GetModelRatios()
enableAbilities, err := GetAllEnableAbilityWithChannels()
if err != nil {
common.SysError(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
common.SysLog(fmt.Sprintf("GetAllEnableAbilityWithChannels error: %v", err))
return
}
// 预加载模型元数据与供应商一次，避免循环查询
var allMeta []Model
_ = DB.Find(&allMeta).Error
metaMap := make(map[string]*Model)
prefixList := make([]*Model, 0)
suffixList := make([]*Model, 0)
containsList := make([]*Model, 0)
for i := range allMeta {
m := &allMeta[i]
if m.NameRule == NameRuleExact {
metaMap[m.ModelName] = m
} else {
switch m.NameRule {
case NameRulePrefix:
prefixList = append(prefixList, m)
case NameRuleSuffix:
suffixList = append(suffixList, m)
case NameRuleContains:
containsList = append(containsList, m)
}
}
}
// 将非精确规则模型匹配到 metaMap
for _, m := range prefixList {
for _, pricingModel := range enableAbilities {
if strings.HasPrefix(pricingModel.Model, m.ModelName) {
if _, exists := metaMap[pricingModel.Model]; !exists {
metaMap[pricingModel.Model] = m
}
}
}
}
for _, m := range suffixList {
for _, pricingModel := range enableAbilities {
if strings.HasSuffix(pricingModel.Model, m.ModelName) {
if _, exists := metaMap[pricingModel.Model]; !exists {
metaMap[pricingModel.Model] = m
}
}
}
}
for _, m := range containsList {
for _, pricingModel := range enableAbilities {
if strings.Contains(pricingModel.Model, m.ModelName) {
if _, exists := metaMap[pricingModel.Model]; !exists {
metaMap[pricingModel.Model] = m
}
}
}
}
// 预加载供应商
var vendors []Vendor
_ = DB.Find(&vendors).Error
vendorMap := make(map[int]*Vendor)
for i := range vendors {
vendorMap[vendors[i].Id] = &vendors[i]
}
// 初始化默认供应商映射
initDefaultVendorMapping(metaMap, vendorMap, enableAbilities)
// 构建对前端友好的供应商列表
vendorsList = make([]PricingVendor, 0, len(vendorMap))
for _, v := range vendorMap {
vendorsList = append(vendorsList, PricingVendor{
ID: v.Id,
Name: v.Name,
Description: v.Description,
Icon: v.Icon,
})
}
modelGroupsMap := make(map[string]*types.Set[string])
for _, ability := range enableAbilities {
@@ -79,12 +183,9 @@ func updatePricing() {
//这里使用切片而不是Set，因为一个模型可能支持多个端点类型，并且第一个端点是优先使用端点
modelSupportEndpointsStr := make(map[string][]string)
// 先根据已有能力填充原生端点
for _, ability := range enableAbilities {
endpoints, ok := modelSupportEndpointsStr[ability.Model]
if !ok {
endpoints = make([]string, 0)
modelSupportEndpointsStr[ability.Model] = endpoints
}
endpoints := modelSupportEndpointsStr[ability.Model]
channelTypes := common.GetEndpointTypesByChannelType(ability.ChannelType, ability.Model)
for _, channelType := range channelTypes {
if !common.StringsContains(endpoints, string(channelType)) {
@@ -94,6 +195,23 @@ func updatePricing() {
modelSupportEndpointsStr[ability.Model] = endpoints
}
// 再补充模型自定义端点
for modelName, meta := range metaMap {
if strings.TrimSpace(meta.Endpoints) == "" {
continue
}
var raw map[string]interface{}
if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
endpoints := modelSupportEndpointsStr[modelName]
for k := range raw {
if !common.StringsContains(endpoints, k) {
endpoints = append(endpoints, k)
}
}
modelSupportEndpointsStr[modelName] = endpoints
}
}
modelSupportEndpointTypes = make(map[string][]constant.EndpointType)
for model, endpoints := range modelSupportEndpointsStr {
supportedEndpoints := make([]constant.EndpointType, 0)
@@ -104,6 +222,45 @@ func updatePricing() {
modelSupportEndpointTypes[model] = supportedEndpoints
}
// 构建全局 supportedEndpointMap（默认 + 自定义覆盖）
supportedEndpointMap = make(map[string]common.EndpointInfo)
// 1. 默认端点
for _, endpoints := range modelSupportEndpointTypes {
for _, et := range endpoints {
if info, ok := common.GetDefaultEndpointInfo(et); ok {
if _, exists := supportedEndpointMap[string(et)]; !exists {
supportedEndpointMap[string(et)] = info
}
}
}
}
// 2. 自定义端点（models 表）覆盖默认
for _, meta := range metaMap {
if strings.TrimSpace(meta.Endpoints) == "" {
continue
}
var raw map[string]interface{}
if err := json.Unmarshal([]byte(meta.Endpoints), &raw); err == nil {
for k, v := range raw {
switch val := v.(type) {
case string:
supportedEndpointMap[k] = common.EndpointInfo{Path: val, Method: "POST"}
case map[string]interface{}:
ep := common.EndpointInfo{Method: "POST"}
if p, ok := val["path"].(string); ok {
ep.Path = p
}
if m, ok := val["method"].(string); ok {
ep.Method = strings.ToUpper(m)
}
supportedEndpointMap[k] = ep
default:
// ignore unsupported types
}
}
}
}
pricingMap = make([]Pricing, 0)
for model, groups := range modelGroupsMap {
pricing := Pricing{
@@ -111,6 +268,18 @@ func updatePricing() {
EnableGroup: groups.Items(),
SupportedEndpointTypes: modelSupportEndpointTypes[model],
}
// 补充模型元数据（描述、标签、供应商、状态）
if meta, ok := metaMap[model]; ok {
// 若模型被禁用(status!=1)，则直接跳过，不返回给前端
if meta.Status != 1 {
continue
}
pricing.Description = meta.Description
pricing.Icon = meta.Icon
pricing.Tags = meta.Tags
pricing.VendorID = meta.VendorID
}
modelPrice, findPrice := ratio_setting.GetModelPrice(model, false)
if findPrice {
pricing.ModelPrice = modelPrice
@@ -123,5 +292,21 @@ func updatePricing() {
}
pricingMap = append(pricingMap, pricing)
}
// 刷新缓存映射，供高并发快速查询
modelEnableGroupsLock.Lock()
modelEnableGroups = make(map[string][]string)
modelQuotaTypeMap = make(map[string]int)
for _, p := range pricingMap {
modelEnableGroups[p.ModelName] = p.EnableGroup
modelQuotaTypeMap[p.ModelName] = p.QuotaType
}
modelEnableGroupsLock.Unlock()
lastGetPricingTime = time.Now()
}
// GetSupportedEndpointMap 返回全局端点到路径的映射
func GetSupportedEndpointMap() map[string]common.EndpointInfo {
return supportedEndpointMap
}
model/pricing_default.go
+128
View File
@@ -0,0 +1,128 @@
package model
import (
"strings"
)
// 简化的供应商映射规则
var defaultVendorRules = map[string]string{
"gpt": "OpenAI",
"dall-e": "OpenAI",
"whisper": "OpenAI",
"o1": "OpenAI",
"o3": "OpenAI",
"claude": "Anthropic",
"gemini": "Google",
"moonshot": "Moonshot",
"kimi": "Moonshot",
"chatglm": "智谱",
"glm-": "智谱",
"qwen": "阿里巴巴",
"deepseek": "DeepSeek",
"abab": "MiniMax",
"ernie": "百度",
"spark": "讯飞",
"hunyuan": "腾讯",
"command": "Cohere",
"@cf/": "Cloudflare",
"360": "360",
"yi": "零一万物",
"jina": "Jina",
"mistral": "Mistral",
"grok": "xAI",
"llama": "Meta",
"doubao": "字节跳动",
"kling": "快手",
"jimeng": "即梦",
"vidu": "Vidu",
}
// 供应商默认图标映射
var defaultVendorIcons = map[string]string{
"OpenAI": "OpenAI",
"Anthropic": "Claude.Color",
"Google": "Gemini.Color",
"Moonshot": "Moonshot",
"智谱": "Zhipu.Color",
"阿里巴巴": "Qwen.Color",
"DeepSeek": "DeepSeek.Color",
"MiniMax": "Minimax.Color",
"百度": "Wenxin.Color",
"讯飞": "Spark.Color",
"腾讯": "Hunyuan.Color",
"Cohere": "Cohere.Color",
"Cloudflare": "Cloudflare.Color",
"360": "Ai360.Color",
"零一万物": "Yi.Color",
"Jina": "Jina",
"Mistral": "Mistral.Color",
"xAI": "XAI",
"Meta": "Ollama",
"字节跳动": "Doubao.Color",
"快手": "Kling.Color",
"即梦": "Jimeng.Color",
"Vidu": "Vidu",
"微软": "AzureAI",
"Microsoft": "AzureAI",
"Azure": "AzureAI",
}
// initDefaultVendorMapping 简化的默认供应商映射
func initDefaultVendorMapping(metaMap map[string]*Model, vendorMap map[int]*Vendor, enableAbilities []AbilityWithChannel) {
for _, ability := range enableAbilities {
modelName := ability.Model
if _, exists := metaMap[modelName]; exists {
continue
}
// 匹配供应商
vendorID := 0
modelLower := strings.ToLower(modelName)
for pattern, vendorName := range defaultVendorRules {
if strings.Contains(modelLower, pattern) {
vendorID = getOrCreateVendor(vendorName, vendorMap)
break
}
}
// 创建模型元数据
metaMap[modelName] = &Model{
ModelName: modelName,
VendorID: vendorID,
Status: 1,
NameRule: NameRuleExact,
}
}
}
// 查找或创建供应商
func getOrCreateVendor(vendorName string, vendorMap map[int]*Vendor) int {
// 查找现有供应商
for id, vendor := range vendorMap {
if vendor.Name == vendorName {
return id
}
}
// 创建新供应商
newVendor := &Vendor{
Name: vendorName,
Status: 1,
Icon: getDefaultVendorIcon(vendorName),
}
if err := newVendor.Insert(); err != nil {
return 0
}
vendorMap[newVendor.Id] = newVendor
return newVendor.Id
}
// 获取供应商默认图标
func getDefaultVendorIcon(vendorName string) string {
if icon, exists := defaultVendorIcons[vendorName]; exists {
return icon
}
return ""
}
model/pricing_refresh.go
+14
View File
@@ -0,0 +1,14 @@
package model
// RefreshPricing 强制立即重新计算与定价相关的缓存。
// 该方法用于需要最新数据的内部管理 API，
// 因此会绕过默认的 1 分钟延迟刷新。
func RefreshPricing() {
updatePricingLock.Lock()
defer updatePricingLock.Unlock()
modelSupportEndpointsLock.Lock()
defer modelSupportEndpointsLock.Unlock()
updatePricing()
}
model/redemption.go
+2 -1
View File
@@ -4,6 +4,7 @@ import (
"errors"
"fmt"
"one-api/common"
"one-api/logger"
"strconv"
"gorm.io/gorm"
@@ -148,7 +149,7 @@ func Redeem(key string, userId int) (quota int, err error) {
if err != nil {
return 0, errors.New("兑换失败，" + err.Error())
}
RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s，兑换码ID %d", common.LogQuota(redemption.Quota), redemption.Id))
RecordLog(userId, LogTypeTopup, fmt.Sprintf("通过兑换码充值 %s，兑换码ID %d", logger.LogQuota(redemption.Quota), redemption.Id))
return redemption.Quota, nil
}
model/task.go
+1 -1
View File
@@ -77,7 +77,7 @@ type SyncTaskQueryParams struct {
UserIDs []int
}
func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.TaskRelayInfo) *Task {
func InitTask(platform constant.TaskPlatform, relayInfo *commonRelay.RelayInfo) *Task {
t := &Task{
UserId: relayInfo.UserId,
SubmitTime: time.Now().Unix(),
model/token.go
+9 -9
View File
@@ -91,7 +91,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())
common.SysLog("failed to update token status" + err.Error())
}
}
return token, errors.New("该令牌已过期")
@@ -102,7 +102,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())
common.SysLog("failed to update token status" + err.Error())
}
}
keyPrefix := key[:3]
@@ -134,7 +134,7 @@ func GetTokenById(id int) (*Token, error) {
if shouldUpdateRedis(true, err) {
gopool.Go(func() {
if err := cacheSetToken(token); err != nil {
common.SysError("failed to update user status cache: " + err.Error())
common.SysLog("failed to update user status cache: " + err.Error())
}
})
}
@@ -147,7 +147,7 @@ func GetTokenByKey(key string, fromDB bool) (token *Token, err error) {
if shouldUpdateRedis(fromDB, err) && token != nil {
gopool.Go(func() {
if err := cacheSetToken(*token); err != nil {
common.SysError("failed to update user status cache: " + err.Error())
common.SysLog("failed to update user status cache: " + err.Error())
}
})
}
@@ -178,7 +178,7 @@ func (token *Token) Update() (err error) {
gopool.Go(func() {
err := cacheSetToken(*token)
if err != nil {
common.SysError("failed to update token cache: " + err.Error())
common.SysLog("failed to update token cache: " + err.Error())
}
})
}
@@ -194,7 +194,7 @@ func (token *Token) SelectUpdate() (err error) {
gopool.Go(func() {
err := cacheSetToken(*token)
if err != nil {
common.SysError("failed to update token cache: " + err.Error())
common.SysLog("failed to update token cache: " + err.Error())
}
})
}
@@ -209,7 +209,7 @@ func (token *Token) Delete() (err error) {
gopool.Go(func() {
err := cacheDeleteToken(token.Key)
if err != nil {
common.SysError("failed to delete token cache: " + err.Error())
common.SysLog("failed to delete token cache: " + err.Error())
}
})
}
@@ -269,7 +269,7 @@ func IncreaseTokenQuota(id int, key string, quota int) (err error) {
gopool.Go(func() {
err := cacheIncrTokenQuota(key, int64(quota))
if err != nil {
common.SysError("failed to increase token quota: " + err.Error())
common.SysLog("failed to increase token quota: " + err.Error())
}
})
}
@@ -299,7 +299,7 @@ func DecreaseTokenQuota(id int, key string, quota int) (err error) {
gopool.Go(func() {
err := cacheDecrTokenQuota(key, int64(quota))
if err != nil {
common.SysError("failed to decrease token quota: " + err.Error())
common.SysLog("failed to decrease token quota: " + err.Error())
}
})
}
model/topup.go
+3 -2
View File
@@ -4,6 +4,7 @@ import (
"errors"
"fmt"
"one-api/common"
"one-api/logger"
"gorm.io/gorm"
)
@@ -94,7 +95,7 @@ func Recharge(referenceId string, customerId string) (err error) {
return errors.New("充值失败，" + err.Error())
}
RecordLog(topUp.UserId, LogTypeTopup, fmt.Sprintf("使用在线充值成功，充值金额: %v，支付金额：%d", common.FormatQuota(int(quota)), topUp.Amount))
RecordLog(topUp.UserId, LogTypeTopup, fmt.Sprintf("使用在线充值成功，充值金额: %v，支付金额：%d", logger.FormatQuota(int(quota)), topUp.Amount))
return nil
}
@@ -163,7 +164,7 @@ func RechargeCreem(referenceId string, customerEmail string, customerName string
return errors.New("充值失败，" + err.Error())
}
RecordLog(topUp.UserId, LogTypeTopup, fmt.Sprintf("使用Creem充值成功，充值额度: %v，支付金额：%.2f", common.FormatQuota(int(quota)), topUp.Money))
RecordLog(topUp.UserId, LogTypeTopup, fmt.Sprintf("使用Creem充值成功，充值额度: %v，支付金额：%.2f", quota, topUp.Money))
return nil
}
model/twofa.go
+322
View File
@@ -0,0 +1,322 @@
package model
import (
"errors"
"fmt"
"one-api/common"
"time"
"gorm.io/gorm"
)
var ErrTwoFANotEnabled = errors.New("用户未启用2FA")
// TwoFA 用户2FA设置表
type TwoFA struct {
Id int `json:"id" gorm:"primaryKey"`
UserId int `json:"user_id" gorm:"unique;not null;index"`
Secret string `json:"-" gorm:"type:varchar(255);not null"` // TOTP密钥，不返回给前端
IsEnabled bool `json:"is_enabled"`
FailedAttempts int `json:"failed_attempts" gorm:"default:0"`
LockedUntil *time.Time `json:"locked_until,omitempty"`
LastUsedAt *time.Time `json:"last_used_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
UpdatedAt time.Time `json:"updated_at"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
}
// TwoFABackupCode 备用码使用记录表
type TwoFABackupCode struct {
Id int `json:"id" gorm:"primaryKey"`
UserId int `json:"user_id" gorm:"not null;index"`
CodeHash string `json:"-" gorm:"type:varchar(255);not null"` // 备用码哈希
IsUsed bool `json:"is_used"`
UsedAt *time.Time `json:"used_at,omitempty"`
CreatedAt time.Time `json:"created_at"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index"`
}
// GetTwoFAByUserId 根据用户ID获取2FA设置
func GetTwoFAByUserId(userId int) (*TwoFA, error) {
if userId == 0 {
return nil, errors.New("用户ID不能为空")
}
var twoFA TwoFA
err := DB.Where("user_id = ?", userId).First(&twoFA).Error
if err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return nil, nil // 返回nil表示未设置2FA
}
return nil, err
}
return &twoFA, nil
}
// IsTwoFAEnabled 检查用户是否启用了2FA
func IsTwoFAEnabled(userId int) bool {
twoFA, err := GetTwoFAByUserId(userId)
if err != nil || twoFA == nil {
return false
}
return twoFA.IsEnabled
}
// CreateTwoFA 创建2FA设置
func (t *TwoFA) Create() error {
// 检查用户是否已存在2FA设置
existing, err := GetTwoFAByUserId(t.UserId)
if err != nil {
return err
}
if existing != nil {
return errors.New("用户已存在2FA设置")
}
// 验证用户存在
var user User
if err := DB.First(&user, t.UserId).Error; err != nil {
if errors.Is(err, gorm.ErrRecordNotFound) {
return errors.New("用户不存在")
}
return err
}
return DB.Create(t).Error
}
// Update 更新2FA设置
func (t *TwoFA) Update() error {
if t.Id == 0 {
return errors.New("2FA记录ID不能为空")
}
return DB.Save(t).Error
}
// Delete 删除2FA设置
func (t *TwoFA) Delete() error {
if t.Id == 0 {
return errors.New("2FA记录ID不能为空")
}
// 使用事务确保原子性
return DB.Transaction(func(tx *gorm.DB) error {
// 同时删除相关的备用码记录（硬删除）
if err := tx.Unscoped().Where("user_id = ?", t.UserId).Delete(&TwoFABackupCode{}).Error; err != nil {
return err
}
// 硬删除2FA记录
return tx.Unscoped().Delete(t).Error
})
}
// ResetFailedAttempts 重置失败尝试次数
func (t *TwoFA) ResetFailedAttempts() error {
t.FailedAttempts = 0
t.LockedUntil = nil
return t.Update()
}
// IncrementFailedAttempts 增加失败尝试次数
func (t *TwoFA) IncrementFailedAttempts() error {
t.FailedAttempts++
// 检查是否需要锁定
if t.FailedAttempts >= common.MaxFailAttempts {
lockUntil := time.Now().Add(time.Duration(common.LockoutDuration) * time.Second)
t.LockedUntil = &lockUntil
}
return t.Update()
}
// IsLocked 检查账户是否被锁定
func (t *TwoFA) IsLocked() bool {
if t.LockedUntil == nil {
return false
}
return time.Now().Before(*t.LockedUntil)
}
// CreateBackupCodes 创建备用码
func CreateBackupCodes(userId int, codes []string) error {
return DB.Transaction(func(tx *gorm.DB) error {
// 先删除现有的备用码
if err := tx.Where("user_id = ?", userId).Delete(&TwoFABackupCode{}).Error; err != nil {
return err
}
// 创建新的备用码记录
for _, code := range codes {
hashedCode, err := common.HashBackupCode(code)
if err != nil {
return err
}
backupCode := TwoFABackupCode{
UserId: userId,
CodeHash: hashedCode,
IsUsed: false,
}
if err := tx.Create(&backupCode).Error; err != nil {
return err
}
}
return nil
})
}
// ValidateBackupCode 验证并使用备用码
func ValidateBackupCode(userId int, code string) (bool, error) {
if !common.ValidateBackupCode(code) {
return false, errors.New("验证码或备用码不正确")
}
normalizedCode := common.NormalizeBackupCode(code)
// 查找未使用的备用码
var backupCodes []TwoFABackupCode
if err := DB.Where("user_id = ? AND is_used = false", userId).Find(&backupCodes).Error; err != nil {
return false, err
}
// 验证备用码
for _, bc := range backupCodes {
if common.ValidatePasswordAndHash(normalizedCode, bc.CodeHash) {
// 标记为已使用
now := time.Now()
bc.IsUsed = true
bc.UsedAt = &now
if err := DB.Save(&bc).Error; err != nil {
return false, err
}
return true, nil
}
}
return false, nil
}
// GetUnusedBackupCodeCount 获取未使用的备用码数量
func GetUnusedBackupCodeCount(userId int) (int, error) {
var count int64
err := DB.Model(&TwoFABackupCode{}).Where("user_id = ? AND is_used = false", userId).Count(&count).Error
return int(count), err
}
// DisableTwoFA 禁用用户的2FA
func DisableTwoFA(userId int) error {
twoFA, err := GetTwoFAByUserId(userId)
if err != nil {
return err
}
if twoFA == nil {
return ErrTwoFANotEnabled
}
// 删除2FA设置和备用码
return twoFA.Delete()
}
// EnableTwoFA 启用2FA
func (t *TwoFA) Enable() error {
t.IsEnabled = true
t.FailedAttempts = 0
t.LockedUntil = nil
return t.Update()
}
// ValidateTOTPAndUpdateUsage 验证TOTP并更新使用记录
func (t *TwoFA) ValidateTOTPAndUpdateUsage(code string) (bool, error) {
// 检查是否被锁定
if t.IsLocked() {
return false, fmt.Errorf("账户已被锁定，请在%v后重试", t.LockedUntil.Format("2006-01-02 15:04:05"))
}
// 验证TOTP码
if !common.ValidateTOTPCode(t.Secret, code) {
// 增加失败次数
if err := t.IncrementFailedAttempts(); err != nil {
common.SysLog("更新2FA失败次数失败: " + err.Error())
}
return false, nil
}
// 验证成功，重置失败次数并更新最后使用时间
now := time.Now()
t.FailedAttempts = 0
t.LockedUntil = nil
t.LastUsedAt = &now
if err := t.Update(); err != nil {
common.SysLog("更新2FA使用记录失败: " + err.Error())
}
return true, nil
}
// ValidateBackupCodeAndUpdateUsage 验证备用码并更新使用记录
func (t *TwoFA) ValidateBackupCodeAndUpdateUsage(code string) (bool, error) {
// 检查是否被锁定
if t.IsLocked() {
return false, fmt.Errorf("账户已被锁定，请在%v后重试", t.LockedUntil.Format("2006-01-02 15:04:05"))
}
// 验证备用码
valid, err := ValidateBackupCode(t.UserId, code)
if err != nil {
return false, err
}
if !valid {
// 增加失败次数
if err := t.IncrementFailedAttempts(); err != nil {
common.SysLog("更新2FA失败次数失败: " + err.Error())
}
return false, nil
}
// 验证成功，重置失败次数并更新最后使用时间
now := time.Now()
t.FailedAttempts = 0
t.LockedUntil = nil
t.LastUsedAt = &now
if err := t.Update(); err != nil {
common.SysLog("更新2FA使用记录失败: " + err.Error())
}
return true, nil
}
// GetTwoFAStats 获取2FA统计信息（管理员使用）
func GetTwoFAStats() (map[string]interface{}, error) {
var totalUsers, enabledUsers int64
// 总用户数
if err := DB.Model(&User{}).Count(&totalUsers).Error; err != nil {
return nil, err
}
// 启用2FA的用户数
if err := DB.Model(&TwoFA{}).Where("is_enabled = true").Count(&enabledUsers).Error; err != nil {
return nil, err
}
enabledRate := float64(0)
if totalUsers > 0 {
enabledRate = float64(enabledUsers) / float64(totalUsers) * 100
}
return map[string]interface{}{
"total_users": totalUsers,
"enabled_users": enabledUsers,
"enabled_rate": fmt.Sprintf("%.1f%%", enabledRate),
}, nil
}
model/usedata.go
-6
View File
@@ -21,12 +21,6 @@ type QuotaData struct {
}
func UpdateQuotaData() {
// recover
defer func() {
if r := recover(); r != nil {
common.SysLog(fmt.Sprintf("UpdateQuotaData panic: %s", r))
}
}()
for {
if common.DataExportEnabled {
common.SysLog("正在更新数据看板数据...")
model/user.go
+103 -16
View File
@@ -6,6 +6,7 @@ import (
"fmt"
"one-api/common"
"one-api/dto"
"one-api/logger"
"strconv"
"strings"
@@ -75,7 +76,7 @@ func (user *User) GetSetting() dto.UserSetting {
if user.Setting != "" {
err := json.Unmarshal([]byte(user.Setting), &setting)
if err != nil {
common.SysError("failed to unmarshal setting: " + err.Error())
common.SysLog("failed to unmarshal setting: " + err.Error())
}
}
return setting
@@ -84,12 +85,74 @@ func (user *User) GetSetting() dto.UserSetting {
func (user *User) SetSetting(setting dto.UserSetting) {
settingBytes, err := json.Marshal(setting)
if err != nil {
common.SysError("failed to marshal setting: " + err.Error())
common.SysLog("failed to marshal setting: " + err.Error())
return
}
user.Setting = string(settingBytes)
}
// 根据用户角色生成默认的边栏配置
func generateDefaultSidebarConfigForRole(userRole int) string {
defaultConfig := map[string]interface{}{}
// 聊天区域 - 所有用户都可以访问
defaultConfig["chat"] = map[string]interface{}{
"enabled": true,
"playground": true,
"chat": true,
}
// 控制台区域 - 所有用户都可以访问
defaultConfig["console"] = map[string]interface{}{
"enabled": true,
"detail": true,
"token": true,
"log": true,
"midjourney": true,
"task": true,
}
// 个人中心区域 - 所有用户都可以访问
defaultConfig["personal"] = map[string]interface{}{
"enabled": true,
"topup": true,
"personal": true,
}
// 管理员区域 - 根据角色决定
if userRole == common.RoleAdminUser {
// 管理员可以访问管理员区域，但不能访问系统设置
defaultConfig["admin"] = map[string]interface{}{
"enabled": true,
"channel": true,
"models": true,
"redemption": true,
"user": true,
"setting": false, // 管理员不能访问系统设置
}
} else if userRole == common.RoleRootUser {
// 超级管理员可以访问所有功能
defaultConfig["admin"] = map[string]interface{}{
"enabled": true,
"channel": true,
"models": true,
"redemption": true,
"user": true,
"setting": true,
}
}
// 普通用户不包含admin区域
// 转换为JSON字符串
configBytes, err := json.Marshal(defaultConfig)
if err != nil {
common.SysLog("生成默认边栏配置失败: " + err.Error())
return ""
}
return string(configBytes)
}
// CheckUserExistOrDeleted check if user exist or deleted, if not exist, return false, nil, if deleted or exist, return true, nil
func CheckUserExistOrDeleted(username string, email string) (bool, error) {
var user User
@@ -274,7 +337,7 @@ func inviteUser(inviterId int) (err error) {
func (user *User) TransferAffQuotaToQuota(quota int) error {
// 检查quota是否小于最小额度
if float64(quota) < common.QuotaPerUnit {
return fmt.Errorf("转移额度最小为%s！", common.LogQuota(int(common.QuotaPerUnit)))
return fmt.Errorf("转移额度最小为%s！", logger.LogQuota(int(common.QuotaPerUnit)))
}
// 开始数据库事务
@@ -319,21 +382,45 @@ func (user *User) Insert(inviterId int) error {
user.Quota = common.QuotaForNewUser
//user.SetAccessToken(common.GetUUID())
user.AffCode = common.GetRandomString(4)
// 初始化用户设置，包括默认的边栏配置
if user.Setting == "" {
defaultSetting := dto.UserSetting{}
// 这里暂时不设置SidebarModules，因为需要在用户创建后根据角色设置
user.SetSetting(defaultSetting)
}
result := DB.Create(user)
if result.Error != nil {
return result.Error
}
// 用户创建成功后，根据角色初始化边栏配置
// 需要重新获取用户以确保有正确的ID和Role
var createdUser User
if err := DB.Where("username = ?", user.Username).First(&createdUser).Error; err == nil {
// 生成基于角色的默认边栏配置
defaultSidebarConfig := generateDefaultSidebarConfigForRole(createdUser.Role)
if defaultSidebarConfig != "" {
currentSetting := createdUser.GetSetting()
currentSetting.SidebarModules = defaultSidebarConfig
createdUser.SetSetting(currentSetting)
createdUser.Update(false)
common.SysLog(fmt.Sprintf("为新用户 %s (角色: %d) 初始化边栏配置", createdUser.Username, createdUser.Role))
}
}
if common.QuotaForNewUser > 0 {
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", common.LogQuota(common.QuotaForNewUser)))
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("新用户注册赠送 %s", logger.LogQuota(common.QuotaForNewUser)))
}
if inviterId != 0 {
if common.QuotaForInvitee > 0 {
_ = IncreaseUserQuota(user.Id, common.QuotaForInvitee, true)
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", common.LogQuota(common.QuotaForInvitee)))
RecordLog(user.Id, LogTypeSystem, fmt.Sprintf("使用邀请码赠送 %s", logger.LogQuota(common.QuotaForInvitee)))
}
if common.QuotaForInviter > 0 {
//_ = IncreaseUserQuota(inviterId, common.QuotaForInviter)
RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", common.LogQuota(common.QuotaForInviter)))
RecordLog(inviterId, LogTypeSystem, fmt.Sprintf("邀请用户赠送 %s", logger.LogQuota(common.QuotaForInviter)))
_ = inviteUser(inviterId)
}
}
@@ -517,7 +604,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())
common.SysLog("no such user " + err.Error())
return false
}
return user.Role >= common.RoleAdminUser
@@ -572,7 +659,7 @@ func GetUserQuota(id int, fromDB bool) (quota int, err error) {
if shouldUpdateRedis(fromDB, err) {
gopool.Go(func() {
if err := updateUserQuotaCache(id, quota); err != nil {
common.SysError("failed to update user quota cache: " + err.Error())
common.SysLog("failed to update user quota cache: " + err.Error())
}
})
}
@@ -610,7 +697,7 @@ func GetUserGroup(id int, fromDB bool) (group string, err error) {
if shouldUpdateRedis(fromDB, err) {
gopool.Go(func() {
if err := updateUserGroupCache(id, group); err != nil {
common.SysError("failed to update user group cache: " + err.Error())
common.SysLog("failed to update user group cache: " + err.Error())
}
})
}
@@ -639,7 +726,7 @@ func GetUserSetting(id int, fromDB bool) (settingMap dto.UserSetting, err error)
if shouldUpdateRedis(fromDB, err) {
gopool.Go(func() {
if err := updateUserSettingCache(id, setting); err != nil {
common.SysError("failed to update user setting cache: " + err.Error())
common.SysLog("failed to update user setting cache: " + err.Error())
}
})
}
@@ -669,7 +756,7 @@ func IncreaseUserQuota(id int, quota int, db bool) (err error) {
gopool.Go(func() {
err := cacheIncrUserQuota(id, int64(quota))
if err != nil {
common.SysError("failed to increase user quota: " + err.Error())
common.SysLog("failed to increase user quota: " + err.Error())
}
})
if !db && common.BatchUpdateEnabled {
@@ -694,7 +781,7 @@ func DecreaseUserQuota(id int, quota int) (err error) {
gopool.Go(func() {
err := cacheDecrUserQuota(id, int64(quota))
if err != nil {
common.SysError("failed to decrease user quota: " + err.Error())
common.SysLog("failed to decrease user quota: " + err.Error())
}
})
if common.BatchUpdateEnabled {
@@ -750,7 +837,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())
common.SysLog("failed to update user used quota and request count: " + err.Error())
return
}
@@ -767,14 +854,14 @@ func updateUserUsedQuota(id int, quota int) {
},
).Error
if err != nil {
common.SysError("failed to update user used quota: " + err.Error())
common.SysLog("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())
common.SysLog("failed to update user request count: " + err.Error())
}
}
@@ -785,7 +872,7 @@ func GetUsernameById(id int, fromDB bool) (username string, err error) {
if shouldUpdateRedis(fromDB, err) {
gopool.Go(func() {
if err := updateUserNameCache(id, username); err != nil {
common.SysError("failed to update user name cache: " + err.Error())
common.SysLog("failed to update user name cache: " + err.Error())
}
})
}
model/user_cache.go
+2 -2
View File
@@ -37,7 +37,7 @@ func (user *UserBase) GetSetting() dto.UserSetting {
if user.Setting != "" {
err := common.Unmarshal([]byte(user.Setting), &setting)
if err != nil {
common.SysError("failed to unmarshal setting: " + err.Error())
common.SysLog("failed to unmarshal setting: " + err.Error())
}
}
return setting
@@ -78,7 +78,7 @@ func GetUserCache(userId int) (userCache *UserBase, err error) {
if shouldUpdateRedis(fromDB, err) && user != nil {
gopool.Go(func() {
if err := updateUserCache(*user); err != nil {
common.SysError("failed to update user status cache: " + err.Error())
common.SysLog("failed to update user status cache: " + err.Error())
}
})
}
model/utils.go
+2 -2
View File
@@ -77,12 +77,12 @@ func batchUpdate() {
case BatchUpdateTypeUserQuota:
err := increaseUserQuota(key, value)
if err != nil {
common.SysError("failed to batch update user quota: " + err.Error())
common.SysLog("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())
common.SysLog("failed to batch update token quota: " + err.Error())
}
case BatchUpdateTypeUsedQuota:
updateUserUsedQuota(key, value)
model/vendor_meta.go
+88
View File
@@ -0,0 +1,88 @@
package model
import (
"one-api/common"
"gorm.io/gorm"
)
// Vendor 用于存储供应商信息，供模型引用
// Name 唯一，用于在模型中关联
// Icon 采用 @lobehub/icons 的图标名，前端可直接渲染
// Status 预留字段，1 表示启用
// 本表同样遵循 3NF 设计范式
type Vendor struct {
Id int `json:"id"`
Name string `json:"name" gorm:"size:128;not null;uniqueIndex:uk_vendor_name_delete_at,priority:1"`
Description string `json:"description,omitempty" gorm:"type:text"`
Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"`
Status int `json:"status" gorm:"default:1"`
CreatedTime int64 `json:"created_time" gorm:"bigint"`
UpdatedTime int64 `json:"updated_time" gorm:"bigint"`
DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_vendor_name_delete_at,priority:2"`
}
// Insert 创建新的供应商记录
func (v *Vendor) Insert() error {
now := common.GetTimestamp()
v.CreatedTime = now
v.UpdatedTime = now
return DB.Create(v).Error
}
// IsVendorNameDuplicated 检查供应商名称是否重复（排除自身 ID）
func IsVendorNameDuplicated(id int, name string) (bool, error) {
if name == "" {
return false, nil
}
var cnt int64
err := DB.Model(&Vendor{}).Where("name = ? AND id <> ?", name, id).Count(&cnt).Error
return cnt > 0, err
}
// Update 更新供应商记录
func (v *Vendor) Update() error {
v.UpdatedTime = common.GetTimestamp()
return DB.Save(v).Error
}
// Delete 软删除供应商
func (v *Vendor) Delete() error {
return DB.Delete(v).Error
}
// GetVendorByID 根据 ID 获取供应商
func GetVendorByID(id int) (*Vendor, error) {
var v Vendor
err := DB.First(&v, id).Error
if err != nil {
return nil, err
}
return &v, nil
}
// GetAllVendors 获取全部供应商（分页）
func GetAllVendors(offset int, limit int) ([]*Vendor, error) {
var vendors []*Vendor
err := DB.Offset(offset).Limit(limit).Find(&vendors).Error
return vendors, err
}
// SearchVendors 按关键字搜索供应商
func SearchVendors(keyword string, offset int, limit int) ([]*Vendor, int64, error) {
db := DB.Model(&Vendor{})
if keyword != "" {
like := "%" + keyword + "%"
db = db.Where("name LIKE ? OR description LIKE ?", like, like)
}
var total int64
if err := db.Count(&total).Error; err != nil {
return nil, 0, err
}
var vendors []*Vendor
if err := db.Offset(offset).Limit(limit).Order("id DESC").Find(&vendors).Error; err != nil {
return nil, 0, err
}
return vendors, total, nil
}