xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-16 21:23:44 +08:00
Code Issues 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 18:12:53 +08:00
parent 00d3a78bef 30187cebe8
commit 8956e2fd60
6 changed files with 29 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
model/ability.go
View File

@@ -1,6 +1,7 @@
package model
import (
"fmt"
"one-api/common"
"strings"
)
@@ -19,9 +20,9 @@ func GetRandomSatisfiedChannel(group string, model string, stream bool) (*Channe
cmd := "`group` = ? and model = ? and enabled = 1"
if stream {
cmd += " and allow_streaming = 1"
cmd += fmt.Sprintf(" and allow_streaming = %d", common.ChannelAllowStreamEnabled)
} else {
cmd += " and allow_non_streaming = 1"
cmd += fmt.Sprintf(" and allow_non_streaming = %d", common.ChannelAllowNonStreamEnabled)
}
if common.UsingSQLite {

2
model/cache.go
View File

@@ -173,7 +173,7 @@ func CacheGetRandomSatisfiedChannel(group string, model string, stream bool) (*C
var filteredChannels []*Channel
for _, channel := range channels {
if (stream && channel.AllowStreaming) || (!stream && channel.AllowNonStreaming) {
if (stream && channel.AllowStreaming == common.ChannelAllowStreamEnabled) || (!stream && channel.AllowNonStreaming == common.ChannelAllowNonStreamEnabled) {
filteredChannels = append(filteredChannels, channel)
}
}

36
model/channel.go
View File

@@ -1,7 +1,6 @@
package model
import (
"encoding/json"
"one-api/common"
"gorm.io/gorm"
@@ -25,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 bool `json:"allow_streaming"`
AllowNonStreaming bool `json:"allow_non_streaming"`
AllowStreaming int `json:"allow_streaming" gorm:"default:2"`
AllowNonStreaming int `json:"allow_non_streaming" gorm:"default:2"`
}
func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {
@@ -84,19 +83,7 @@ func BatchInsertChannels(channels []Channel) error {
func (channel *Channel) Insert() error {
var err error
// turn channel into a map
channelMap := make(map[string]interface{})
// Convert channel struct to a map
channelBytes, err := json.Marshal(channel)
if err != nil {
return err
}
err = json.Unmarshal(channelBytes, &channelMap)
if err != nil {
return err
}
err = DB.Create(channelMap).Error
err = DB.Create(channel).Error
if err != nil {
return err
}
@@ -106,24 +93,11 @@ func (channel *Channel) Insert() error {
func (channel *Channel) Update() error {
var err error
// turn channel into a map
channelMap := make(map[string]interface{})
// Convert channel struct to a map
channelBytes, err := json.Marshal(channel)
err = DB.Model(channel).Updates(channel).Error
if err != nil {
return err
}
err = json.Unmarshal(channelBytes, &channelMap)
if err != nil {
return err
}
err = DB.Model(channel).Updates(channelMap).Error
if err != nil {
return err
}
DB.Model(channel).First(channelMap, "id = ?", channel.Id)
DB.Model(channel).First(channel, "id = ?", channel.Id)
err = channel.UpdateAbilities()
return err
}