feat(claude): forward Claude count_tokens to AWS Bedrock

Signed-off-by: B1F030 <b1fzhang@gmail.com>

Author: B1F030

constant/context_key.go

@@ -37,6 +37,7 @@ const (
 	ContextKeyChannelIsMultiKey        ContextKey = "channel_is_multi_key"
 	ContextKeyChannelMultiKeyIndex     ContextKey = "channel_multi_key_index"
 	ContextKeyChannelKey               ContextKey = "channel_key"
+	ContextKeyAllowedChannelTypes      ContextKey = "allowed_channel_types"
 
 	ContextKeyAutoGroup           ContextKey = "auto_group"
 	ContextKeyAutoGroupIndex      ContextKey = "auto_group_index"

controller/relay.go

@@ -123,6 +123,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		return
 	}
 
+	if relayFormat == types.RelayFormatClaude && relayInfo.RelayMode == relayconstant.RelayModeClaudeCountTokens {
+		newAPIError = relayClaudeCountTokens(c, relayInfo, request)
+		return
+	}
+
 	needSensitiveCheck := setting.ShouldCheckPromptSensitive()
 	needCountToken := constant.CountToken
 	// Avoid building huge CombineText (strings.Join) when token counting and sensitive check are both disabled.
@@ -179,10 +184,11 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 	}()
 
 	retryParam := &service.RetryParam{
-		Ctx:        c,
-		TokenGroup: relayInfo.TokenGroup,
-		ModelName:  relayInfo.OriginModelName,
-		Retry:      common.GetPointer(0),
+		Ctx:                 c,
+		TokenGroup:          relayInfo.TokenGroup,
+		ModelName:           relayInfo.OriginModelName,
+		Retry:               common.GetPointer(0),
+		AllowedChannelTypes: service.GetAllowedChannelTypes(c),
 	}
 	relayInfo.RetryIndex = 0
 	relayInfo.LastError = nil
@@ -247,6 +253,72 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 	}
 }
 
