xinyin025/one-api
1
0
Fork 0
mirror of https://github.com/songquanpeng/one-api.git synced 2025-11-07 01:03:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

chore: do not hardcode context key

Browse Source
This commit is contained in:
JustSong
2024-04-21 19:43:23 +08:00
parent 83517f687c
commit 3d149fedf4
15 changed files with 98 additions and 73 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
middleware/distributor.go
View File

@@ -17,12 +17,12 @@ type ModelRequest struct {
func Distribute() func(c *gin.Context) {
return func(c *gin.Context) {
userId := c.GetInt("id")
userId := c.GetInt(ctxkey.Id)
userGroup, _ := model.CacheGetUserGroup(userId)
c.Set("group", userGroup)
c.Set(ctxkey.Group, userGroup)
var requestModel string
var channel *model.Channel
channelId, ok := c.Get("specific_channel_id")
channelId, ok := c.Get(ctxkey.SpecificChannelId)
if ok {
id, err := strconv.Atoi(channelId.(string))
if err != nil {
@@ -39,7 +39,7 @@ func Distribute() func(c *gin.Context) {
return
}
} else {
requestModel = c.GetString("request_model")
requestModel = c.GetString(ctxkey.RequestModel)
var err error
channel, err = model.CacheGetRandomSatisfiedChannel(userGroup, requestModel, false)
if err != nil {
@@ -58,13 +58,13 @@ func Distribute() func(c *gin.Context) {
}
func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) {
c.Set("channel", channel.Type)
c.Set("channel_id", channel.Id)
c.Set("channel_name", channel.Name)
c.Set("model_mapping", channel.GetModelMapping())
c.Set(ctxkey.Channel, channel.Type)
c.Set(ctxkey.ChannelId, channel.Id)
c.Set(ctxkey.ChannelName, channel.Name)
c.Set(ctxkey.ModelMapping, channel.GetModelMapping())
c.Set(ctxkey.OriginalModel, modelName) // for retry
c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", channel.Key))
c.Set("base_url", channel.GetBaseURL())
c.Set(ctxkey.BaseURL, channel.GetBaseURL())
// this is for backward compatibility
switch channel.Type {
case channeltype.Azure: