diff --git a/apisix/plugins/ai-rate-limiting.lua b/apisix/plugins/ai-rate-limiting.lua index cdf2d9fb140d..489176a63358 100644 --- a/apisix/plugins/ai-rate-limiting.lua +++ b/apisix/plugins/ai-rate-limiting.lua @@ -25,6 +25,7 @@ local math_floor = math.floor local math_huge = math.huge local core = require("apisix.core") local limit_count = require("apisix.plugins.limit-count.init") +local policy_to_additional_properties = limit_count.policy_to_additional_properties local plugin_name = "ai-rate-limiting" @@ -90,6 +91,12 @@ local schema = { rejected_msg = { type = "string", minLength = 1 }, + policy = { + type = "string", + enum = {"local", "redis", "redis-cluster", "redis-sentinel"}, + default = "local", + }, + allow_degradation = {type = "boolean", default = false}, rules = { type = "array", items = { @@ -135,7 +142,32 @@ local schema = { { required = {"rules"}, } - } + }, + ["if"] = { + properties = { + policy = { + enum = {"redis"}, + }, + }, + }, + ["then"] = policy_to_additional_properties.redis, + ["else"] = { + ["if"] = { + properties = { + policy = { + enum = {"redis-cluster"}, + }, + }, + }, + ["then"] = policy_to_additional_properties["redis-cluster"], + ["else"] = { + ["if"] = { + properties = { policy = { enum = { "redis-sentinel" } } }, + }, + ["then"] = policy_to_additional_properties["redis-sentinel"], + }, + }, + encrypt_fields = {"redis_password", "sentinel_password"}, } local _M = { @@ -191,43 +223,49 @@ end local function transform_limit_conf(plugin_conf, instance_conf, instance_name) local limit_conf = { + _meta = plugin_conf._meta, rejected_code = plugin_conf.rejected_code, rejected_msg = plugin_conf.rejected_msg, show_limit_quota_header = plugin_conf.show_limit_quota_header, - -- we may expose those fields to ai-rate-limiting later - policy = "local", + -- counters can be shared across nodes via redis policies + policy = plugin_conf.policy or "local", key_type = "constant", - allow_degradation = false, + allow_degradation = plugin_conf.allow_degradation, sync_interval = -1, limit_header = "X-AI-RateLimit-Limit", remaining_header = "X-AI-RateLimit-Remaining", reset_header = "X-AI-RateLimit-Reset", } + local name = instance_name or "" if plugin_conf.rules and #plugin_conf.rules > 0 then limit_conf.rules = plugin_conf.rules - limit_conf._meta = plugin_conf._meta - return limit_conf + else + local key = plugin_name .. "#global" + local limit = plugin_conf.limit + local time_window = plugin_conf.time_window + if instance_conf then + name = instance_conf.name + key = instance_conf.name + limit = instance_conf.limit + time_window = instance_conf.time_window + end + limit_conf._vid = key + limit_conf.key = key + limit_conf.count = limit + limit_conf.time_window = time_window + limit_conf.limit_header = "X-AI-RateLimit-Limit-" .. name + limit_conf.remaining_header = "X-AI-RateLimit-Remaining-" .. name + limit_conf.reset_header = "X-AI-RateLimit-Reset-" .. name end - local key = plugin_name .. "#global" - local limit = plugin_conf.limit - local time_window = plugin_conf.time_window - local name = instance_name or "" - if instance_conf then - name = instance_conf.name - key = instance_conf.name - limit = instance_conf.limit - time_window = instance_conf.time_window + -- copy the redis fields straight from the policy's schema so no field is missed + local extra = policy_to_additional_properties[plugin_conf.policy] + if extra then + for k in pairs(extra.properties) do + limit_conf[k] = plugin_conf[k] + end end - limit_conf._vid = key - limit_conf.key = key - limit_conf._meta = plugin_conf._meta - limit_conf.count = limit - limit_conf.time_window = time_window - limit_conf.limit_header = "X-AI-RateLimit-Limit-" .. name - limit_conf.remaining_header = "X-AI-RateLimit-Remaining-" .. name - limit_conf.reset_header = "X-AI-RateLimit-Reset-" .. name return limit_conf end diff --git a/apisix/plugins/limit-count/init.lua b/apisix/plugins/limit-count/init.lua index 2cdfdac8571c..48d93cfd9278 100644 --- a/apisix/plugins/limit-count/init.lua +++ b/apisix/plugins/limit-count/init.lua @@ -227,6 +227,7 @@ local schema = { local schema_copy = core.table.deepcopy(schema) local _M = { + policy_to_additional_properties = policy_to_additional_properties, schema = schema, metadata_schema = metadata_schema, } diff --git a/docs/en/latest/plugins/ai-rate-limiting.md b/docs/en/latest/plugins/ai-rate-limiting.md index 15ec2314586b..acf1367a9bb1 100644 --- a/docs/en/latest/plugins/ai-rate-limiting.md +++ b/docs/en/latest/plugins/ai-rate-limiting.md @@ -55,6 +55,28 @@ The `ai-rate-limiting` Plugin enforces token-based rate limiting for requests se | instances.time_window | integer | True | | >0 | The time interval corresponding to the rate limiting `limit` in seconds for an instance. | | rejected_code | integer | False | 503 | [200, 599] | The HTTP status code returned when a request exceeding the quota is rejected. | | rejected_msg | string | False | | | The response body returned when a request exceeding the quota is rejected. | +| policy | string | False | local | [`local`, `redis`, `redis-cluster`, `redis-sentinel`] | The policy for the rate limiting counter. If it is `local`, the counter is stored in local memory. If it is `redis`, the counter is stored on a Redis instance. If it is `redis-cluster`, the counter is stored in a Redis cluster. If it is `redis-sentinel`, the counter is stored on the Redis master discovered through Sentinel. Use a Redis-based policy to share token counters across gateway nodes in a distributed deployment. | +| allow_degradation | boolean | False | false | | If true, allow APISIX to continue serving requests without the Plugin when the Redis backend is unavailable. | +| redis_host | string | False | | | The address of the Redis node. Required when `policy` is `redis`. | +| redis_port | integer | False | 6379 | [1,...] | The port of the Redis node when `policy` is `redis`. | +| redis_username | string | False | | | The username for Redis if Redis ACL is used. If you use the legacy authentication method `requirepass`, configure only the `redis_password`. Used when `policy` is `redis` or `redis-sentinel`. | +| redis_password | string | False | | | The password of the Redis node when `policy` is `redis`, `redis-cluster`, or `redis-sentinel`. | +| redis_database | integer | False | 0 | >= 0 | The database number in Redis when `policy` is `redis` or `redis-sentinel`. | +| redis_timeout | integer | False | 1000 | [1,...] | The Redis timeout value in milliseconds when `policy` is `redis` or `redis-cluster`. | +| redis_ssl | boolean | False | false | | If true, use SSL to connect to Redis when `policy` is `redis`. | +| redis_ssl_verify | boolean | False | false | | If true, verify the server SSL certificate when `policy` is `redis`. | +| redis_cluster_nodes | array[string] | False | | | The list of Redis cluster nodes with at least one address. Required when `policy` is `redis-cluster`. | +| redis_cluster_name | string | False | | | The name of the Redis cluster. Required when `policy` is `redis-cluster`. | +| redis_cluster_ssl | boolean | False | false | | If true, use SSL to connect to Redis when `policy` is `redis-cluster`. | +| redis_cluster_ssl_verify | boolean | False | false | | If true, verify the server SSL certificate when `policy` is `redis-cluster`. | +| redis_sentinels | array[object] | False | | | The list of Sentinel nodes. Required when `policy` is `redis-sentinel`. Each item must contain `host` and `port`. | +| redis_master_name | string | False | | | The Redis master name monitored by Sentinel. Required when `policy` is `redis-sentinel`. | +| sentinel_username | string | False | | | The username for Sentinel when `policy` is `redis-sentinel`. | +| sentinel_password | string | False | | | The password for Sentinel when `policy` is `redis-sentinel`. | +| redis_role | string | False | master | [`master`, `slave`] | The Redis role selected through Sentinel when `policy` is `redis-sentinel`. | +| redis_connect_timeout | integer | False | 1000 | [1,...] | The Redis connection timeout in milliseconds when `policy` is `redis-sentinel`. | +| redis_read_timeout | integer | False | 1000 | [1,...] | The Redis read timeout in milliseconds when `policy` is `redis-sentinel`. | +| redis_keepalive_timeout | integer | False | 60000 | [1,...] | The Redis keepalive timeout in milliseconds when `policy` is `redis-sentinel`. | | rules | array[object] | False | | | An array of rate-limiting rules that are applied sequentially. If configured, this takes precedence over `limit` and `time_window`. | | rules.count | integer or string | True | | >0 or variable expression | The maximum number of tokens allowed within a given time interval. Can be a static integer or a variable expression like `$http_custom_limit`. | | rules.time_window | integer or string | True | | >0 or variable expression | The time interval corresponding to the rate limiting `count` in seconds. Can be a static integer or a variable expression. | @@ -286,6 +308,63 @@ You should receive a response similar to the following: If the rate limiting quota of 300 prompt tokens has been consumed in a 30-second window, all additional requests will be rejected. +### Share Rate Limiting Counters Across Nodes with Redis + +By default, the `local` policy keeps token counters in each node's shared memory, so counters are not shared across a multi-node deployment and the effective quota scales with the number of nodes. Set `policy` to `redis`, `redis-cluster`, or `redis-sentinel` to centralize counters and enforce a single quota across all nodes. + +The following example demonstrates how you can use the `redis` policy to share the rate limiting counter across gateway nodes. + +Create a Route as such and update with your LLM providers, models, API keys, endpoints, and Redis address, if applicable: + +```shell +curl "http://127.0.0.1:9180/apisix/admin/routes" -X PUT \ + -H "X-API-KEY: ${admin_key}" \ + -d '{ + "id": "ai-rate-limiting-route", + "uri": "/anything", + "methods": ["POST"], + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { + "header": { + "Authorization": "Bearer '"$OPENAI_API_KEY"'" + } + }, + "options": { + "model": "gpt-35-turbo-instruct", + "max_tokens": 512, + "temperature": 1.0 + } + }, + "ai-rate-limiting": { + "limit": 300, + "time_window": 30, + "limit_strategy": "total_tokens", + "policy": "redis", + "redis_host": "127.0.0.1", + "redis_port": 6379, + "redis_database": 1 + } + } + }' +``` + +Send a POST request to the Route with a sample question in the request body from any node: + +```shell +curl "http://127.0.0.1:9080/anything" -X POST \ + -H "Content-Type: application/json" \ + -d '{ + "messages": [ + { "role": "system", "content": "You are a mathematician" }, + { "role": "user", "content": "What is 1+1?" } + ] + }' +``` + +The consumed tokens are counted against the shared Redis counter. Once the quota of 300 tokens is consumed in a 30-second window, additional requests are rejected on every node. + ### Rate Limit One Instance Among Multiple The following example demonstrates how you can use `ai-proxy-multi` to configure two models for load balancing, forwarding 80% of the traffic to one instance and 20% to the other. Additionally, use `ai-rate-limiting` to configure token-based rate limiting on the instance that receives 80% of the traffic, such that when the configured quota is fully consumed, the additional traffic will be forwarded to the other instance. diff --git a/t/plugin/ai-rate-limiting.t b/t/plugin/ai-rate-limiting.t index 4ca49b5e9317..9fcadbc1c785 100644 --- a/t/plugin/ai-rate-limiting.t +++ b/t/plugin/ai-rate-limiting.t @@ -1538,3 +1538,218 @@ X-AI-Fixture: openai/chat-model-echo.json --- error_code: 200 --- no_error_log [error] + + + +=== TEST 35: schema check, redis policy requires redis_host +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-rate-limiting") + local ok, err = plugin.check_schema({ + limit = 30, + time_window = 60, + policy = "redis", + }) + if not ok then + ngx.say(err) + else + ngx.say("passed") + end + } + } +--- response_body +then clause did not match + + + +=== TEST 36: flush redis and set route with redis policy +--- config + location /t { + content_by_lua_block { + local redis = require("resty.redis") + local red = redis:new() + red:set_timeout(1000) + local ok, err = red:connect("127.0.0.1", 6379) + if not ok then + ngx.say("failed to connect: ", err) + return + end + red:flushall() + + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/ai", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { + "header": { + "Authorization": "Bearer token" + } + }, + "options": { + "model": "gpt-35-turbo-instruct", + "max_tokens": 512, + "temperature": 1.0 + }, + "override": { + "endpoint": "http://127.0.0.1:1980" + }, + "ssl_verify": false + }, + "ai-rate-limiting": { + "limit": 30, + "time_window": 60, + "policy": "redis", + "redis_host": "127.0.0.1", + "redis_port": 6379, + "redis_database": 1 + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 37: redis policy shares counter and rejects the 4th request +--- pipelined_requests eval +[ + "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }", + "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }", + "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }", + "POST /ai\n" . "{ \"messages\": [ { \"role\": \"system\", \"content\": \"You are a mathematician\" }, { \"role\": \"user\", \"content\": \"What is 1+1?\"} ] }", +] +--- more_headers +Authorization: Bearer token +X-AI-Fixture: openai/chat-model-echo.json +--- error_code eval +[200, 200, 200, 503] + + + +=== TEST 38: counter is stored in redis, not local memory +--- config + location /t { + content_by_lua_block { + local redis = require("resty.redis") + local red = redis:new() + red:set_timeout(1000) + local ok, err = red:connect("127.0.0.1", 6379) + if not ok then + ngx.say("failed to connect: ", err) + return + end + red:select(1) + local keys, err = red:keys("*ai-rate-limiting*") + if not keys then + ngx.say("failed to get keys: ", err) + return + end + ngx.say(#keys > 0 and "counter in redis" or "no counter") + } + } +--- response_body +counter in redis + + + +=== TEST 39: schema check, redis-sentinel accepts redis master auth +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-rate-limiting") + local ok, err = plugin.check_schema({ + limit = 30, + time_window = 60, + policy = "redis-sentinel", + redis_sentinels = { + { host = "127.0.0.1", port = 26379 }, + }, + redis_master_name = "mymaster", + redis_username = "alice", + redis_password = "somepassword", + sentinel_username = "bob", + sentinel_password = "sentinelpass", + }) + ngx.say(ok and "passed" or err) + } + } +--- response_body +passed + + + +=== TEST 40: redis_password is encrypted at rest +--- yaml_config +apisix: + data_encryption: + enable_encrypt_fields: true + keyring: + - edd1c9f0985e76a2 +--- config + location /t { + content_by_lua_block { + local json = require("toolkit.json") + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/ai", + "plugins": { + "ai-rate-limiting": { + "limit": 30, + "time_window": 60, + "policy": "redis", + "redis_host": "127.0.0.1", + "redis_port": 6379, + "redis_password": "somepassword" + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + if code >= 300 then + ngx.status = code + ngx.say(body) + return + end + ngx.sleep(0.1) + + -- admin api returns the decrypted password + local code, _, res = t('/apisix/admin/routes/1', ngx.HTTP_GET) + res = json.decode(res) + ngx.say(res.value.plugins["ai-rate-limiting"].redis_password) + + -- etcd stores the encrypted password + local etcd = require("apisix.core.etcd") + local res = assert(etcd.get('/routes/1')) + local stored = res.body.node.value.plugins["ai-rate-limiting"].redis_password + ngx.say(stored ~= "somepassword" and "encrypted" or "plaintext") + } + } +--- response_body +somepassword +encrypted