feat: 为渠道添加上下文限制功能

Author: StageDog

constant/context_key.go

@@ -31,6 +31,7 @@ const (
 	ContextKeyChannelIsMultiKey        ContextKey = "channel_is_multi_key"
 	ContextKeyChannelMultiKeyIndex     ContextKey = "channel_multi_key_index"
 	ContextKeyChannelKey               ContextKey = "channel_key"
+	ContextKeyChannelTokenLimit        ContextKey = "channel_token_limit"
 
 	/* user related keys */
 	ContextKeyUserId      ContextKey = "id"

middleware/distributor.go

@@ -267,6 +267,7 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
 	common.SetContextKey(c, constant.ContextKeyChannelAutoBan, channel.GetAutoBan())
 	common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
 	common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
+	common.SetContextKey(c, constant.ContextKeyChannelTokenLimit, channel.GetTokenLimit())
 
 	key, index, newAPIError := channel.GetNextEnabledKey()
 	if newAPIError != nil {

model/channel.go

@@ -44,6 +44,7 @@ type Channel struct {
 	Tag               *string `json:"tag" gorm:"index"`
 	Setting           *string `json:"setting" gorm:"type:text"` // 渠道额外设置
 	ParamOverride     *string `json:"param_override" gorm:"type:text"`
+	TokenLimit        *int    `json:"token_limit" gorm:"default:0"`
 	// add after v0.8.5
 	ChannelInfo ChannelInfo `json:"channel_info" gorm:"type:json"`
 }
@@ -397,6 +398,14 @@ func (channel *Channel) GetStatusCodeMapping() string {
 	return *channel.StatusCodeMapping
 }
 
+func (channel *Channel) GetTokenLimit() int {
+	if channel.TokenLimit == nil || *channel.TokenLimit <= 0 {
+		return 0
+	}
+	return *channel.TokenLimit
+}
+
+
 func (channel *Channel) Insert() error {
 	var err error
 	err = DB.Create(channel).Error

relay/common/relay_info.go

@@ -62,6 +62,7 @@ type ResponsesUsageInfo struct {
 type RelayInfo struct {
 	ChannelType       int
 	ChannelId         int
+	ChannelTokenLimit int
 	TokenId           int
 	TokenKey          string
 	UserId            int
@@ -215,6 +216,7 @@ func GenRelayInfoImage(c *gin.Context) *RelayInfo {
 func GenRelayInfo(c *gin.Context) *RelayInfo {
 	channelType := common.GetContextKeyInt(c, constant.ContextKeyChannelType)
 	channelId := common.GetContextKeyInt(c, constant.ContextKeyChannelId)
+	channelTokenLimit := common.GetContextKeyInt(c, constant.ContextKeyChannelTokenLimit)
 	paramOverride := common.GetContextKeyStringMap(c, constant.ContextKeyChannelParamOverride)
 
 	tokenId := common.GetContextKeyInt(c, constant.ContextKeyTokenId)
@@ -235,6 +237,7 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 		RequestURLPath:    c.Request.URL.String(),
 		ChannelType:       channelType,
 		ChannelId:         channelId,
+		ChannelTokenLimit: channelTokenLimit,
 		TokenId:           tokenId,
 		TokenKey:          tokenKey,
 		UserId:            userId,

relay/relay-text.go

@@ -127,6 +127,11 @@ func TextHelper(c *gin.Context) (newAPIError *types.NewAPIError) {
 		c.Set("prompt_tokens", promptTokens)
 	}
 
+	err = checkPromptTokensInBotChannel(promptTokens, relayInfo)
+	if err != nil {
+		return types.NewError(err, types.ErrorCodePromptTokensTooLarge)
+	}
+
 	priceData, err := helper.ModelPriceHelper(c, relayInfo, promptTokens, int(math.Max(float64(textRequest.MaxTokens), float64(textRequest.MaxCompletionTokens))))
 	if err != nil {
 		return types.NewError(err, types.ErrorCodeModelPriceError)
@@ -261,6 +266,13 @@ func getPromptTokens(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.Re
 	return promptTokens, err
 }
 
+func checkPromptTokensInBotChannel(promptTokens int, info *relaycommon.RelayInfo) error {
+	if info.ChannelTokenLimit > 0 && promptTokens > info.ChannelTokenLimit {
+		return fmt.Errorf("prompt tokens (%d) is greater than channel token limit (%d)", promptTokens, info.ChannelTokenLimit)
+	}
+	return nil
+}
+
 func checkRequestSensitive(textRequest *dto.GeneralOpenAIRequest, info *relaycommon.RelayInfo) ([]string, error) {
 	var err error
 	var words []string

types/error.go

@@ -38,6 +38,7 @@ const (
 
 	// new api error
 	ErrorCodeCountTokenFailed  ErrorCode = "count_token_failed"
+	ErrorCodePromptTokensTooLarge ErrorCode = "prompt_tokens_too_large"
 	ErrorCodeModelPriceError   ErrorCode = "model_price_error"
 	ErrorCodeInvalidApiType    ErrorCode = "invalid_api_type"
 	ErrorCodeJsonMarshalFailed ErrorCode = "json_marshal_failed"

web/src/pages/Channel/EditChannel.js

@@ -102,6 +102,7 @@ const EditChannel = (props) => {
     groups: ['default'],
     priority: 0,
     weight: 0,
+    token_limit: 0,
     tag: '',
     multi_key_mode: 'random',
   };
@@ -1371,6 +1372,16 @@ const EditChannel = (props) => {
                         style={{ width: '100%' }}
                       />
                     </Col>
+                    <Col span={12}>
+                      <Form.InputNumber
+                        field='token_limit'
+                        label={t('最大上下文')}
+                        placeholder={t('最大上下文')}
+                        min={0}
+                        onNumberChange={(value) => handleInputChange('token_limit', value)}
+                        style={{ width: '100%' }}
+                      />
+                    </Col>
                   </Row>
 
                   <Form.Switch