Merge pull request #3009 from seefs001/feature/improve-param-override
feat: improve channel override ui/ux
This commit is contained in:
@@ -45,6 +45,7 @@ type channelAffinityMeta struct {
|
||||
TTLSeconds int
|
||||
RuleName string
|
||||
SkipRetry bool
|
||||
ParamTemplate map[string]interface{}
|
||||
KeySourceType string
|
||||
KeySourceKey string
|
||||
KeySourcePath string
|
||||
@@ -415,6 +416,84 @@ func buildChannelAffinityKeyHint(s string) string {
|
||||
return s[:4] + "..." + s[len(s)-4:]
|
||||
}
|
||||
|
||||
func cloneStringAnyMap(src map[string]interface{}) map[string]interface{} {
|
||||
if len(src) == 0 {
|
||||
return map[string]interface{}{}
|
||||
}
|
||||
dst := make(map[string]interface{}, len(src))
|
||||
for k, v := range src {
|
||||
dst[k] = v
|
||||
}
|
||||
return dst
|
||||
}
|
||||
|
||||
func mergeChannelOverride(base map[string]interface{}, tpl map[string]interface{}) map[string]interface{} {
|
||||
if len(base) == 0 && len(tpl) == 0 {
|
||||
return map[string]interface{}{}
|
||||
}
|
||||
if len(tpl) == 0 {
|
||||
return base
|
||||
}
|
||||
out := cloneStringAnyMap(base)
|
||||
for k, v := range tpl {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func appendChannelAffinityTemplateAdminInfo(c *gin.Context, meta channelAffinityMeta) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if len(meta.ParamTemplate) == 0 {
|
||||
return
|
||||
}
|
||||
|
||||
templateInfo := map[string]interface{}{
|
||||
"applied": true,
|
||||
"rule_name": meta.RuleName,
|
||||
"param_override_keys": len(meta.ParamTemplate),
|
||||
}
|
||||
if anyInfo, ok := c.Get(ginKeyChannelAffinityLogInfo); ok {
|
||||
if info, ok := anyInfo.(map[string]interface{}); ok {
|
||||
info["override_template"] = templateInfo
|
||||
c.Set(ginKeyChannelAffinityLogInfo, info)
|
||||
return
|
||||
}
|
||||
}
|
||||
c.Set(ginKeyChannelAffinityLogInfo, map[string]interface{}{
|
||||
"reason": meta.RuleName,
|
||||
"rule_name": meta.RuleName,
|
||||
"using_group": meta.UsingGroup,
|
||||
"model": meta.ModelName,
|
||||
"request_path": meta.RequestPath,
|
||||
"key_source": meta.KeySourceType,
|
||||
"key_key": meta.KeySourceKey,
|
||||
"key_path": meta.KeySourcePath,
|
||||
"key_hint": meta.KeyHint,
|
||||
"key_fp": meta.KeyFingerprint,
|
||||
"override_template": templateInfo,
|
||||
})
|
||||
}
|
||||
|
||||
// ApplyChannelAffinityOverrideTemplate merges per-rule channel override templates onto the selected channel override config.
|
||||
func ApplyChannelAffinityOverrideTemplate(c *gin.Context, paramOverride map[string]interface{}) (map[string]interface{}, bool) {
|
||||
if c == nil {
|
||||
return paramOverride, false
|
||||
}
|
||||
meta, ok := getChannelAffinityMeta(c)
|
||||
if !ok {
|
||||
return paramOverride, false
|
||||
}
|
||||
if len(meta.ParamTemplate) == 0 {
|
||||
return paramOverride, false
|
||||
}
|
||||
|
||||
mergedParam := mergeChannelOverride(paramOverride, meta.ParamTemplate)
|
||||
appendChannelAffinityTemplateAdminInfo(c, meta)
|
||||
return mergedParam, true
|
||||
}
|
||||
|
||||
func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup string) (int, bool) {
|
||||
setting := operation_setting.GetChannelAffinitySetting()
|
||||
if setting == nil || !setting.Enabled {
|
||||
@@ -466,6 +545,7 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup
|
||||
TTLSeconds: ttlSeconds,
|
||||
RuleName: rule.Name,
|
||||
SkipRetry: rule.SkipRetryOnFailure,
|
||||
ParamTemplate: cloneStringAnyMap(rule.ParamOverrideTemplate),
|
||||
KeySourceType: strings.TrimSpace(usedSource.Type),
|
||||
KeySourceKey: strings.TrimSpace(usedSource.Key),
|
||||
KeySourcePath: strings.TrimSpace(usedSource.Path),
|
||||
|
||||
145
service/channel_affinity_template_test.go
Normal file
145
service/channel_affinity_template_test.go
Normal file
@@ -0,0 +1,145 @@
|
||||
package service
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
"github.com/QuantumNous/new-api/setting/operation_setting"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func buildChannelAffinityTemplateContextForTest(meta channelAffinityMeta) *gin.Context {
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
setChannelAffinityContext(ctx, meta)
|
||||
return ctx
|
||||
}
|
||||
|
||||
func TestApplyChannelAffinityOverrideTemplate_NoTemplate(t *testing.T) {
|
||||
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
||||
RuleName: "rule-no-template",
|
||||
})
|
||||
base := map[string]interface{}{
|
||||
"temperature": 0.7,
|
||||
}
|
||||
|
||||
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
||||
require.False(t, applied)
|
||||
require.Equal(t, base, merged)
|
||||
}
|
||||
|
||||
func TestApplyChannelAffinityOverrideTemplate_MergeTemplate(t *testing.T) {
|
||||
ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
|
||||
RuleName: "rule-with-template",
|
||||
ParamTemplate: map[string]interface{}{
|
||||
"temperature": 0.2,
|
||||
"top_p": 0.95,
|
||||
},
|
||||
UsingGroup: "default",
|
||||
ModelName: "gpt-4.1",
|
||||
RequestPath: "/v1/responses",
|
||||
KeySourceType: "gjson",
|
||||
KeySourcePath: "prompt_cache_key",
|
||||
KeyHint: "abcd...wxyz",
|
||||
KeyFingerprint: "abcd1234",
|
||||
})
|
||||
base := map[string]interface{}{
|
||||
"temperature": 0.7,
|
||||
"max_tokens": 2000,
|
||||
}
|
||||
|
||||
merged, applied := ApplyChannelAffinityOverrideTemplate(ctx, base)
|
||||
require.True(t, applied)
|
||||
require.Equal(t, 0.2, merged["temperature"])
|
||||
require.Equal(t, 0.95, merged["top_p"])
|
||||
require.Equal(t, 2000, merged["max_tokens"])
|
||||
require.Equal(t, 0.7, base["temperature"])
|
||||
|
||||
anyInfo, ok := ctx.Get(ginKeyChannelAffinityLogInfo)
|
||||
require.True(t, ok)
|
||||
info, ok := anyInfo.(map[string]interface{})
|
||||
require.True(t, ok)
|
||||
overrideInfoAny, ok := info["override_template"]
|
||||
require.True(t, ok)
|
||||
overrideInfo, ok := overrideInfoAny.(map[string]interface{})
|
||||
require.True(t, ok)
|
||||
require.Equal(t, true, overrideInfo["applied"])
|
||||
require.Equal(t, "rule-with-template", overrideInfo["rule_name"])
|
||||
require.EqualValues(t, 2, overrideInfo["param_override_keys"])
|
||||
}
|
||||
|
||||
func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
setting := operation_setting.GetChannelAffinitySetting()
|
||||
require.NotNil(t, setting)
|
||||
|
||||
var codexRule *operation_setting.ChannelAffinityRule
|
||||
for i := range setting.Rules {
|
||||
rule := &setting.Rules[i]
|
||||
if strings.EqualFold(strings.TrimSpace(rule.Name), "codex cli trace") {
|
||||
codexRule = rule
|
||||
break
|
||||
}
|
||||
}
|
||||
require.NotNil(t, codexRule)
|
||||
|
||||
affinityValue := fmt.Sprintf("pc-hit-%d", time.Now().UnixNano())
|
||||
cacheKeySuffix := buildChannelAffinityCacheKeySuffix(*codexRule, "default", affinityValue)
|
||||
|
||||
cache := getChannelAffinityCache()
|
||||
require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute))
|
||||
t.Cleanup(func() {
|
||||
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(fmt.Sprintf(`{"prompt_cache_key":"%s"}`, affinityValue)))
|
||||
ctx.Request.Header.Set("Content-Type", "application/json")
|
||||
|
||||
channelID, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
|
||||
require.True(t, found)
|
||||
require.Equal(t, 9527, channelID)
|
||||
|
||||
baseOverride := map[string]interface{}{
|
||||
"temperature": 0.2,
|
||||
}
|
||||
mergedOverride, applied := ApplyChannelAffinityOverrideTemplate(ctx, baseOverride)
|
||||
require.True(t, applied)
|
||||
require.Equal(t, 0.2, mergedOverride["temperature"])
|
||||
|
||||
info := &relaycommon.RelayInfo{
|
||||
RequestHeaders: map[string]string{
|
||||
"Originator": "Codex CLI",
|
||||
"Session_id": "sess-123",
|
||||
"User-Agent": "codex-cli-test",
|
||||
},
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ParamOverride: mergedOverride,
|
||||
HeadersOverride: map[string]interface{}{
|
||||
"X-Static": "legacy-static",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
_, err := relaycommon.ApplyParamOverrideWithRelayInfo([]byte(`{"model":"gpt-5"}`), info)
|
||||
require.NoError(t, err)
|
||||
require.True(t, info.UseRuntimeHeadersOverride)
|
||||
|
||||
require.Equal(t, "legacy-static", info.RuntimeHeadersOverride["x-static"])
|
||||
require.Equal(t, "Codex CLI", info.RuntimeHeadersOverride["originator"])
|
||||
require.Equal(t, "sess-123", info.RuntimeHeadersOverride["session_id"])
|
||||
require.Equal(t, "codex-cli-test", info.RuntimeHeadersOverride["user-agent"])
|
||||
|
||||
_, exists := info.RuntimeHeadersOverride["x-codex-beta-features"]
|
||||
require.False(t, exists)
|
||||
_, exists = info.RuntimeHeadersOverride["x-codex-turn-metadata"]
|
||||
require.False(t, exists)
|
||||
}
|
||||
Reference in New Issue
Block a user