Skip to content

Commit

Permalink
feat: add AI quota plugin (alibaba#1200)
Browse files Browse the repository at this point in the history
  • Loading branch information
2456868764 authored Aug 14, 2024
1 parent daa374d commit d31c978
Show file tree
Hide file tree
Showing 6 changed files with 582 additions and 0 deletions.
58 changes: 58 additions & 0 deletions plugins/wasm-go/extensions/ai-quota/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,58 @@
# 功能说明

`ai-qutoa` 插件实现给特定 consumer 根据分配固定的 quota 进行 quota 策略限流,同时支持 quota 管理能力,包括查询 quota 、刷新 quota、增减 quota。

`ai-quota` 插件需要配合 认证插件比如 `key-auth``jwt-auth` 等插件获取认证身份的 consumer 名称,同时需要配合 `ai-statatistics` 插件获取 AI Token 统计信息。

# 配置说明

| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|--------------------|-----------------|--------------------------------------| ---- |--------------------------------------------|
| `redis_key_prefix` | string | 选填 | chat_quota: | qutoa redis key 前缀 |
| `admin_consumer` | string | 必填 | | 管理 quota 管理身份的 consumer 名称 |
| `admin_path` | string | 选填 | /quota | 管理 quota 请求 path 前缀 |
| `redis` | object || | redis相关配置 |

`redis`中每一项的配置字段说明

| 配置项 | 类型 | 必填 | 默认值 | 说明 |
| ------------ | ------ | ---- | ---------------------------------------------------------- | --------------------------- |
| service_name | string | 必填 | - | redis 服务名称,带服务类型的完整 FQDN 名称,例如 my-redis.dns、redis.my-ns.svc.cluster.local |
| service_port | int || 服务类型为固定地址(static service)默认值为80,其他为6379 | 输入redis服务的服务端口 |
| username | string || - | redis用户名 |
| password | string || - | redis密码 |
| timeout | int || 1000 | redis连接超时时间,单位毫秒 |



# 配置示例

## 识别请求参数 apikey,进行区别限流
```yaml
redis_key_prefix: "chat_quota:"
admin_consumer: consumer3
admin_path: /quota
redis:
service_name: redis-service.default.svc.cluster.local
service_port: 6379
timeout: 2000
```
## 刷新 quota
如果当前请求 url 的后缀符合 admin_path,例如插件在 example.com/v1/chat/completions 这个路由上生效,那么更新 quota 可以通过
curl https://example.com/v1/chat/completions/quota/refresh -H "Authorization: Bearer credential3" -d "consumer=consumer1&quota=10000"
Redis 中 key 为 chat_quota:consumer1 的值就会被刷新为 10000
## 查询 quota
查询特定用户的 quota 可以通过 curl https://example.com/v1/chat/completions/quota?consumer=consumer1 -H "Authorization: Bearer credential3"
将返回: {"quota": 10000, "consumer": "consumer1"}
## 增减 quota
增减特定用户的 quota 可以通过 curl https://example.com/v1/chat/completions/quota/delta -d "consumer=consumer1&value=100" -H "Authorization: Bearer credential3"
这样 Redis 中 Key 为 chat_quota:consumer1 的值就会增加100,可以支持负数,则减去对应值。
20 changes: 20 additions & 0 deletions plugins/wasm-go/extensions/ai-quota/go.mod
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
module github.com/alibaba/higress/plugins/wasm-go/extensions/ai-quota

go 1.19

//replace github.com/alibaba/higress/plugins/wasm-go => ../..

require (
github.com/alibaba/higress/plugins/wasm-go v1.4.3-0.20240808022948-34f5722d93de
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f
github.com/tidwall/gjson v1.17.3
github.com/tidwall/resp v0.1.1
)

require (
github.com/google/uuid v1.3.0 // indirect
github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 // indirect
github.com/magefile/mage v1.14.0 // indirect
github.com/tidwall/match v1.1.1 // indirect
github.com/tidwall/pretty v1.2.0 // indirect
)
22 changes: 22 additions & 0 deletions plugins/wasm-go/extensions/ai-quota/go.sum
Original file line number Diff line number Diff line change
@@ -0,0 +1,22 @@
github.com/alibaba/higress/plugins/wasm-go v1.4.3-0.20240808022948-34f5722d93de h1:lDLqj7Hw41ox8VdsP7oCTPhjPa3+QJUCKApcLh2a45Y=
github.com/alibaba/higress/plugins/wasm-go v1.4.3-0.20240808022948-34f5722d93de/go.mod h1:359don/ahMxpfeLMzr29Cjwcu8IywTTDUzWlBPRNLHw=
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 h1:IHDghbGQ2DTIXHBHxWfqCYQW1fKjyJ/I7W1pMyUDeEA=
github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520/go.mod h1:Nz8ORLaFiLWotg6GeKlJMhv8cci8mM43uEnLA5t8iew=
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f h1:ZIiIBRvIw62gA5MJhuwp1+2wWbqL9IGElQ499rUsYYg=
github.com/higress-group/proxy-wasm-go-sdk v0.0.0-20240711023527-ba358c48772f/go.mod h1:hNFjhrLUIq+kJ9bOcs8QtiplSQ61GZXtd2xHKx4BYRo=
github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
github.com/tidwall/gjson v1.17.3 h1:bwWLZU7icoKRG+C+0PNwIKC6FCJO/Q3p2pZvuP0jN94=
github.com/tidwall/gjson v1.17.3/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
github.com/tidwall/resp v0.1.1 h1:Ly20wkhqKTmDUPlyM1S7pWo5kk0tDu8OoC/vFArXmwE=
github.com/tidwall/resp v0.1.1/go.mod h1:3/FrruOBAxPTPtundW0VXgmsQ4ZBA0Aw714lVYgwFa0=
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
Loading

0 comments on commit d31c978

Please sign in to comment.