feat(model-router): add auto routing based on user message content (#3403)

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

@@ -8,6 +8,7 @@ import (
 	"mime/multipart"
 	"net/http"
 	"net/textproto"
+	"regexp"
 	"strings"
 
 	"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
@@ -20,6 +21,7 @@ import (
 
 const (
 	DefaultMaxBodyBytes = 100 * 1024 * 1024 // 100MB
+	AutoModelPrefix     = "higress/auto"
 )
 
 func main() {}
@@ -35,11 +37,21 @@ func init() {
 	)
 }
 
+// AutoRoutingRule defines a regex-based routing rule for auto model selection
+type AutoRoutingRule struct {
+	Pattern *regexp.Regexp
+	Model   string
+}
+
 type ModelRouterConfig struct {
 	modelKey           string
 	addProviderHeader  string
 	modelToHeader      string
 	enableOnPathSuffix []string
+	// Auto routing configuration
+	enableAutoRouting bool
+	autoRoutingRules  []AutoRoutingRule
+	defaultModel      string
 }
 
 func parseConfig(json gjson.Result, config *ModelRouterConfig) error {
@@ -70,6 +82,36 @@ func parseConfig(json gjson.Result, config *ModelRouterConfig) error {
 			"/messages",
 		}
 	}
+
+	// Parse auto routing configuration
+	autoRouting := json.Get("autoRouting")
+	if autoRouting.Exists() {
+		config.enableAutoRouting = autoRouting.Get("enable").Bool()
+		config.defaultModel = autoRouting.Get("defaultModel").String()
+
+		rules := autoRouting.Get("rules")
+		if rules.Exists() && rules.IsArray() {
+			for _, rule := range rules.Array() {
+				patternStr := rule.Get("pattern").String()
+				model := rule.Get("model").String()
+				if patternStr == "" || model == "" {
+					log.Warnf("skipping invalid auto routing rule: pattern=%s, model=%s", patternStr, model)
+					continue
+				}
+				compiled, err := regexp.Compile(patternStr)
+				if err != nil {
+					log.Warnf("failed to compile regex pattern '%s': %v", patternStr, err)
+					continue
+				}
+				config.autoRoutingRules = append(config.autoRoutingRules, AutoRoutingRule{
+					Pattern: compiled,
+					Model:   model,
+				})
+				log.Debugf("loaded auto routing rule: pattern=%s, model=%s", patternStr, model)
+			}
+		}
+	}
+
 	return nil
 }
 
@@ -120,6 +162,43 @@ func onHttpRequestBody(ctx wrapper.HttpContext, config ModelRouterConfig, body [
 	return types.ActionContinue
 }
 
+// extractLastUserMessage extracts the content of the last message with role "user" from the messages array
+func extractLastUserMessage(body []byte) string {
+	messages := gjson.GetBytes(body, "messages")
+	if !messages.Exists() || !messages.IsArray() {
+		return ""
+	}
+
+	var lastUserContent string
+	for _, msg := range messages.Array() {
+		if msg.Get("role").String() == "user" {
+			content := msg.Get("content")
+			if content.IsArray() {
+				// Handle array content (e.g., multimodal messages with text and images)
+				for _, item := range content.Array() {
+					if item.Get("type").String() == "text" {
+						lastUserContent = item.Get("text").String()
+					}
+				}
+			} else {
+				lastUserContent = content.String()
+			}
+		}
+	}
+	return lastUserContent
+}
+
+// matchAutoRoutingRule matches the user message against auto routing rules and returns the matched model
+func matchAutoRoutingRule(config ModelRouterConfig, userMessage string) (string, bool) {
+	for _, rule := range config.autoRoutingRules {
+		if rule.Pattern.MatchString(userMessage) {
+			log.Debugf("auto routing rule matched: pattern=%s, model=%s", rule.Pattern.String(), rule.Model)
+			return rule.Model, true
+		}
+	}
+	return "", false
+}
+
 func handleJsonBody(ctx wrapper.HttpContext, config ModelRouterConfig, body []byte) types.Action {
 	if !json.Valid(body) {
 		log.Error("invalid json body")
@@ -130,6 +209,27 @@ func handleJsonBody(ctx wrapper.HttpContext, config ModelRouterConfig, body []by
 		return types.ActionContinue
 	}
 
+	// Check if auto routing should be triggered
+	if config.enableAutoRouting && modelValue == AutoModelPrefix {
+		userMessage := extractLastUserMessage(body)
+		if userMessage != "" {
+			if matchedModel, found := matchAutoRoutingRule(config, userMessage); found {
+				// Set the matched model to the header for routing
+				_ = proxywasm.ReplaceHttpRequestHeader("x-higress-llm-model", matchedModel)
+				log.Infof("auto routing: user message matched, routing to model: %s", matchedModel)
+				return types.ActionContinue
+			}
+		}
+		// No rule matched, use default model if configured
+		if config.defaultModel != "" {
+			_ = proxywasm.ReplaceHttpRequestHeader("x-higress-llm-model", config.defaultModel)
+			log.Infof("auto routing: no rule matched, using default model: %s", config.defaultModel)
+		} else {
+			log.Warnf("auto routing: no rule matched and no default model configured")
+		}
+		return types.ActionContinue
+	}
+
 	if config.modelToHeader != "" {
 		_ = proxywasm.ReplaceHttpRequestHeader(config.modelToHeader, modelValue)
 	}