diff --git a/plugins/wasm-go/extensions/ai-quota/README.md b/plugins/wasm-go/extensions/ai-quota/README.md index 4b0d362f..6bfa6de1 100644 --- a/plugins/wasm-go/extensions/ai-quota/README.md +++ b/plugins/wasm-go/extensions/ai-quota/README.md @@ -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`中每一项的配置字段说明 diff --git a/plugins/wasm-go/extensions/ai-quota/README_EN.md b/plugins/wasm-go/extensions/ai-quota/README_EN.md index 0eff19ae..070cb3e5 100644 --- a/plugins/wasm-go/extensions/ai-quota/README_EN.md +++ b/plugins/wasm-go/extensions/ai-quota/README_EN.md @@ -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 | diff --git a/plugins/wasm-go/extensions/ai-quota/main.go b/plugins/wasm-go/extensions/ai-quota/main.go index a7647652..1b7dcbf9 100644 --- a/plugins/wasm-go/extensions/ai-quota/main.go +++ b/plugins/wasm-go/extensions/ai-quota/main.go @@ -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 } diff --git a/plugins/wasm-go/extensions/ai-quota/main_test.go b/plugins/wasm-go/extensions/ai-quota/main_test.go index e6b70b7e..bb2473b7 100644 --- a/plugins/wasm-go/extensions/ai-quota/main_test.go +++ b/plugins/wasm-go/extensions/ai-quota/main_test.go @@ -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) })