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

feat: support group now (close #17, close #72, close #85, close #104, close #136)

Browse Source 
Co-authored-by: quzard <1191890118@qq.com>
This commit is contained in:
JustSong 2023-06-07 23:26:00 +08:00
parent 502515bbbd
commit 2ad22e1425
15 changed files with 235 additions and 29 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
common/gin.go Normal file
View File

@ -0,0 +1,26 @@
package common
import (
"bytes"
"encoding/json"
"github.com/gin-gonic/gin"
"io"
)
func UnmarshalBodyReusable(c *gin.Context, v any) error {
requestBody, err := io.ReadAll(c.Request.Body)
if err != nil {
return err
}
err = c.Request.Body.Close()
if err != nil {
return err
}
err = json.Unmarshal(requestBody, &v)
if err != nil {
return err
}
// Reset request body
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
return nil
}

14
controller/relay.go
View File

@ -116,20 +116,10 @@ func relayHelper(c *gin.Context) *OpenAIErrorWithStatusCode {
consumeQuota := c.GetBool("consume_quota")
var textRequest GeneralOpenAIRequest
if consumeQuota || channelType == common.ChannelTypeAzure || channelType == common.ChannelTypePaLM {
requestBody, err := io.ReadAll(c.Request.Body)
err := common.UnmarshalBodyReusable(c, &textRequest)
if err != nil {
return errorWrapper(err, "read_request_body_failed", http.StatusBadRequest)
return errorWrapper(err, "bind_request_body_failed", http.StatusBadRequest)
}
err = c.Request.Body.Close()
if err != nil {
return errorWrapper(err, "close_request_body_failed", http.StatusBadRequest)
}
err = json.Unmarshal(requestBody, &textRequest)
if err != nil {
return errorWrapper(err, "unmarshal_request_body_failed", http.StatusBadRequest)
}
// Reset request body
c.Request.Body = io.NopCloser(bytes.NewBuffer(requestBody))
}
baseURL := common.ChannelBaseURLs[channelType]
requestURL := c.Request.URL.String()

21
middleware/distributor.go
View File

@ -9,6 +9,10 @@ import (
"strconv"
)
type ModelRequest struct {
Model string `json:"model"`
}
func Distribute() func(c *gin.Context) {
return func(c *gin.Context) {
var channel *model.Channel
@ -48,8 +52,21 @@ func Distribute() func(c *gin.Context) {
}
} else {
// Select a channel for the user
var err error
channel, err = model.GetRandomChannel()
var modelRequest ModelRequest
err := common.UnmarshalBodyReusable(c, &modelRequest)
if err != nil {
c.JSON(200, gin.H{
"error": gin.H{
"message": "无效的请求",
"type": "one_api_error",
},
})
c.Abort()
return
}
userId := c.GetInt("id")
userGroup, _ := model.GetUserGroup(userId)
channel, err = model.GetRandomSatisfiedChannel(userGroup, modelRequest.Model)
if err != nil {
c.JSON(200, gin.H{
"error": gin.H{

72
model/ability.go Normal file
View 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
}

37
model/channel.go
View File

@ -1,7 +1,6 @@
package model
import (
_ "gorm.io/driver/sqlite"
"one-api/common"
)
@ -19,6 +18,8 @@ type Channel struct {
Other string `json:"other"`
Balance float64 `json:"balance"` // in USD
BalanceUpdatedTime int64 `json:"balance_updated_time" gorm:"bigint"`
Models string `json:"models"`
Group string `json:"group" gorm:"type:varchar(32);default:'default'"`
}
func GetAllChannels(startIdx int, num int, selectAll bool) ([]*Channel, error) {
@ -49,13 +50,12 @@ func GetChannelById(id int, selectAll bool) (*Channel, error) {
}
func GetRandomChannel() (*Channel, error) {
// TODO: consider weight
channel := Channel{}
var err error = nil
if common.UsingSQLite {
err = DB.Where("status = ?", common.ChannelStatusEnabled).Order("RANDOM()").Limit(1).First(&channel).Error
err = DB.Where("status = ? and `group` = ?", common.ChannelStatusEnabled, "default").Order("RANDOM()").Limit(1).First(&channel).Error
} else {
err = DB.Where("status = ?", common.ChannelStatusEnabled).Order("RAND()").Limit(1).First(&channel).Error
err = DB.Where("status = ? and `group` = ?", common.ChannelStatusEnabled, "default").Order("RAND()").Limit(1).First(&channel).Error
}
return &channel, err
}
@ -63,18 +63,35 @@ func GetRandomChannel() (*Channel, error) {
func BatchInsertChannels(channels []Channel) error {
var err error
err = DB.Create(&channels).Error
return err
if err != nil {
return err
}
for _, channel_ := range channels {
err = channel_.AddAbilities()
if err != nil {
return err
}
}
return nil
}
func (channel *Channel) Insert() error {
var err error
err = DB.Create(channel).Error
if err != nil {
return err
}
err = channel.AddAbilities()
return err
}
func (channel *Channel) Update() error {
var err error
err = DB.Model(channel).Updates(channel).Error
if err != nil {
return err
}
err = channel.UpdateAbilities()
return err
}
@ -101,11 +118,19 @@ func (channel *Channel) UpdateBalance(balance float64) {
func (channel *Channel) Delete() error {
var err error
err = DB.Delete(channel).Error
if err != nil {
return err
}
err = channel.DeleteAbilities()
return err
}
func UpdateChannelStatusById(id int, status int) {
err := DB.Model(&Channel{}).Where("id = ?", id).Update("status", status).Error
err := UpdateAbilityStatus(id, status == common.ChannelStatusEnabled)
if err != nil {
common.SysError("failed to update ability status: " + err.Error())
}
err = DB.Model(&Channel{}).Where("id = ?", id).Update("status", status).Error
if err != nil {
common.SysError("failed to update channel status: " + err.Error())
}

4
model/main.go
View File

@ -75,6 +75,10 @@ func InitDB() (err error) {
if err != nil {
return err
}
err = db.AutoMigrate(&Ability{})
if err != nil {
return err
}
err = createRootAccountIfNeed()
return err
} else {

1
model/redemption.go
View File

@ -2,7 +2,6 @@ package model
import (
"errors"
_ "gorm.io/driver/sqlite"
"one-api/common"
)

1
model/token.go
View File

@ -3,7 +3,6 @@ package model
import (
"errors"
"fmt"
_ "gorm.io/driver/sqlite"
"gorm.io/gorm"
"one-api/common"
)

6
model/user.go
View File

@ -22,6 +22,7 @@ type User struct {
VerificationCode string `json:"verification_code" gorm:"-:all"` // this field is only for Email verification, don't save it to database!
AccessToken string `json:"access_token" gorm:"type:char(32);column:access_token;uniqueIndex"` // this token is for system management
Quota int `json:"quota" gorm:"type:int;default:0"`
Group string `json:"group" gorm:"type:varchar(32);default:'default'"`
}
func GetMaxUserId() int {
@ -229,6 +230,11 @@ func GetUserEmail(id int) (email string, err error) {
return email, err
}
func GetUserGroup(id int) (group string, err error) {
err = DB.Model(&User{}).Where("id = ?", id).Select("`group`").Find(&group).Error
return group, err
}
func IncreaseUserQuota(id int, quota int) (err error) {
if quota < 0 {
return errors.New("quota 不能为负数！")

1
router/api-router.go
View File

@ -63,6 +63,7 @@ func SetApiRouter(router *gin.Engine) {
{
channelRoute.GET("/", controller.GetAllChannels)
channelRoute.GET("/search", controller.SearchChannels)
channelRoute.GET("/models", controller.ListModels)
channelRoute.GET("/:id", controller.GetChannel)
channelRoute.GET("/test", controller.TestAllChannels)
channelRoute.GET("/test/:id", controller.TestChannel)

8
router/relay-router.go
View File

@ -8,11 +8,15 @@ import (
func SetRelayRouter(router *gin.Engine) {
// https://platform.openai.com/docs/api-reference/introduction
modelsRouter := router.Group("/v1/models")
modelsRouter.Use(middleware.TokenAuth())
{
modelsRouter.GET("/", controller.ListModels)
modelsRouter.GET("/:model", controller.RetrieveModel)
}
relayV1Router := router.Group("/v1")
relayV1Router.Use(middleware.TokenAuth(), middleware.Distribute())
{
relayV1Router.GET("/models", controller.ListModels)
relayV1Router.GET("/models/:model", controller.RetrieveModel)
relayV1Router.POST("/completions", controller.RelayNotImplemented)
relayV1Router.POST("/chat/completions", controller.Relay)
relayV1Router.POST("/edits", controller.RelayNotImplemented)

12
web/src/components/ChannelsTable.js
View File

@ -4,6 +4,7 @@ import { Link } from 'react-router-dom';
import { API, showError, showInfo, showSuccess, timestamp2string } from '../helpers';
import { CHANNEL_OPTIONS, ITEMS_PER_PAGE } from '../constants';
import { renderGroup } from '../helpers/render';
function renderTimestamp(timestamp) {
return (
@ -264,6 +265,14 @@ const ChannelsTable = () => {
>
名称
</Table.HeaderCell>
<Table.HeaderCell
style={{ cursor: 'pointer' }}
onClick={() => {
sortChannel('group');
}}
>
分组
</Table.HeaderCell>
<Table.HeaderCell
style={{ cursor: 'pointer' }}
onClick={() => {
@ -312,6 +321,7 @@ const ChannelsTable = () => {
<Table.Row key={channel.id}>
<Table.Cell>{channel.id}</Table.Cell>
<Table.Cell>{channel.name ? channel.name : '无'}</Table.Cell>
<Table.Cell>{renderGroup(channel.group)}</Table.Cell>
<Table.Cell>{renderType(channel.type)}</Table.Cell>
<Table.Cell>{renderStatus(channel.status)}</Table.Cell>
<Table.Cell>
@ -398,7 +408,7 @@ const ChannelsTable = () => {
<Table.Footer>
<Table.Row>
<Table.HeaderCell colSpan='7'>
<Table.HeaderCell colSpan='8'>
<Button size='small' as={Link} to='/channel/add' loading={loading}>
添加新的渠道
</Button>

13
web/src/components/UsersTable.js
View File

@ -4,7 +4,7 @@ import { Link } from 'react-router-dom';
import { API, showError, showSuccess } from '../helpers';
import { ITEMS_PER_PAGE } from '../constants';
import { renderText } from '../helpers/render';
import { renderGroup, renderText } from '../helpers/render';
function renderRole(role) {
switch (role) {
@ -175,6 +175,14 @@ const UsersTable = () => {
>
用户名
</Table.HeaderCell>
<Table.HeaderCell
style={{ cursor: 'pointer' }}
onClick={() => {
sortUser('group');
}}
>
分组
</Table.HeaderCell>
<Table.HeaderCell
style={{ cursor: 'pointer' }}
onClick={() => {
@ -231,6 +239,7 @@ const UsersTable = () => {
hoverable
/>
</Table.Cell>
<Table.Cell>{renderGroup(user.group)}</Table.Cell>
<Table.Cell>{user.email ? renderText(user.email, 30) : '无'}</Table.Cell>
<Table.Cell>{user.quota}</Table.Cell>
<Table.Cell>{renderRole(user.role)}</Table.Cell>
@ -306,7 +315,7 @@ const UsersTable = () => {
<Table.Footer>
<Table.Row>
<Table.HeaderCell colSpan='7'>
<Table.HeaderCell colSpan='8'>
<Button size='small' as={Link} to='/user/add' loading={loading}>
添加新的用户
</Button>

9
web/src/helpers/render.js
View File

@ -1,6 +1,15 @@
import { Label } from 'semantic-ui-react';
export function renderText(text, limit) {
if (text.length > limit) {
return text.slice(0, limit - 3) + '...';
}
return text;
}
export function renderGroup(group) {
if (group === "") {
return <Label>default</Label>
}
return <Label>{group}</Label>
}

39
web/src/pages/Channel/EditChannel.js
View File

@ -14,10 +14,12 @@ const EditChannel = () => {
type: 1,
key: '',
base_url: '',
other: ''
other: '',
models: [],
};
const [batch, setBatch] = useState(false);
const [inputs, setInputs] = useState(originInputs);
const [modelOptions, setModelOptions] = useState([]);
const handleInputChange = (e, { name, value }) => {
console.log(name, value);
setInputs((inputs) => ({ ...inputs, [name]: value }));
@ -27,17 +29,36 @@ const EditChannel = () => {
let res = await API.get(`/api/channel/${channelId}`);
const { success, message, data } = res.data;
if (success) {
data.password = '';
if (data.models === "") {
data.models = []
} else {
data.models = data.models.split(",")
}
setInputs(data);
} else {
showError(message);
}
setLoading(false);
};
const fetchModels = async () => {
try {
let res = await API.get(`/api/channel/models`);
setModelOptions(res.data.data.map((model) => ({
key: model.id,
text: model.id,
value: model.id,
})));
} catch (error) {
console.error('Error fetching models:', error);
}
};
useEffect(() => {
if (isEdit) {
loadChannel().then();
}
fetchModels().then();
}, []);
const submit = async () => {
@ -50,6 +71,7 @@ const EditChannel = () => {
localInputs.other = '2023-03-15-preview';
}
let res;
localInputs.models = localInputs.models.join(",")
if (isEdit) {
res = await API.put(`/api/channel/`, { ...localInputs, id: parseInt(channelId) });
} else {
@ -137,6 +159,19 @@ const EditChannel = () => {
autoComplete='new-password'
/>
</Form.Field>
<Form.Field>
<Form.Dropdown
label='支持的模型'
name='models'
fluid
multiple
selection
onChange={handleInputChange}
value={inputs.models}
autoComplete='new-password'
options={modelOptions}
/>
</Form.Field>
{
batch ? <Form.Field>
<Form.TextArea