jiazhizhong/higress
1
0
Fork 0
mirror of https://github.com/alibaba/higress.git synced 2026-05-21 11:17:28 +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

1
plugins/wasm-go/extensions/ai-quota/README.md
View File

@@ -22,6 +22,7 @@ description: AI 配额管理插件配置参考
| `redis_key_prefix` | string | 选填 | chat_quota: | qutoa redis key 前缀 |
| `admin_consumer` | string | 必填 | | 管理 quota 管理身份的 consumer 名称 |
| `admin_path` | string | 选填 | /quota | 管理 quota 请求 path 前缀 |
| `enable_path_suffixes` | []string | 选填 | ["/v1/chat/completions", "/v1/messages"] | 启用配额校验的请求路径后缀（仅用于 completion 请求，不影响管理接口路径） |
| `redis` | object | 是 | | redis相关配置 |
`redis`中每一项的配置字段说明

1
plugins/wasm-go/extensions/ai-quota/README_EN.md
View File

@@ -16,6 +16,7 @@ Plugin execution priority: `750`
| `redis_key_prefix` | string | Optional | chat_quota: | Quota redis key prefix |
| `admin_consumer` | string | Required | | Consumer name for managing quota management identity |
| `admin_path` | string | Optional | /quota | Prefix for the path to manage quota requests |
| `enable_path_suffixes` | []string | Optional | ["/v1/chat/completions", "/v1/messages"] | Enabled path suffixes for completion quota checks only; does not affect admin API path |
| `redis` | object | Yes | | Redis related configuration |
Explanation of each configuration field in `redis`
| Configuration Item | Type | Required | Default Value | Explanation |

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
}

68
plugins/wasm-go/extensions/ai-quota/main_test.go
View File

@@ -30,6 +30,10 @@ var basicConfig = func() json.RawMessage {
"admin_consumer": "admin",
"redis_key_prefix": "chat_quota:",
"admin_path": "/quota",
"enable_path_suffixes": []string{
"/v1/chat/completions",
"/v1/messages",
},
"redis": map[string]interface{}{
"service_name": "redis.static",
"service_port": 6379,
@@ -51,6 +55,17 @@ var missingAdminConsumerConfig = func() json.RawMessage {
return data
}()
var defaultPathSuffixesConfig = func() json.RawMessage {
data, _ := json.Marshal(map[string]interface{}{
"admin_consumer": "admin",
"redis": map[string]interface{}{
"service_name": "redis.static",
"service_port": 6379,
},
})
return data
}()
func TestParseConfig(t *testing.T) {
test.RunGoTest(t, func(t *testing.T) {
// 测试基础配置解析
@@ -66,6 +81,7 @@ func TestParseConfig(t *testing.T) {
require.Equal(t, "admin", quotaConfig.AdminConsumer)
require.Equal(t, "chat_quota:", quotaConfig.RedisKeyPrefix)
require.Equal(t, "/quota", quotaConfig.AdminPath)
require.Equal(t, []string{"/v1/chat/completions", "/v1/messages"}, quotaConfig.EnablePathSuffixes)
})
// 测试缺少admin_consumer的配置
@@ -74,6 +90,18 @@ func TestParseConfig(t *testing.T) {
defer host.Reset()
require.Equal(t, types.OnPluginStartStatusFailed, status)
})
t.Run("default path suffixes", func(t *testing.T) {
host, status := test.NewTestHost(defaultPathSuffixesConfig)
defer host.Reset()
require.Equal(t, types.OnPluginStartStatusOK, status)
config, err := host.GetMatchConfig()
require.NoError(t, err)
require.NotNil(t, config)
quotaConfig := config.(*QuotaConfig)
require.Equal(t, []string{"/v1/chat/completions", "/v1/messages"}, quotaConfig.EnablePathSuffixes)
})
})
}
@@ -278,6 +306,7 @@ func TestGetOperationMode(t *testing.T) {
name string
path string
adminPath string
suffixes []string
chatMode ChatMode
adminMode AdminMode
}{
@@ -285,6 +314,7 @@ func TestGetOperationMode(t *testing.T) {
name: "chat completion mode",
path: "/v1/chat/completions",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeCompletion,
adminMode: AdminModeNone,
},
@@ -292,6 +322,7 @@ func TestGetOperationMode(t *testing.T) {
name: "admin query mode",
path: "/v1/chat/completions/quota",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeAdmin,
adminMode: AdminModeQuery,
},
@@ -299,6 +330,7 @@ func TestGetOperationMode(t *testing.T) {
name: "admin refresh mode",
path: "/v1/chat/completions/quota/refresh",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeAdmin,
adminMode: AdminModeRefresh,
},
@@ -306,13 +338,47 @@ func TestGetOperationMode(t *testing.T) {
name: "admin delta mode",
path: "/v1/chat/completions/quota/delta",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeAdmin,
adminMode: AdminModeDelta,
},
{
name: "anthropic messages completion mode",
path: "/v1/messages",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeCompletion,
adminMode: AdminModeNone,
},
{
name: "custom suffix completion mode",
path: "/llm/invoke",
adminPath: "/quota",
suffixes: []string{"/invoke"},
chatMode: ChatModeCompletion,
adminMode: AdminModeNone,
},
{
name: "admin path fixed to chat completions",
path: "/v1/chat/completions/quota",
adminPath: "/quota",
suffixes: []string{"/invoke"},
chatMode: ChatModeAdmin,
adminMode: AdminModeQuery,
},
{
name: "messages admin path not supported",
path: "/v1/messages/quota",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeNone,
adminMode: AdminModeNone,
},
{
name: "none mode",
path: "/other/path",
adminPath: "/quota",
suffixes: []string{"/v1/chat/completions", "/v1/messages"},
chatMode: ChatModeNone,
adminMode: AdminModeNone,
},
@@ -320,7 +386,7 @@ func TestGetOperationMode(t *testing.T) {
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
chatMode, adminMode := getOperationMode(tt.path, tt.adminPath)
chatMode, adminMode := getOperationMode(tt.path, tt.adminPath, tt.suffixes)
require.Equal(t, tt.chatMode, chatMode)
require.Equal(t, tt.adminMode, adminMode)
})