jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-05-27 22:27:29 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(ai-quota): add enable_path_suffixes configuration and update rel… (#3748)

Browse Source 
Signed-off-by: zat366 <authentic.zhao@gmail.com>
Co-authored-by: Kent Dong <ch3cho@qq.com>
Co-authored-by: EndlessSeeker <153817598+EndlessSeeker@users.noreply.github.com>
This commit is contained in:
zat366
2026-05-12 17:44:10 +08:00
committed by GitHub
parent 29da03c371
commit f8d81a7eb4
4 changed files with 101 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
plugins/wasm-go/extensions/ai-quota/main.go
View File

@@ -54,12 +54,13 @@ func init() {
}
type QuotaConfig struct {
redisInfo RedisInfo `yaml:"redis"`
RedisKeyPrefix string `yaml:"redis_key_prefix"`
AdminConsumer string `yaml:"admin_consumer"`
AdminPath string `yaml:"admin_path"`
credential2Name map[string]string `yaml:"-"`
redisClient wrapper.RedisClient
redisInfo RedisInfo `yaml:"redis"`
RedisKeyPrefix string `yaml:"redis_key_prefix"`
AdminConsumer string `yaml:"admin_consumer"`
AdminPath string `yaml:"admin_path"`
EnablePathSuffixes []string `yaml:"enable_path_suffixes"`
credential2Name map[string]string `yaml:"-"`
redisClient wrapper.RedisClient
}
type Consumer struct {
@@ -84,6 +85,25 @@ func parseConfig(json gjson.Result, config *QuotaConfig) error {
if config.AdminPath == "" {
config.AdminPath = "/quota"
}
suffixResult := json.Get("enable_path_suffixes")
if !suffixResult.Exists() {
config.EnablePathSuffixes = []string{"/v1/chat/completions", "/v1/messages"}
} else if !suffixResult.IsArray() {
return errors.New("enable_path_suffixes must be an array")
} else {
pathSuffixes := suffixResult.Array()
config.EnablePathSuffixes = make([]string, 0, len(pathSuffixes))
for _, suffix := range pathSuffixes {
suffixStr := strings.TrimSpace(suffix.String())
if suffixStr == "" {
continue
}
config.EnablePathSuffixes = append(config.EnablePathSuffixes, suffixStr)
}
}
if len(config.EnablePathSuffixes) == 0 {
return errors.New("enable_path_suffixes must not be empty")
}
if config.AdminConsumer == "" {
return errors.New("missing admin_consumer in config")
}
@@ -144,7 +164,7 @@ func onHttpRequestHeaders(context wrapper.HttpContext, config QuotaConfig) types
rawPath := context.Path()
path, _ := url.Parse(rawPath)
chatMode, adminMode := getOperationMode(path.Path, config.AdminPath)
chatMode, adminMode := getOperationMode(path.Path, config.AdminPath, config.EnablePathSuffixes)
context.SetContext("chatMode", chatMode)
context.SetContext("adminMode", adminMode)
context.SetContext("consumer", consumer)
@@ -257,7 +277,7 @@ func deniedUnauthorizedConsumer() types.Action {
return types.ActionContinue
}
func getOperationMode(path string, adminPath string) (ChatMode, AdminMode) {
func getOperationMode(path string, adminPath string, pathSuffixes []string) (ChatMode, AdminMode) {
fullAdminPath := "/v1/chat/completions" + adminPath
if strings.HasSuffix(path, fullAdminPath+"/refresh") {
return ChatModeAdmin, AdminModeRefresh
@@ -268,8 +288,10 @@ func getOperationMode(path string, adminPath string) (ChatMode, AdminMode) {
if strings.HasSuffix(path, fullAdminPath) {
return ChatModeAdmin, AdminModeQuery
}
if strings.HasSuffix(path, "/v1/chat/completions") {
return ChatModeCompletion, AdminModeNone
for _, suffix := range pathSuffixes {
if strings.HasSuffix(path, suffix) {
return ChatModeCompletion, AdminModeNone
}
}
return ChatModeNone, AdminModeNone
}