feat(model-mapper): 新增 modelToHeader 配置项并优化 header 更新逻辑 || feat(model-mapper): Added modelToHeader configuration item and optimized header update logic (#3689)

2026-05-26 13:47:27 +08:00 · 2026-04-08 17:12:46 +08:00
parent 228eb27e6a
commit 60ce07d297
2 changed files with 216 additions and 0 deletions
--- a/plugins/wasm-go/extensions/model-mapper/main_test.go
+++ b/plugins/wasm-go/extensions/model-mapper/main_test.go
@@ -42,6 +42,20 @@ var (
 		})
 		return data
 	}()
+
+	customHeaderConfig = func() json.RawMessage {
+		data, _ := json.Marshal(map[string]interface{}{
+			"modelKey":     "model",
+			"modelToHeader": "x-custom-model-header",
+			"modelMapping": map[string]string{
+				"gpt-3.5-turbo": "gpt-4",
+			},
+			"enableOnPathSuffix": []string{
+				"/v1/chat/completions",
+			},
+		})
+		return data
+	}()
 )

 func TestParseConfig(t *testing.T) {
@@ -95,6 +109,21 @@ func TestParseConfig(t *testing.T) {
 			require.Contains(t, cfg.enableOnPathSuffix, "/v1/embeddings")
 		})

+		t.Run("custom modelToHeader", func(t *testing.T) {
+			var cfg Config
+			jsonData := []byte(`{
+				"modelKey": "model",
+				"modelToHeader": "x-custom-model-header",
+				"modelMapping": {
+					"gpt-3.5-turbo": "gpt-4"
+				}
+			}`)
+			err := parseConfig(gjson.ParseBytes(jsonData), &cfg)
+			require.NoError(t, err)
+
+			require.Equal(t, "model", cfg.modelKey)
+		})
+
 		t.Run("modelMapping must be object", func(t *testing.T) {
 			var cfg Config
 			jsonData := []byte(`{
@@ -246,5 +275,179 @@ func TestOnHttpRequestBody_ModelMapping(t *testing.T) {
 			require.NotNil(t, processed)
 			require.Equal(t, "gpt-4-mini", gjson.GetBytes(processed, "request.model").String())
 		})
+
+		t.Run("update model header when model changes", func(t *testing.T) {
+			host, status := test.NewTestHost(basicConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+				{"content-type", "application/json"},
+				{"x-higress-llm-model", "gpt-3.5-turbo-fallback"},
+			})
+
+			origBody := []byte(`{
+				"model": "gpt-3.5-turbo",
+				"messages": [{"role": "user", "content": "hello"}]
+			}`)
+			action := host.CallOnHttpRequestBody(origBody)
+			require.Equal(t, types.ActionContinue, action)
+
+			// verify x-higress-llm-model header was updated to the mapped target
+			newHeaders := host.GetRequestHeaders()
+			foundUpdatedHeader := false
+			for _, h := range newHeaders {
+				if strings.ToLower(h[0]) == "x-higress-llm-model" {
+					require.Equal(t, "gpt-4", h[1])
+					foundUpdatedHeader = true
+					break
+				}
+			}
+			require.True(t, foundUpdatedHeader, "x-higress-llm-model header should be updated")
+		})
+
+		t.Run("skip model header update when header not set", func(t *testing.T) {
+			host, status := test.NewTestHost(basicConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+				{"content-type", "application/json"},
+			})
+
+			origBody := []byte(`{
+				"model": "gpt-3.5-turbo",
+				"messages": [{"role": "user", "content": "hello"}]
+			}`)
+			action := host.CallOnHttpRequestBody(origBody)
+			require.Equal(t, types.ActionContinue, action)
+
+			// verify x-higress-llm-model header was NOT added (should not exist)
+			newHeaders := host.GetRequestHeaders()
+			for _, h := range newHeaders {
+				require.NotEqual(t, strings.ToLower(h[0]), "x-higress-llm-model")
+			}
+		})
+
+		t.Run("skip model header update when header already matches new model", func(t *testing.T) {
+			host, status := test.NewTestHost(basicConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+				{"content-type", "application/json"},
+				{"x-higress-llm-model", "gpt-4"},
+			})
+
+			origBody := []byte(`{
+				"model": "gpt-3.5-turbo",
+				"messages": [{"role": "user", "content": "hello"}]
+			}`)
+			action := host.CallOnHttpRequestBody(origBody)
+			require.Equal(t, types.ActionContinue, action)
+
+			// verify x-higress-llm-model header has the correct value
+			newHeaders := host.GetRequestHeaders()
+			for _, h := range newHeaders {
+				if strings.ToLower(h[0]) == "x-higress-llm-model" {
+					require.Equal(t, "gpt-4", h[1])
+					break
+				}
+			}
+		})
+
+		t.Run("no model mapping keeps header unchanged", func(t *testing.T) {
+			host, status := test.NewTestHost(basicConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+				{"content-type", "application/json"},
+				{"x-higress-llm-model", "some-other-model"},
+			})
+
+			origBody := []byte(`{
+				"model": "unknown-model",
+				"messages": [{"role": "user", "content": "hello"}]
+			}`)
+			action := host.CallOnHttpRequestBody(origBody)
+			require.Equal(t, types.ActionContinue, action)
+
+			// model should remain unchanged (no mapping)
+			processed := host.GetRequestBody()
+			require.NotNil(t, processed)
+			require.Equal(t, "unknown-model", gjson.GetBytes(processed, "model").String())
+		})
+
+		t.Run("use custom modelToHeader config", func(t *testing.T) {
+			host, status := test.NewTestHost(customHeaderConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+				{"content-type", "application/json"},
+				{"x-custom-model-header", "original-model"},
+			})
+
+			origBody := []byte(`{
+				"model": "gpt-3.5-turbo",
+				"messages": [{"role": "user", "content": "hello"}]
+			}`)
+			action := host.CallOnHttpRequestBody(origBody)
+			require.Equal(t, types.ActionContinue, action)
+
+			// verify custom header was updated to the mapped target
+			newHeaders := host.GetRequestHeaders()
+			foundUpdatedHeader := false
+			for _, h := range newHeaders {
+				if strings.ToLower(h[0]) == "x-custom-model-header" {
+					require.Equal(t, "gpt-4", h[1])
+					foundUpdatedHeader = true
+					break
+				}
+			}
+			require.True(t, foundUpdatedHeader, "x-custom-model-header should be updated")
+		})
+
+		t.Run("use custom modelToHeader with empty header value", func(t *testing.T) {
+			host, status := test.NewTestHost(customHeaderConfig)
+			defer host.Reset()
+			require.Equal(t, types.OnPluginStartStatusOK, status)
+
+			host.CallOnHttpRequestHeaders([][2]string{
+				{":authority", "example.com"},
+				{":path", "/v1/chat/completions"},
+				{":method", "POST"},
+				{"content-type", "application/json"},
+			})
+
+			origBody := []byte(`{
+				"model": "gpt-3.5-turbo",
+				"messages": [{"role": "user", "content": "hello"}]
+			}`)
+			action := host.CallOnHttpRequestBody(origBody)
+			require.Equal(t, types.ActionContinue, action)
+
+			// verify custom header was NOT added when not present
+			newHeaders := host.GetRequestHeaders()
+			for _, h := range newHeaders {
+				require.NotEqual(t, strings.ToLower(h[0]), "x-custom-model-header")
+			}
+		})
 	})
 }