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

feat: add IncludeModelName option to channel affinity rules for per-model affinity tracking

Clansty 2 месяцев назад
Родитель
Сommit
70560d5371

+ 17 - 4
service/channel_affinity.go

@@ -166,12 +166,22 @@ func GetChannelAffinityCacheStats() ChannelAffinityCacheStats {
 			unknown++
 			continue
 		}
-		if rule.IncludeUsingGroup {
+		if rule.IncludeModelName {
 			if len(parts) < 3 {
 				unknown++
 				continue
 			}
 		}
+		if rule.IncludeUsingGroup {
+			minParts := 3
+			if rule.IncludeModelName {
+				minParts = 4
+			}
+			if len(parts) < minParts {
+				unknown++
+				continue
+			}
+		}
 		byRuleName[ruleName]++
 	}
 
@@ -319,11 +329,14 @@ func extractChannelAffinityValue(c *gin.Context, src operation_setting.ChannelAf
 	}
 }
 
-func buildChannelAffinityCacheKeySuffix(rule operation_setting.ChannelAffinityRule, usingGroup string, affinityValue string) string {
-	parts := make([]string, 0, 3)
+func buildChannelAffinityCacheKeySuffix(rule operation_setting.ChannelAffinityRule, modelName string, usingGroup string, affinityValue string) string {
+	parts := make([]string, 0, 4)
 	if rule.IncludeRuleName && rule.Name != "" {
 		parts = append(parts, rule.Name)
 	}
+	if rule.IncludeModelName && modelName != "" {
+		parts = append(parts, modelName)
+	}
 	if rule.IncludeUsingGroup && usingGroup != "" {
 		parts = append(parts, usingGroup)
 	}
@@ -573,7 +586,7 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup
 		if ttlSeconds <= 0 {
 			ttlSeconds = setting.DefaultTTLSeconds
 		}
-		cacheKeySuffix := buildChannelAffinityCacheKeySuffix(rule, usingGroup, affinityValue)
+		cacheKeySuffix := buildChannelAffinityCacheKeySuffix(rule, modelName, usingGroup, affinityValue)
 		cacheKeyFull := channelAffinityCacheNamespace + ":" + cacheKeySuffix
 		setChannelAffinityContext(c, channelAffinityMeta{
 			CacheKey:       cacheKeyFull,

+ 1 - 1
service/channel_affinity_template_test.go

@@ -193,7 +193,7 @@ func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
 	require.NotNil(t, codexRule)
 
 	affinityValue := fmt.Sprintf("pc-hit-%d", time.Now().UnixNano())
-	cacheKeySuffix := buildChannelAffinityCacheKeySuffix(*codexRule, "default", affinityValue)
+	cacheKeySuffix := buildChannelAffinityCacheKeySuffix(*codexRule, "gpt-5", "default", affinityValue)
 
 	cache := getChannelAffinityCache()
 	require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute))

+ 1 - 0
setting/operation_setting/channel_affinity_setting.go

@@ -23,6 +23,7 @@ type ChannelAffinityRule struct {
 	SkipRetryOnFailure bool `json:"skip_retry_on_failure,omitempty"`
 
 	IncludeUsingGroup bool `json:"include_using_group"`
+	IncludeModelName  bool `json:"include_model_name"`
 	IncludeRuleName   bool `json:"include_rule_name"`
 }