+func relayClaudeCountTokens(c *gin.Context, relayInfo *relaycommon.RelayInfo, request dto.Request) *types.NewAPIError {
+	if setting.ShouldCheckPromptSensitive() {
+		meta := request.GetTokenCountMeta()
+		if meta != nil {
+			contains, words := service.CheckSensitiveText(meta.CombineText)
+			if contains {
+				logger.LogWarn(c, fmt.Sprintf("user sensitive words detected: %s", strings.Join(words, ", ")))
+				return types.NewError(errors.New("sensitive words detected"), types.ErrorCodeSensitiveWordsDetected)
+			}
+		}
+	}
+
+	retryParam := &service.RetryParam{
+		Ctx:                 c,
+		TokenGroup:          relayInfo.TokenGroup,
+		ModelName:           relayInfo.OriginModelName,
+		Retry:               common.GetPointer(0),
+		AllowedChannelTypes: service.GetAllowedChannelTypes(c),
+	}
+	relayInfo.RetryIndex = 0
+	relayInfo.LastError = nil
+
+	var newAPIError *types.NewAPIError
+	for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
+		relayInfo.RetryIndex = retryParam.GetRetry()
+		channel, channelErr := getChannel(c, relayInfo, retryParam)
+		if channelErr != nil {
+			logger.LogError(c, channelErr.Error())
+			newAPIError = channelErr
+			break
+		}
+
+		addUsedChannel(c, channel.Id)
+		bodyStorage, bodyErr := common.GetBodyStorage(c)
+		if bodyErr != nil {
+			if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
+				newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
+			} else {
+				newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry())
+			}
+			break
+		}
+		c.Request.Body = io.NopCloser(bodyStorage)
+
+		newAPIError = relay.ClaudeCountTokensHelper(c, relayInfo)
+		if newAPIError == nil {
+			relayInfo.LastError = nil
+			return nil
+		}
+
+		relayInfo.LastError = newAPIError
+		processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
+
+		if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
+			break
+		}
+	}
+
+	useChannel := c.GetStringSlice("use_channel")
+	if len(useChannel) > 1 {
+		retryLogStr := fmt.Sprintf("重试：%s", strings.Trim(strings.Join(strings.Fields(fmt.Sprint(useChannel)), "->"), "[]"))
+		logger.LogInfo(c, retryLogStr)
+	}
+	return newAPIError
+}
+
 var upgrader = websocket.Upgrader{
 	Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
 	CheckOrigin: func(r *http.Request) bool {
@@ -296,9 +368,13 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
 		if !autoBan {
 			autoBanInt = 0
 		}
+		channelType := c.GetInt("channel_type")
+		if !service.IsChannelTypeAllowed(channelType, retryParam.AllowedChannelTypes) {
+			return nil, types.NewErrorWithStatusCode(fmt.Errorf("channel type %d is not allowed for this route", channelType), types.ErrorCodeGetChannelFailed, http.StatusServiceUnavailable, types.ErrOptionWithSkipRetry())
+		}
 		return &model.Channel{
 			Id:      c.GetInt("channel_id"),
-			Type:    c.GetInt("channel_type"),
+			Type:    channelType,
 			Name:    c.GetString("channel_name"),
 			AutoBan: &autoBanInt,
 		}, nil

i18n/keys.go

@@ -309,6 +309,7 @@ const (
 	MsgDistributorInvalidChannelId        = "distributor.invalid_channel_id"
 	MsgDistributorChannelDisabled         = "distributor.channel_disabled"
 	MsgDistributorAffinityChannelDisabled = "distributor.affinity_channel_disabled"
+	MsgDistributorChannelTypeNotAllowed   = "distributor.channel_type_not_allowed"
 	MsgDistributorTokenNoModelAccess      = "distributor.token_no_model_access"
 	MsgDistributorTokenModelForbidden     = "distributor.token_model_forbidden"
 	MsgDistributorModelNameRequired       = "distributor.model_name_required"

i18n/locales/en.yaml

@@ -259,6 +259,7 @@ distributor.invalid_request: "Invalid request: {{.Error}}"
 distributor.invalid_channel_id: "Invalid channel ID"
 distributor.channel_disabled: "This channel has been disabled"
 distributor.affinity_channel_disabled: "The channel selected by channel affinity has been disabled, and retry was stopped by rule. Please contact the administrator"
+distributor.channel_type_not_allowed: "Channel type is not allowed for this route"
 distributor.token_no_model_access: "This token has no access to any models"
 distributor.token_model_forbidden: "This token has no access to model {{.Model}}"
 distributor.model_name_required: "Model name not specified, model name cannot be empty"

i18n/locales/zh-CN.yaml

@@ -260,6 +260,7 @@ distributor.invalid_request: "无效的请求，{{.Error}}"
 distributor.invalid_channel_id: "无效的渠道 Id"
 distributor.channel_disabled: "该渠道已被禁用"
 distributor.affinity_channel_disabled: "渠道亲和性命中的渠道已被禁用，已按规则停止重试，请联系管理员处理"
+distributor.channel_type_not_allowed: "该渠道类型不允许用于此路由"
 distributor.token_no_model_access: "该令牌无权访问任何模型"
 distributor.token_model_forbidden: "该令牌无权访问模型 {{.Model}}"
 distributor.model_name_required: "未指定模型名称，模型名称不能为空"

i18n/locales/zh-TW.yaml

@@ -260,6 +260,7 @@ distributor.invalid_request: "無效的請求，{{.Error}}"
 distributor.invalid_channel_id: "無效的管道 Id"
 distributor.channel_disabled: "該管道已被禁用"
 distributor.affinity_channel_disabled: "管道親和性命中的管道已被禁用，已按規則停止重試，請聯絡管理員處理"
+distributor.channel_type_not_allowed: "該管道類型不允許用於此路由"
 distributor.token_no_model_access: "該令牌無權存取任何模型"
 distributor.token_model_forbidden: "該令牌無權存取模型 {{.Model}}"
 distributor.model_name_required: "未指定模型名稱，模型名稱不能為空"

middleware/channel_type_limit_test.go

@@ -0,0 +1,30 @@
+package middleware
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/QuantumNous/new-api/service"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func TestLimitChannelTypesStoresAllowedTypes(t *testing.T) {
+	gin.SetMode(gin.TestMode)
+
+	router := gin.New()
+	router.Use(LimitChannelTypes(constant.ChannelTypeAws))
+	router.GET("/test", func(c *gin.Context) {
+		require.True(t, service.IsChannelTypeAllowed(constant.ChannelTypeAws, service.GetAllowedChannelTypes(c)))
+		require.False(t, service.IsChannelTypeAllowed(constant.ChannelTypeOpenAI, service.GetAllowedChannelTypes(c)))
+		c.Status(http.StatusNoContent)
+	})
+
+	recorder := httptest.NewRecorder()
+	request := httptest.NewRequest(http.MethodGet, "/test", nil)
+	router.ServeHTTP(recorder, request)
+
+	require.Equal(t, http.StatusNoContent, recorder.Code)
+}

middleware/distributor.go

@@ -29,10 +29,18 @@ type ModelRequest struct {
 	Group string `json:"group,omitempty"`
 }
 
+func LimitChannelTypes(channelTypes ...int) func(c *gin.Context) {
+	return func(c *gin.Context) {
+		common.SetContextKey(c, constant.ContextKeyAllowedChannelTypes, channelTypes)
+		c.Next()
+	}
+}
+
 func Distribute() func(c *gin.Context) {
 	return func(c *gin.Context) {
 		var channel *model.Channel
 		channelId, ok := common.GetContextKey(c, constant.ContextKeyTokenSpecificChannelId)
+		allowedChannelTypes := service.GetAllowedChannelTypes(c)
 		modelRequest, shouldSelectChannel, err := getModelRequest(c)
 		if err != nil {
 			abortWithOpenAiMessage(c, http.StatusBadRequest, i18n.T(c, i18n.MsgDistributorInvalidRequest, map[string]any{"Error": err.Error()}))
@@ -53,6 +61,10 @@ func Distribute() func(c *gin.Context) {
 				abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelDisabled))
 				return
 			}
+			if !service.IsChannelTypeAllowed(channel.Type, allowedChannelTypes) {
+				abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorChannelTypeNotAllowed))
+				return
+			}
 		} else {
 			// Select a channel for the user
 			// check token model mapping
@@ -104,8 +116,16 @@ func Distribute() func(c *gin.Context) {
 				if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
 					affinityUsable := false
 					preferred, err := model.CacheGetChannel(preferredChannelID)
-					if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
-						if usingGroup == "auto" {
+					if err == nil && preferred != nil {
+						if preferred.Status != common.ChannelStatusEnabled {
+							if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
+								abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorAffinityChannelDisabled))
+								return
+							}
+						} else if !service.IsChannelTypeAllowed(preferred.Type, allowedChannelTypes) {
+							// Affinity is only a preference. If it points to a disallowed channel type,
+							// keep selecting from the normal candidate pool.
+						} else if usingGroup == "auto" {
 							userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
 							autoGroups := service.GetUserAutoGroup(userGroup)
 							for _, g := range autoGroups {
@@ -132,10 +152,11 @@ func Distribute() func(c *gin.Context) {
 
 				if channel == nil {
 					channel, selectGroup, err = service.CacheGetRandomSatisfiedChannel(&service.RetryParam{
-						Ctx:        c,
-						ModelName:  modelRequest.Model,
-						TokenGroup: usingGroup,
-						Retry:      common.GetPointer(0),
+						Ctx:                 c,
+						ModelName:           modelRequest.Model,
+						TokenGroup:          usingGroup,
+						Retry:               common.GetPointer(0),
+						AllowedChannelTypes: allowedChannelTypes,
 					})
 					if err != nil {
 						showGroup := usingGroup

model/ability.go

@@ -3,6 +3,7 @@ package model
 import (
 	"errors"
 	"fmt"
+	"sort"
 	"strings"
 	"sync"
 
@@ -104,6 +105,14 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
 }
 
 func GetChannel(group string, model string, retry int) (*Channel, error) {
+	return GetChannelWithChannelTypes(group, model, retry, nil)
+}
+
+func GetChannelWithChannelTypes(group string, model string, retry int, allowedChannelTypes []int) (*Channel, error) {
+	if len(allowedChannelTypes) > 0 {
+		return getChannelWithChannelTypes(group, model, retry, allowedChannelTypes)
+	}
+
 	var abilities []Ability
 
 	var err error = nil
@@ -143,6 +152,74 @@ func GetChannel(group string, model string, retry int) (*Channel, error) {
 	return &channel, err
 }
 
+func getChannelWithChannelTypes(group string, model string, retry int, allowedChannelTypes []int) (*Channel, error) {
+	var abilities []AbilityWithChannel
+	err := DB.Table("abilities").
+		Select("abilities.*, channels.type as channel_type").
+		Joins("left join channels on abilities.channel_id = channels.id").
+		Where("abilities."+commonGroupCol+" = ? and abilities.model = ? and abilities.enabled = ?", group, model, true).
+		Where("channels.type IN ?", allowedChannelTypes).
+		Scan(&abilities).Error
+	if err != nil {
+		return nil, err
+	}
+	if len(abilities) == 0 {
+		return nil, nil
+	}
+
+	uniquePriorities := make(map[int]bool)
+	for _, ability := range abilities {
+		priority := int64(0)
+		if ability.Priority != nil {
+			priority = *ability.Priority
+		}
+		uniquePriorities[int(priority)] = true
+	}
+	var sortedUniquePriorities []int
+	for priority := range uniquePriorities {
+		sortedUniquePriorities = append(sortedUniquePriorities, priority)
+	}
+	sort.Sort(sort.Reverse(sort.IntSlice(sortedUniquePriorities)))
+
+	if retry >= len(uniquePriorities) {
+		retry = len(uniquePriorities) - 1
+	}
+	targetPriority := int64(sortedUniquePriorities[retry])
+
+	targetAbilities := make([]AbilityWithChannel, 0)
+	weightSum := uint(0)
+	for _, ability := range abilities {
+		priority := int64(0)
+		if ability.Priority != nil {
+			priority = *ability.Priority
+		}
+		if priority != targetPriority {
+			continue
+		}
+		targetAbilities = append(targetAbilities, ability)
+		weightSum += ability.Weight + 10
+	}
+	if len(targetAbilities) == 0 {
+		return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority))
+	}
+
+	weight := common.GetRandomInt(int(weightSum))
+	channelID := 0
+	for _, ability := range targetAbilities {
+		weight -= int(ability.Weight) + 10
+		if weight <= 0 {
+			channelID = ability.ChannelId
+			break
+		}
+	}
+	if channelID == 0 {
+		return nil, errors.New("channel not found")
+	}
+	channel := Channel{}
+	err = DB.First(&channel, "id = ?", channelID).Error
+	return &channel, err
+}
+
 func (channel *Channel) AddAbilities(tx *gorm.DB) error {
 	models_ := strings.Split(channel.Models, ",")
 	groups_ := strings.Split(channel.Group, ",")