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

fix: channel issue

Browse Source
This commit is contained in:
ckt1031 2023-07-24 22:05:21 +08:00
parent 30187cebe8
commit 13b3bfee2a
3 changed files with 16 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
middleware/distributor.go
View File

@ -2,6 +2,7 @@ package middleware
import (
"fmt"
"log"
"net/http"
"one-api/common"
"one-api/model"
@ -85,6 +86,7 @@ func Distribute() func(c *gin.Context) {
modelRequest.Model = "dall-e"
}
}
log.Print(modelRequest.Stream)
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model, modelRequest.Stream)
if err != nil {
message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model)

20
model/ability.go
View File

@ -7,10 +7,12 @@ import (
)
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"`
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"`
AllowStreaming int `json:"allow_streaming" gorm:"default:1"`
AllowNonStreaming int `json:"allow_non_streaming" gorm:"default:1"`
}
func GetRandomSatisfiedChannel(group string, model string, stream bool) (*Channel, error) {
@ -46,10 +48,12 @@ func (channel *Channel) AddAbilities() error {
for _, model := range models_ {
for _, group := range groups_ {
ability := Ability{
Group: group,
Model: model,
ChannelId: channel.Id,
Enabled: channel.Status == common.ChannelStatusEnabled,
Group: group,
Model: model,
ChannelId: channel.Id,
Enabled: channel.Status == common.ChannelStatusEnabled,
AllowStreaming: channel.AllowStreaming,
AllowNonStreaming: channel.AllowNonStreaming,
}
abilities = append(abilities, ability)
}

4
model/channel.go
View File

@ -24,8 +24,8 @@ type Channel struct {
Group string `json:"group" gorm:"type:varchar(32);default:'default'"`
UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
ModelMapping string `json:"model_mapping" gorm:"type:varchar(1024);default:''"`
AllowStreaming int `json:"allow_streaming" gorm:"default:2"`
AllowNonStreaming int `json:"allow_non_streaming" gorm:"default:2"`
AllowStreaming int `json:"allow_streaming" gorm:"default:1"`
AllowNonStreaming int `json:"allow_non_streaming" gorm:"default:1"`
}
func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {