Merge commit '3d149fedf45472eff92910324974c762fc37dad6'

middleware/auth.go

@@ -2,14 +2,16 @@ package middleware
 
 import (
 	"fmt"
+	"net/http"
+	"strings"
+
 	"github.com/Laisky/one-api/common/blacklist"
+	"github.com/Laisky/one-api/common/ctxkey"
 	"github.com/Laisky/one-api/common/logger"
 	"github.com/Laisky/one-api/common/network"
 	"github.com/Laisky/one-api/model"
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-gonic/gin"
-	"net/http"
-	"strings"
 )
 
 func authHelper(c *gin.Context, minRole int) {
@@ -123,20 +125,20 @@ func TokenAuth() func(c *gin.Context) {
 			abortWithMessage(c, http.StatusBadRequest, err.Error())
 			return
 		}
-		c.Set("request_model", requestModel)
+		c.Set(ctxkey.RequestModel, requestModel)
 		if token.Models != nil && *token.Models != "" {
-			c.Set("available_models", *token.Models)
+			c.Set(ctxkey.AvailableModels, *token.Models)
 			if requestModel != "" && !isModelInList(requestModel, *token.Models) {
 				abortWithMessage(c, http.StatusForbidden, fmt.Sprintf("该令牌无权使用模型：%s", requestModel))
 				return
 			}
 		}
-		c.Set("id", token.UserId)
-		c.Set("token_id", token.Id)
-		c.Set("token_name", token.Name)
+		c.Set(ctxkey.Id, token.UserId)
+		c.Set(ctxkey.TokenId, token.Id)
+		c.Set(ctxkey.TokenName, token.Name)
 		if len(parts) > 1 {
 			if model.IsAdmin(token.UserId) {
-				c.Set("specific_channel_id", parts[1])
+				c.Set(ctxkey.SpecificChannelId, parts[1])
 			} else {
 				abortWithMessage(c, http.StatusForbidden, "普通用户不支持指定渠道")
 				return

middleware/distributor.go

@@ -6,8 +6,6 @@ import (
 	"strconv"
 	"strings"
 
-	"github.com/Laisky/one-api/common"
-	"github.com/Laisky/one-api/common/config"
 	"github.com/Laisky/one-api/common/ctxkey"
 	"github.com/Laisky/one-api/common/logger"
 	"github.com/Laisky/one-api/model"
@@ -22,12 +20,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 {
@@ -44,7 +42,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 {
@@ -74,30 +72,31 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
 		}
 	}
 	logger.Info(c.Request.Context(), fmt.Sprintf("set channel %s ratio to %f", channel.Name, minimalRatio))
-	c.Set("channel_ratio", minimalRatio)
-	c.Set(common.CtxKeyChannel, channel)
-	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.ChannelRatio, minimalRatio)
+	c.Set(ctxkey.ChannelModel, channel)
+
+	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:
-		c.Set(config.KeyAPIVersion, channel.Other)
+		c.Set(ctxkey.ConfigAPIVersion, channel.Other)
 	case channeltype.Xunfei:
-		c.Set(config.KeyAPIVersion, channel.Other)
+		c.Set(ctxkey.ConfigAPIVersion, channel.Other)
 	case channeltype.Gemini:
-		c.Set(config.KeyAPIVersion, channel.Other)
+		c.Set(ctxkey.ConfigAPIVersion, channel.Other)
 	case channeltype.AIProxyLibrary:
-		c.Set(config.KeyLibraryID, channel.Other)
+		c.Set(ctxkey.ConfigLibraryID, channel.Other)
 	case channeltype.Ali:
-		c.Set(config.KeyPlugin, channel.Other)
+		c.Set(ctxkey.ConfigPlugin, channel.Other)
 	}
 	cfg, _ := channel.LoadConfig()
 	for k, v := range cfg {
-		c.Set(config.KeyPrefix+k, v)
+		c.Set(ctxkey.ConfigPrefix+k, v)
 	}
 }