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

plugins/wasm-go/extensions/model-mapper/main.go

@@ -42,6 +42,7 @@ type Config struct {
 	prefixModelMapping []ModelMapping
 	defaultModel       string
 	enableOnPathSuffix []string
+	modelToHeader      string
 }
 
 func parseConfig(json gjson.Result, config *Config) error {
@@ -49,6 +50,10 @@ func parseConfig(json gjson.Result, config *Config) error {
 	if config.modelKey == "" {
 		config.modelKey = "model"
 	}
+	config.modelToHeader = json.Get("modelToHeader").String()
+	if config.modelToHeader == "" {
+		config.modelToHeader = "x-higress-llm-model"
+	}
 
 	modelMapping := json.Get("modelMapping")
 	if modelMapping.Exists() && !modelMapping.IsObject() {
@@ -144,6 +149,8 @@ func onHttpRequestHeaders(ctx wrapper.HttpContext, config Config) types.Action {
 		return types.ActionContinue
 	}
 
+	// Disable re-route since the plugin may modify some headers related to the chosen route.
+	ctx.DisableReroute()
 	// Prepare for body processing
 	proxywasm.RemoveHttpRequestHeader("content-length")
 	// 100MB buffer limit
@@ -183,6 +190,12 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config Config, body []byte) type
 	}
 
 	if newModel != "" && newModel != oldModel {
+		// if x-higress-llm-model header is set, and it is not the same as the new model, update it
+		// this is to support fallback and token rate limit
+		model, _ := proxywasm.GetHttpRequestHeader(config.modelToHeader)
+		if model != "" && model != newModel {
+			proxywasm.ReplaceHttpRequestHeader(config.modelToHeader, newModel)
+		}
 		newBody, err := sjson.SetBytes(body, config.modelKey, newModel)
 		if err != nil {
 			log.Errorf("failed to update model: %v", err)