File size: 4,397 Bytes
4674012 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 |
package model
import (
"strconv"
"github.com/QuantumNous/new-api/common"
"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
}
|