mirror of
https://github.com/songquanpeng/one-api.git
synced 2025-10-14 23:00:27 +00:00
Co-authored-by: quzard <1191890118@qq.com>
This commit is contained in:
72
model/ability.go
Normal file
72
model/ability.go
Normal file
@@ -0,0 +1,72 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"one-api/common"
|
||||
"strings"
|
||||
)
|
||||
|
||||
type Ability struct {
|
||||
Group string `json:"group" gorm:"type:varchar(32);primaryKey;autoIncrement:false"`
|
||||
Model string `json:"model" gorm:"primaryKey;autoIncrement:false"`
|
||||
ChannelId int `json:"channel_id" gorm:"primaryKey;autoIncrement:false;index"`
|
||||
Enabled bool `json:"enabled" gorm:"default:1"`
|
||||
}
|
||||
|
||||
func GetRandomSatisfiedChannel(group string, model string) (*Channel, error) {
|
||||
if group == "default" {
|
||||
return GetRandomChannel()
|
||||
}
|
||||
ability := Ability{}
|
||||
var err error = nil
|
||||
if common.UsingSQLite {
|
||||
err = DB.Where("`group` = ? and model = ? and enabled = 1", group, model).Order("RANDOM()").Limit(1).First(&ability).Error
|
||||
} else {
|
||||
err = DB.Where("`group` = ? and model = ? and enabled = 1", group, model).Order("RAND()").Limit(1).First(&ability).Error
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
channel := Channel{}
|
||||
err = DB.First(&channel, "id = ?", ability.ChannelId).Error
|
||||
return &channel, err
|
||||
}
|
||||
|
||||
func (channel *Channel) AddAbilities() error {
|
||||
models_ := strings.Split(channel.Models, ",")
|
||||
abilities := make([]Ability, 0, len(models_))
|
||||
for _, model := range models_ {
|
||||
ability := Ability{
|
||||
Group: channel.Group,
|
||||
Model: model,
|
||||
ChannelId: channel.Id,
|
||||
Enabled: channel.Status == common.ChannelStatusEnabled,
|
||||
}
|
||||
abilities = append(abilities, ability)
|
||||
}
|
||||
return DB.Create(&abilities).Error
|
||||
}
|
||||
|
||||
func (channel *Channel) DeleteAbilities() error {
|
||||
return DB.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error
|
||||
}
|
||||
|
||||
// UpdateAbilities updates abilities of this channel.
|
||||
// Make sure the channel is completed before calling this function.
|
||||
func (channel *Channel) UpdateAbilities() error {
|
||||
// A quick and dirty way to update abilities
|
||||
// First delete all abilities of this channel
|
||||
err := channel.DeleteAbilities()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// Then add new abilities
|
||||
err = channel.AddAbilities()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func UpdateAbilityStatus(channelId int, status bool) error {
|
||||
return DB.Model(&Ability{}).Where("channel_id = ?", channelId).Update("enabled", status).Error
|
||||
}
|
Reference in New Issue
Block a user