Просмотр исходного кода

feat: add support for multi-key channels in RelayInfo and access token caching

CaIon 7 месяцев назад
Родитель
Сommit
c056a7ad7c
2 измененных файлов с 22 добавлено и 12 удалено
  1. 6 1
      relay/channel/vertex/service_account.go
  2. 16 11
      relay/common/relay_info.go

+ 6 - 1
relay/channel/vertex/service_account.go

@@ -36,7 +36,12 @@ var Cache = asynccache.NewAsyncCache(asynccache.Options{
 })
 
 func getAccessToken(a *Adaptor, info *relaycommon.RelayInfo) (string, error) {
-	cacheKey := fmt.Sprintf("access-token-%d", info.ChannelId)
+	var cacheKey string
+	if info.ChannelIsMultiKey {
+		cacheKey = fmt.Sprintf("access-token-%d-%d", info.ChannelId, info.ChannelMultiKeyIndex)
+	} else {
+		cacheKey = fmt.Sprintf("access-token-%d", info.ChannelId)
+	}
 	val, err := Cache.Get(cacheKey)
 	if err == nil {
 		return val.(string), nil

+ 16 - 11
relay/common/relay_info.go

@@ -60,17 +60,19 @@ type ResponsesUsageInfo struct {
 }
 
 type RelayInfo struct {
-	ChannelType       int
-	ChannelId         int
-	TokenId           int
-	TokenKey          string
-	UserId            int
-	UsingGroup        string // 使用的分组
-	UserGroup         string // 用户所在分组
-	TokenUnlimited    bool
-	StartTime         time.Time
-	FirstResponseTime time.Time
-	isFirstResponse   bool
+	ChannelType          int
+	ChannelId            int
+	ChannelIsMultiKey    bool // 是否多密钥
+	ChannelMultiKeyIndex int  // 多密钥索引
+	TokenId              int
+	TokenKey             string
+	UserId               int
+	UsingGroup           string // 使用的分组
+	UserGroup            string // 用户所在分组
+	TokenUnlimited       bool
+	StartTime            time.Time
+	FirstResponseTime    time.Time
+	isFirstResponse      bool
 	//SendLastReasoningResponse bool
 	ApiType           int
 	IsStream          bool
@@ -260,6 +262,9 @@ func GenRelayInfo(c *gin.Context) *RelayInfo {
 			IsFirstThinkingContent:  true,
 			SendLastThinkingContent: false,
 		},
+
+		ChannelIsMultiKey:    common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey),
+		ChannelMultiKeyIndex: common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex),
 	}
 	if strings.HasPrefix(c.Request.URL.Path, "/pg") {
 		info.IsPlayground = true