🐛 fix: claude api restricts channel type (lobehub#318)

coulsontl · Dec 3, 2024 · d9b9ae9 · d9b9ae9
1 parent 8ed1a9f
commit d9b9ae9
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 0 deletions.
diff --git a/model/balancer.go b/model/balancer.go
@@ -34,6 +34,12 @@ func FilterChannelId(skipChannelIds []int) ChannelsFilterFunc {
 	}
 }
 
+func FilterChannelTypes(channelTypes []int) ChannelsFilterFunc {
+	return func(_ int, choice *ChannelChoice) bool {
+		return !utils.Contains(choice.Channel.Type, channelTypes)
+	}
+}
+
 func FilterOnlyChat() ChannelsFilterFunc {
 	return func(channelId int, choice *ChannelChoice) bool {
 		return choice.Channel.OnlyChat

diff --git a/relay/claude.go b/relay/claude.go
@@ -18,13 +18,16 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
+var AllowChannelType = []int{config.ChannelTypeAnthropic, config.ChannelTypeVertexAI}
+
 func RelaycClaudeOnly(c *gin.Context) {
 	request := &claude.ClaudeRequest{}
 
 	if err := common.UnmarshalBodyReusable(c, request); err != nil {
 		common.AbortWithErr(c, http.StatusBadRequest, claude.ErrorToClaudeErr(err))
 		return
 	}
+	c.Set("allow_channel_type", AllowChannelType)
 
 	cacheProps := relay_util.NewChatCacheProps(c, true)
 	cacheProps.SetHash(request)

diff --git a/relay/common.go b/relay/common.go
@@ -120,6 +120,12 @@ func fetchChannelByModel(c *gin.Context, modelName string) (*model.Channel, erro
 		filters = append(filters, model.FilterChannelId(skipChannelIds))
 	}
 
+	if types, exists := c.Get("allow_channel_type"); exists {
+		if allowTypes, ok := types.([]int); ok {
+			filters = append(filters, model.FilterChannelTypes(allowTypes))
+		}
+	}
+
 	channel, err := model.ChannelGroup.Next(group, modelName, filters...)
 	if err != nil {
 		message := fmt.Sprintf("当前分组 %s 下对于模型 %s 无可用渠道", group, modelName)