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

Merge branch 'channel-stream-mode' into refactor-main

Browse Source
This commit is contained in:
ckt1031 2023-07-24 22:05:54 +08:00
commit c6c070b8bd
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 ( import (
"fmt" "fmt"
"log"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
@ -85,6 +86,7 @@ func Distribute() func(c *gin.Context) {
modelRequest.Model = "dall-e" modelRequest.Model = "dall-e"
} }
} }
log.Print(modelRequest.Stream)
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model, modelRequest.Stream) channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, modelRequest.Model, modelRequest.Stream)
if err != nil { if err != nil {
message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model) message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", userGroup, modelRequest.Model)

20
model/ability.go
View File

@ -7,10 +7,12 @@ import (
) )
type Ability struct { type Ability struct {
Group string `json:"group" gorm:"type:varchar(32);primaryKey;autoIncrement:false"` Group string `json:"group" gorm:"type:varchar(32);primaryKey;autoIncrement:false"`
Model string `json:"model" gorm:"primaryKey;autoIncrement:false"` Model string `json:"model" gorm:"primaryKey;autoIncrement:false"`
ChannelId int `json:"channel_id" gorm:"primaryKey;autoIncrement:false;index"` ChannelId int `json:"channel_id" gorm:"primaryKey;autoIncrement:false;index"`
Enabled bool `json:"enabled"` 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) { func GetRandomSatisfiedChannel(group string, model string, stream bool) (*Channel, error) {
@ -46,10 +48,12 @@ func (channel *Channel) AddAbilities() error {
for _, model := range models_ { for _, model := range models_ {
for _, group := range groups_ { for _, group := range groups_ {
ability := Ability{ ability := Ability{
Group: group, Group: group,
Model: model, Model: model,
ChannelId: channel.Id, ChannelId: channel.Id,
Enabled: channel.Status == common.ChannelStatusEnabled, Enabled: channel.Status == common.ChannelStatusEnabled,
AllowStreaming: channel.AllowStreaming,
AllowNonStreaming: channel.AllowNonStreaming,
} }
abilities = append(abilities, ability) 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'"` Group string `json:"group" gorm:"type:varchar(32);default:'default'"`
UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"` UsedQuota int64 `json:"used_quota" gorm:"bigint;default:0"`
ModelMapping string `json:"model_mapping" gorm:"type:varchar(1024);default:''"` ModelMapping string `json:"model_mapping" gorm:"type:varchar(1024);default:''"`
AllowStreaming int `json:"allow_streaming" gorm:"default:2"` AllowStreaming int `json:"allow_streaming" gorm:"default:1"`
AllowNonStreaming int `json:"allow_non_streaming" gorm:"default:2"` AllowNonStreaming int `json:"allow_non_streaming" gorm:"default:1"`
} }
func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) { func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {