diff --git a/constant/context_key.go b/constant/context_key.go index b856bc3dda1..0dbc224bf04 100644 --- a/constant/context_key.go +++ b/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" diff --git a/controller/relay.go b/controller/relay.go index 1d14dcc6f88..dd332fbff63 100644 --- a/controller/relay.go +++ b/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 diff --git a/i18n/keys.go b/i18n/keys.go index 6cf5c1bdc33..4ed0dc9a9b7 100644 --- a/i18n/keys.go +++ b/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" diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 198aa27412b..9031b8005eb 100644 --- a/i18n/locales/en.yaml +++ b/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" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index 45a10a584e5..b6064d69c5a 100644 --- a/i18n/locales/zh-CN.yaml +++ b/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: "未指定模型名称，模型名称不能为空" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index ad5a6eae322..ce9da7f6591 100644 --- a/i18n/locales/zh-TW.yaml +++ b/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: "未指定模型名稱，模型名稱不能為空" diff --git a/middleware/channel_type_limit_test.go b/middleware/channel_type_limit_test.go new file mode 100644 index 00000000000..5d373b9c85b --- /dev/null +++ b/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) +} diff --git a/middleware/distributor.go b/middleware/distributor.go index 258aebb5703..82070018cfc 100644 --- a/middleware/distributor.go +++ b/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 diff --git a/model/ability.go b/model/ability.go index 1d7c53fa580..ffec9220095 100644 --- a/model/ability.go +++ b/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, ",") diff --git a/model/channel_cache.go b/model/channel_cache.go index 03740d2cd3a..f118a87664d 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -95,9 +95,13 @@ func SyncChannelCache(frequency int) { } func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, error) { + return GetRandomSatisfiedChannelWithTypes(group, model, retry, nil) +} + +func GetRandomSatisfiedChannelWithTypes(group string, model string, retry int, allowedChannelTypes []int) (*Channel, error) { // if memory cache is disabled, get channel directly from database if !common.MemoryCacheEnabled { - return GetChannel(group, model, retry) + return GetChannelWithChannelTypes(group, model, retry, allowedChannelTypes) } channelSyncLock.RLock() @@ -116,20 +120,21 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, return nil, nil } - if len(channels) == 1 { - if channel, ok := channelsIDM[channels[0]]; ok { - return channel, nil - } - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channels[0]) + candidateChannels, err := getCachedChannelsByTypes(channels, allowedChannelTypes) + if err != nil { + return nil, err + } + if len(candidateChannels) == 0 { + return nil, nil + } + + if len(candidateChannels) == 1 { + return candidateChannels[0], nil } uniquePriorities := make(map[int]bool) - for _, channelId := range channels { - if channel, ok := channelsIDM[channelId]; ok { - uniquePriorities[int(channel.GetPriority())] = true - } else { - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) - } + for _, channel := range candidateChannels { + uniquePriorities[int(channel.GetPriority())] = true } var sortedUniquePriorities []int for priority := range uniquePriorities { @@ -145,14 +150,10 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, // get the priority for the given retry number var sumWeight = 0 var targetChannels []*Channel - for _, channelId := range channels { - if channel, ok := channelsIDM[channelId]; ok { - if channel.GetPriority() == targetPriority { - sumWeight += channel.GetWeight() - targetChannels = append(targetChannels, channel) - } - } else { - return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelId) + for _, channel := range candidateChannels { + if channel.GetPriority() == targetPriority { + sumWeight += channel.GetWeight() + targetChannels = append(targetChannels, channel) } } @@ -191,6 +192,33 @@ func GetRandomSatisfiedChannel(group string, model string, retry int) (*Channel, return nil, errors.New("channel not found") } +func getCachedChannelsByTypes(channelIDs []int, allowedChannelTypes []int) ([]*Channel, error) { + candidateChannels := make([]*Channel, 0, len(channelIDs)) + for _, channelID := range channelIDs { + channel, ok := channelsIDM[channelID] + if !ok { + return nil, fmt.Errorf("数据库一致性错误，渠道# %d 不存在，请联系管理员修复", channelID) + } + if !isChannelTypeAllowed(channel.Type, allowedChannelTypes) { + continue + } + candidateChannels = append(candidateChannels, channel) + } + return candidateChannels, nil +} + +func isChannelTypeAllowed(channelType int, allowedChannelTypes []int) bool { + if len(allowedChannelTypes) == 0 { + return true + } + for _, allowed := range allowedChannelTypes { + if channelType == allowed { + return true + } + } + return false +} + func CacheGetChannel(id int) (*Channel, error) { if !common.MemoryCacheEnabled { return GetChannelById(id, true) diff --git a/model/channel_cache_test.go b/model/channel_cache_test.go new file mode 100644 index 00000000000..cc95574d744 --- /dev/null +++ b/model/channel_cache_test.go @@ -0,0 +1,69 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/stretchr/testify/require" +) + +func TestGetRandomSatisfiedChannelWithTypesFiltersCandidates(t *testing.T) { + restore := installTestChannelCache(t) + defer restore() + + openAIWeight := uint(100) + awsWeight := uint(100) + priority := int64(10) + group2model2channels = map[string]map[string][]int{ + "default": { + "claude-3-5-sonnet-20240620": {1, 2}, + }, + } + channelsIDM = map[int]*Channel{ + 1: {Id: 1, Type: constant.ChannelTypeOpenAI, Weight: &openAIWeight, Priority: &priority}, + 2: {Id: 2, Type: constant.ChannelTypeAws, Weight: &awsWeight, Priority: &priority}, + } + + channel, err := GetRandomSatisfiedChannelWithTypes("default", "claude-3-5-sonnet-20240620", 0, []int{constant.ChannelTypeAws}) + + require.NoError(t, err) + require.NotNil(t, channel) + require.Equal(t, 2, channel.Id) +} + +func TestGetRandomSatisfiedChannelWithTypesReturnsNilWhenNoCandidateMatches(t *testing.T) { + restore := installTestChannelCache(t) + defer restore() + + weight := uint(100) + priority := int64(10) + group2model2channels = map[string]map[string][]int{ + "default": { + "claude-3-5-sonnet-20240620": {1}, + }, + } + channelsIDM = map[int]*Channel{ + 1: {Id: 1, Type: constant.ChannelTypeOpenAI, Weight: &weight, Priority: &priority}, + } + + channel, err := GetRandomSatisfiedChannelWithTypes("default", "claude-3-5-sonnet-20240620", 0, []int{constant.ChannelTypeAws}) + + require.NoError(t, err) + require.Nil(t, channel) +} + +func installTestChannelCache(t *testing.T) func() { + t.Helper() + + oldMemoryCacheEnabled := common.MemoryCacheEnabled + oldGroup2Model2Channels := group2model2channels + oldChannelsIDM := channelsIDM + common.MemoryCacheEnabled = true + + return func() { + common.MemoryCacheEnabled = oldMemoryCacheEnabled + group2model2channels = oldGroup2Model2Channels + channelsIDM = oldChannelsIDM + } +} diff --git a/relay/channel/aws/relay-aws.go b/relay/channel/aws/relay-aws.go index 1f6ff7e6926..9645fd34f2f 100644 --- a/relay/channel/aws/relay-aws.go +++ b/relay/channel/aws/relay-aws.go @@ -170,6 +170,92 @@ func doAwsClientRequest(c *gin.Context, info *relaycommon.RelayInfo, a *Adaptor, } } +func CountClaudeTokens(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (int, *types.NewAPIError) { + awsReq, awsCli, err := buildAwsCountTokensInput(c, info, requestBody) + if err != nil { + var newAPIError *types.NewAPIError + if errors.As(err, &newAPIError) { + return 0, newAPIError + } + return 0, types.NewErrorWithStatusCode(err, types.ErrorCodeBadRequestBody, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + + ctx, cancel := newAwsInvokeContext() + defer cancel() + + awsResp, err := awsCli.CountTokens(ctx, awsReq) + if err != nil { + statusCode := getAwsErrorStatusCode(err) + return 0, types.NewOpenAIError(errors.Wrap(err, "CountTokens"), types.ErrorCodeAwsInvokeError, statusCode) + } + if awsResp == nil || awsResp.InputTokens == nil { + return 0, types.NewErrorWithStatusCode(errors.New("empty CountTokens response"), types.ErrorCodeBadResponseBody, http.StatusInternalServerError) + } + return int(*awsResp.InputTokens), nil +} + +func buildAwsCountTokensInput(c *gin.Context, info *relaycommon.RelayInfo, requestBody io.Reader) (*bedrockruntime.CountTokensInput, *bedrockruntime.Client, error) { + awsCli, err := newAwsClient(c, info) + if err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeChannelAwsClientError) + } + + awsModelId := getAwsCountTokensModelID(info.UpstreamModelName) + if !isAwsClaudeModel(awsModelId) { + return nil, nil, fmt.Errorf("AWS Bedrock count_tokens only supports Claude models, got %s", info.UpstreamModelName) + } + + requestHeader := http.Header{} + adaptor := &Adaptor{} + if err := adaptor.SetupRequestHeader(c, &requestHeader, info); err != nil { + return nil, nil, types.NewError(err, types.ErrorCodeChannelAwsClientError) + } + headerOverride, err := channel.ResolveHeaderOverride(info, c) + if err != nil { + return nil, nil, err + } + for key, value := range headerOverride { + requestHeader.Set(key, value) + } + + awsClaudeReq, err := formatRequest(requestBody, requestHeader) + if err != nil { + return nil, nil, errors.Wrap(err, "format aws count_tokens request fail") + } + body, err := buildAwsRequestBody(c, info, awsClaudeReq) + if err != nil { + return nil, nil, errors.Wrap(err, "marshal aws count_tokens request fail") + } + + return &bedrockruntime.CountTokensInput{ + ModelId: aws.String(awsModelId), + Input: &bedrockruntimeTypes.CountTokensInputMemberInvokeModel{ + Value: bedrockruntimeTypes.InvokeModelTokensRequest{ + Body: body, + }, + }, + }, awsCli, nil +} + +func isAwsClaudeModel(modelId string) bool { + modelId = strings.ToLower(modelId) + return strings.Contains(modelId, "claude") && !isNovaModel(modelId) +} + +func getAwsCountTokensModelID(requestModel string) string { + return normalizeAwsFoundationModelID(getAwsModelID(requestModel)) +} + +func normalizeAwsFoundationModelID(modelId string) string { + lowerModelId := strings.ToLower(modelId) + for _, prefix := range []string{"global.", "us.", "eu.", "apac.", "au."} { + if strings.HasPrefix(lowerModelId, prefix+"anthropic.") || strings.HasPrefix(lowerModelId, prefix+"amazon.") { + return modelId[len(prefix):] + } + } + return modelId +} + // buildAwsRequestBody prepares the payload for AWS requests, applying passthrough rules when enabled. func buildAwsRequestBody(c *gin.Context, info *relaycommon.RelayInfo, awsClaudeReq any) ([]byte, error) { if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled { diff --git a/relay/channel/aws/relay_aws_test.go b/relay/channel/aws/relay_aws_test.go index 92745ff4092..8eaa09c6f72 100644 --- a/relay/channel/aws/relay_aws_test.go +++ b/relay/channel/aws/relay_aws_test.go @@ -8,7 +8,9 @@ import ( "github.com/QuantumNous/new-api/common" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/types" "github.com/aws/aws-sdk-go-v2/service/bedrockruntime" + bedrockruntimeTypes "github.com/aws/aws-sdk-go-v2/service/bedrockruntime/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -53,3 +55,106 @@ func TestDoAwsClientRequest_AppliesRuntimeHeaderOverrideToAnthropicBeta(t *testi require.True(t, ok) require.Equal(t, []any{"computer-use-2025-01-24"}, values) } + +func TestBuildAwsCountTokensInputUsesInvokeModelBody(t *testing.T) { + t.Parallel() + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + info := &relaycommon.RelayInfo{ + OriginModelName: "claude-3-5-sonnet-20240620", + ChannelMeta: &relaycommon.ChannelMeta{ + ApiKey: "access-key|secret-key|us-east-1", + UpstreamModelName: "claude-3-5-sonnet-20240620", + }, + } + + requestBody := bytes.NewBufferString(`{"model":"claude-3-5-sonnet-20240620","messages":[{"role":"user","content":"hello"}]}`) + input, _, err := buildAwsCountTokensInput(ctx, info, requestBody) + require.NoError(t, err) + + require.Equal(t, "anthropic.claude-3-5-sonnet-20240620-v1:0", *input.ModelId) + invokeModel, ok := input.Input.(*bedrockruntimeTypes.CountTokensInputMemberInvokeModel) + require.True(t, ok) + + var payload map[string]any + require.NoError(t, common.Unmarshal(invokeModel.Value.Body, &payload)) + require.Equal(t, "bedrock-2023-05-31", payload["anthropic_version"]) + require.NotContains(t, payload, "model") + require.NotEmpty(t, payload["messages"]) +} + +func TestBuildAwsCountTokensInputNormalizesGlobalInferenceProfile(t *testing.T) { + t.Parallel() + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + info := &relaycommon.RelayInfo{ + OriginModelName: "global.anthropic.claude-opus-4-6-v1", + ChannelMeta: &relaycommon.ChannelMeta{ + ApiKey: "access-key|secret-key|ap-northeast-1", + UpstreamModelName: "global.anthropic.claude-opus-4-6-v1", + }, + } + + requestBody := bytes.NewBufferString(`{"model":"global.anthropic.claude-opus-4-6-v1","messages":[{"role":"user","content":"hello"}]}`) + input, _, err := buildAwsCountTokensInput(ctx, info, requestBody) + require.NoError(t, err) + + require.Equal(t, "anthropic.claude-opus-4-6-v1", *input.ModelId) + _, ok := input.Input.(*bedrockruntimeTypes.CountTokensInputMemberInvokeModel) + require.True(t, ok) +} + +func TestBuildAwsCountTokensInputRejectsNonClaudeModel(t *testing.T) { + t.Parallel() + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + info := &relaycommon.RelayInfo{ + OriginModelName: "nova-micro-v1:0", + ChannelMeta: &relaycommon.ChannelMeta{ + ApiKey: "access-key|secret-key|us-east-1", + UpstreamModelName: "nova-micro-v1:0", + }, + } + + requestBody := bytes.NewBufferString(`{"model":"nova-micro-v1:0","messages":[{"role":"user","content":"hello"}]}`) + _, _, err := buildAwsCountTokensInput(ctx, info, requestBody) + + require.Error(t, err) + require.Contains(t, err.Error(), "only supports Claude") +} + +func TestCountClaudeTokensPreservesAwsClientError(t *testing.T) { + t.Parallel() + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/messages/count_tokens", nil) + + info := &relaycommon.RelayInfo{ + OriginModelName: "claude-3-5-sonnet-20240620", + ChannelMeta: &relaycommon.ChannelMeta{ + ApiKey: "invalid-secret", + UpstreamModelName: "claude-3-5-sonnet-20240620", + }, + } + + _, err := CountClaudeTokens(ctx, info, bytes.NewBufferString(`{"model":"claude-3-5-sonnet-20240620","messages":[{"role":"user","content":"hello"}]}`)) + + require.NotNil(t, err) + require.Equal(t, types.ErrorCodeChannelAwsClientError, err.GetErrorCode()) + require.False(t, types.IsSkipRetryError(err)) + require.Equal(t, http.StatusInternalServerError, err.StatusCode) +} diff --git a/relay/claude_count_tokens.go b/relay/claude_count_tokens.go new file mode 100644 index 00000000000..854f26d094f --- /dev/null +++ b/relay/claude_count_tokens.go @@ -0,0 +1,94 @@ +package relay + +import ( + "bytes" + "fmt" + "io" + "net/http" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/relay/channel" + awsrelay "github.com/QuantumNous/new-api/relay/channel/aws" + relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/setting/model_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" +) + +func ClaudeCountTokensHelper(c *gin.Context, info *relaycommon.RelayInfo) *types.NewAPIError { + info.InitChannelMeta(c) + + if info.ChannelType != constant.ChannelTypeAws { + return types.NewErrorWithStatusCode(fmt.Errorf("count_tokens only supports AWS Bedrock Claude channels"), types.ErrorCodeInvalidApiType, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + + claudeReq, ok := info.Request.(*dto.ClaudeRequest) + if !ok { + return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected *dto.ClaudeRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + + request, prepareErr := prepareClaudeRequestForUpstream(c, info, claudeReq, false) + if prepareErr != nil { + return prepareErr + } + + adaptor := GetAdaptor(info.ApiType) + if adaptor == nil { + return types.NewError(fmt.Errorf("invalid api type: %d", info.ApiType), types.ErrorCodeInvalidApiType, types.ErrOptionWithSkipRetry()) + } + adaptor.Init(info) + + requestBody, buildErr := buildClaudeCountTokensRequestBody(c, info, adaptor, request) + if buildErr != nil { + return buildErr + } + + inputTokens, countErr := awsrelay.CountClaudeTokens(c, info, requestBody) + if countErr != nil { + return countErr + } + + c.JSON(http.StatusOK, gin.H{"input_tokens": inputTokens}) + return nil +} + +func buildClaudeCountTokensRequestBody(c *gin.Context, info *relaycommon.RelayInfo, adaptor channel.Adaptor, request *dto.ClaudeRequest) (io.Reader, *types.NewAPIError) { + if model_setting.GetGlobalSettings().PassThroughRequestEnabled || info.ChannelSetting.PassThroughBodyEnabled { + storage, err := common.GetBodyStorage(c) + if err != nil { + return nil, types.NewErrorWithStatusCode(err, types.ErrorCodeReadRequestBodyFailed, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) + } + return common.ReaderOnly(storage), nil + } + + convertedRequest, err := adaptor.ConvertClaudeRequest(c, info, request) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + relaycommon.AppendRequestConversionFromRequest(info, convertedRequest) + + jsonData, err := common.Marshal(convertedRequest) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + + jsonData, err = relaycommon.RemoveDisabledFields(jsonData, info.ChannelOtherSettings, info.ChannelSetting.PassThroughBodyEnabled) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeConvertRequestFailed, types.ErrOptionWithSkipRetry()) + } + + if len(info.ParamOverride) > 0 { + jsonData, err = relaycommon.ApplyParamOverrideWithRelayInfo(jsonData, info) + if err != nil { + return nil, newAPIErrorFromParamOverride(err) + } + } + + if common.DebugEnabled { + println("countTokensRequestBody: ", string(jsonData)) + } + return bytes.NewBuffer(jsonData), nil +} diff --git a/relay/claude_handler.go b/relay/claude_handler.go index 527363205a1..8f8b7d16e8b 100644 --- a/relay/claude_handler.go +++ b/relay/claude_handler.go @@ -31,14 +31,9 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ return types.NewErrorWithStatusCode(fmt.Errorf("invalid request type, expected *dto.ClaudeRequest, got %T", info.Request), types.ErrorCodeInvalidRequest, http.StatusBadRequest, types.ErrOptionWithSkipRetry()) } - request, err := common.DeepCopy(claudeReq) - if err != nil { - return types.NewError(fmt.Errorf("failed to copy request to ClaudeRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) - } - - err = helper.ModelMappedHelper(c, info, request) - if err != nil { - return types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry()) + request, prepareErr := prepareClaudeRequestForUpstream(c, info, claudeReq, true) + if prepareErr != nil { + return prepareErr } adaptor := GetAdaptor(info.ApiType) @@ -221,3 +216,99 @@ func ClaudeHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *typ service.PostTextConsumeQuota(c, info, usage.(*dto.Usage), nil) return nil } + +func prepareClaudeRequestForUpstream(c *gin.Context, info *relaycommon.RelayInfo, claudeReq *dto.ClaudeRequest, setDefaultMaxTokens bool) (*dto.ClaudeRequest, *types.NewAPIError) { + request, err := common.DeepCopy(claudeReq) + if err != nil { + return nil, types.NewError(fmt.Errorf("failed to copy request to ClaudeRequest: %w", err), types.ErrorCodeInvalidRequest, types.ErrOptionWithSkipRetry()) + } + + err = helper.ModelMappedHelper(c, info, request) + if err != nil { + return nil, types.NewError(err, types.ErrorCodeChannelModelMappedError, types.ErrOptionWithSkipRetry()) + } + + if request.MaxTokens == nil || *request.MaxTokens == 0 { + if setDefaultMaxTokens { + defaultMaxTokens := uint(model_setting.GetClaudeSettings().GetDefaultMaxTokens(request.Model)) + request.MaxTokens = &defaultMaxTokens + } + } + + if baseModel, effortLevel, ok := reasoning.TrimEffortSuffix(request.Model); ok && effortLevel != "" && + (strings.HasPrefix(request.Model, "claude-opus-4-6") || strings.HasPrefix(request.Model, "claude-opus-4-7")) { + request.Model = baseModel + request.Thinking = &dto.Thinking{ + Type: "adaptive", + } + request.OutputConfig = json.RawMessage(fmt.Sprintf(`{"effort":"%s"}`, effortLevel)) + if strings.HasPrefix(request.Model, "claude-opus-4-7") { + // Opus 4.7 rejects non-default temperature/top_p/top_k with 400 + // and defaults display to "omitted"; restore the 4.6 visible summary. + request.Thinking.Display = "summarized" + request.Temperature = nil + request.TopP = nil + request.TopK = nil + } else { + request.Temperature = common.GetPointer[float64](1.0) + } + info.UpstreamModelName = request.Model + } else if model_setting.GetClaudeSettings().ThinkingAdapterEnabled && + strings.HasSuffix(request.Model, "-thinking") { + if request.Thinking == nil { + baseModel := strings.TrimSuffix(request.Model, "-thinking") + if strings.HasPrefix(baseModel, "claude-opus-4-7") { + // Opus 4.7 rejects thinking.type="enabled"; use adaptive at high effort. + request.Thinking = &dto.Thinking{Type: "adaptive", Display: "summarized"} + request.OutputConfig = json.RawMessage(`{"effort":"high"}`) + request.Temperature = nil + request.TopP = nil + request.TopK = nil + } else { + // 因为BudgetTokens 必须大于1024 + if request.MaxTokens == nil || *request.MaxTokens < 1280 { + request.MaxTokens = common.GetPointer[uint](1280) + } + + // BudgetTokens 为 max_tokens 的 80% + request.Thinking = &dto.Thinking{ + Type: "enabled", + BudgetTokens: common.GetPointer[int](int(float64(*request.MaxTokens) * model_setting.GetClaudeSettings().ThinkingAdapterBudgetTokensPercentage)), + } + // TODO: 临时处理 + // https://docs.anthropic.com/en/docs/build-with-claude/extended-thinking#important-considerations-when-using-extended-thinking + request.Temperature = common.GetPointer[float64](1.0) + } + } + if !model_setting.ShouldPreserveThinkingSuffix(info.OriginModelName) { + request.Model = strings.TrimSuffix(request.Model, "-thinking") + } + info.UpstreamModelName = request.Model + } + + if info.ChannelSetting.SystemPrompt != "" { + if request.System == nil { + request.SetStringSystem(info.ChannelSetting.SystemPrompt) + } else if info.ChannelSetting.SystemPromptOverride { + common.SetContextKey(c, constant.ContextKeySystemPromptOverride, true) + if request.IsStringSystem() { + existing := strings.TrimSpace(request.GetStringSystem()) + if existing == "" { + request.SetStringSystem(info.ChannelSetting.SystemPrompt) + } else { + request.SetStringSystem(info.ChannelSetting.SystemPrompt + "\n" + existing) + } + } else { + systemContents := request.ParseSystem() + newSystem := dto.ClaudeMediaMessage{Type: dto.ContentTypeText} + newSystem.SetText(info.ChannelSetting.SystemPrompt) + if len(systemContents) == 0 { + request.System = []dto.ClaudeMediaMessage{newSystem} + } else { + request.System = append([]dto.ClaudeMediaMessage{newSystem}, systemContents...) + } + } + } + } + return request, nil +} diff --git a/relay/constant/relay_mode.go b/relay/constant/relay_mode.go index 25671567921..ea3a0b5c84e 100644 --- a/relay/constant/relay_mode.go +++ b/relay/constant/relay_mode.go @@ -52,6 +52,8 @@ const ( RelayModeGemini RelayModeResponsesCompact + + RelayModeClaudeCountTokens ) func Path2RelayMode(path string) int { @@ -76,6 +78,8 @@ func Path2RelayMode(path string) int { relayMode = RelayModeResponsesCompact } else if strings.HasPrefix(path, "/v1/responses") { relayMode = RelayModeResponses + } else if strings.HasPrefix(path, "/v1/messages/count_tokens") { + relayMode = RelayModeClaudeCountTokens } else if strings.HasPrefix(path, "/v1/audio/speech") { relayMode = RelayModeAudioSpeech } else if strings.HasPrefix(path, "/v1/audio/transcriptions") { diff --git a/relay/constant/relay_mode_test.go b/relay/constant/relay_mode_test.go new file mode 100644 index 00000000000..279d1f13ff6 --- /dev/null +++ b/relay/constant/relay_mode_test.go @@ -0,0 +1,11 @@ +package constant + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestPath2RelayModeClaudeCountTokens(t *testing.T) { + require.Equal(t, RelayModeClaudeCountTokens, Path2RelayMode("/v1/messages/count_tokens")) +} diff --git a/router/relay-router.go b/router/relay-router.go index 17a13cad7fd..12246e8ced0 100644 --- a/router/relay-router.go +++ b/router/relay-router.go @@ -88,6 +88,19 @@ func SetRelayRouter(router *gin.Engine) { httpRouter.POST("/messages", func(c *gin.Context) { controller.Relay(c, types.RelayFormatClaude) }) + } + { + claudeCountTokensRouter := relayV1Router.Group("") + claudeCountTokensRouter.Use(middleware.LimitChannelTypes(constant.ChannelTypeAws)) + claudeCountTokensRouter.Use(middleware.Distribute()) + claudeCountTokensRouter.POST("/messages/count_tokens", func(c *gin.Context) { + controller.Relay(c, types.RelayFormatClaude) + }) + } + { + //http router + httpRouter := relayV1Router.Group("") + httpRouter.Use(middleware.Distribute()) // chat related routes httpRouter.POST("/completions", func(c *gin.Context) { diff --git a/service/channel_select.go b/service/channel_select.go index a3710ef8cec..64e4fadc8ac 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -12,11 +12,12 @@ import ( ) type RetryParam struct { - Ctx *gin.Context - TokenGroup string - ModelName string - Retry *int - resetNextTry bool + Ctx *gin.Context + TokenGroup string + ModelName string + Retry *int + AllowedChannelTypes []int + resetNextTry bool } func (p *RetryParam) GetRetry() int { @@ -45,6 +46,29 @@ func (p *RetryParam) ResetRetryNextTry() { p.resetNextTry = true } +func GetAllowedChannelTypes(c *gin.Context) []int { + if c == nil { + return nil + } + channelTypes, ok := common.GetContextKeyType[[]int](c, constant.ContextKeyAllowedChannelTypes) + if !ok { + return nil + } + return channelTypes +} + +func IsChannelTypeAllowed(channelType int, allowedChannelTypes []int) bool { + if len(allowedChannelTypes) == 0 { + return true + } + for _, allowed := range allowedChannelTypes { + if channelType == allowed { + return true + } + } + return false +} + // CacheGetRandomSatisfiedChannel tries to get a random channel that satisfies the requirements. // 尝试获取一个满足要求的随机渠道。 // @@ -85,6 +109,10 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, var err error selectGroup := param.TokenGroup userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) + allowedChannelTypes := param.AllowedChannelTypes + if len(allowedChannelTypes) == 0 { + allowedChannelTypes = GetAllowedChannelTypes(param.Ctx) + } if param.TokenGroup == "auto" { if len(setting.GetAutoGroups()) == 0 { @@ -115,7 +143,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, } logger.LogDebug(param.Ctx, "Auto selecting group: %s, priorityRetry: %d", autoGroup, priorityRetry) - channel, _ = model.GetRandomSatisfiedChannel(autoGroup, param.ModelName, priorityRetry) + channel, _ = model.GetRandomSatisfiedChannelWithTypes(autoGroup, param.ModelName, priorityRetry, allowedChannelTypes) if channel == nil { // Current group has no available channel for this model, try next group // 当前分组没有该模型的可用渠道，尝试下一个分组 @@ -153,7 +181,7 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, break } } else { - channel, err = model.GetRandomSatisfiedChannel(param.TokenGroup, param.ModelName, param.GetRetry()) + channel, err = model.GetRandomSatisfiedChannelWithTypes(param.TokenGroup, param.ModelName, param.GetRetry(), allowedChannelTypes) if err != nil { return nil, param.TokenGroup, err }