Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
nic-6443 merged PR #13191: URL: https://github.com/apache/apisix/pull/13191 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
AlinsRan commented on code in PR #13191:
URL: https://github.com/apache/apisix/pull/13191#discussion_r3063271363
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +312,57 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
+-- build environment: safe math + usage variables (missing vars default to
0)
+local env = setmetatable({}, {
+__index = function(_, k)
+local v = expr_safe_env[k]
+if v ~= nil then
+return v
+end
+return 0
+end
+})
+for k, v in pairs(raw) do
+if type(v) == "number" and not expr_safe_env[k] then
+env[k] = v
+end
+end
+local fn, err = load(fn_code, "cost_expr", "t", env)
+if not fn then
+return nil, "failed to compile cost_expr: " .. err
+end
+local ok, result = pcall(fn)
+if not ok then
+return nil, "failed to evaluate cost_expr: " .. result
+end
+if type(result) ~= "number" then
+return nil, "cost_expr must return a number, got: " .. type(result)
+end
+if result ~= result or result == math_huge or result == -math_huge then
Review Comment:
`result ~= result` typo?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
AlinsRan commented on code in PR #13191:
URL: https://github.com/apache/apisix/pull/13191#discussion_r3063271363
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +312,57 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
+-- build environment: safe math + usage variables (missing vars default to
0)
+local env = setmetatable({}, {
+__index = function(_, k)
+local v = expr_safe_env[k]
+if v ~= nil then
+return v
+end
+return 0
+end
+})
+for k, v in pairs(raw) do
+if type(v) == "number" and not expr_safe_env[k] then
+env[k] = v
+end
+end
+local fn, err = load(fn_code, "cost_expr", "t", env)
+if not fn then
+return nil, "failed to compile cost_expr: " .. err
+end
+local ok, result = pcall(fn)
+if not ok then
+return nil, "failed to evaluate cost_expr: " .. result
+end
+if type(result) ~= "number" then
+return nil, "cost_expr must return a number, got: " .. type(result)
+end
+if result ~= result or result == math_huge or result == -math_huge then
Review Comment:
`result ~= result` typo?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
nic-6443 commented on PR #13191: URL: https://github.com/apache/apisix/pull/13191#issuecomment-4214855367 Fixed the CI failure in TEST 13: when the expression evaluates to a negative value that gets clamped to 0, calling `rate_limit()` with `cost=0` triggers an assertion failure in `resty.limit.count`'s `incoming` function (`dict:incr(key, 0, ...)`). The fix skips the `rate_limit()` call entirely when `used_tokens == 0` since there's nothing to deduct. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
nic-6443 commented on code in PR #13191:
URL: https://github.com/apache/apisix/pull/13191#discussion_r3057789724
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+math = math,
+abs = math.abs,
+ceil = math.ceil,
+floor = math.floor,
+max = math.max,
+min = math.min,
+}
+
+local function compile_cost_expr(expr_str)
+local fn_code = "return " .. expr_str
+-- validate syntax by loading first
+local fn, err = load(fn_code, "cost_expr", "t", expr_safe_env)
+if not fn then
+return nil, err
Review Comment:
The expression is configured by admins via the plugin config, not by
end-users or API callers. Admins already have full control over the gateway. A
strict arithmetic-only validator would add significant complexity. This is
consistent with the upstream implementation.
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
Review Comment:
Valid optimization opportunity. The expression strings are typically very
short and load() in LuaJIT is fast. This is consistent with the upstream
implementation. Caching the compiled function can be done as a follow-up
optimization if profiling shows it matters.
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
+-- build environment: safe math + usage variables (missing vars default to
0)
+local env = setmetatable({}, {
+__index = function(_, k)
+local v = expr_safe_env[k]
+if v ~= nil then
+return v
+end
+return 0
+end
+})
+for k, v in pairs(raw) do
+if type(v) == "number" then
+env[k] = v
+end
+end
+local fn, err = load(fn_code, "cost_expr", "t", env)
+if not fn then
+return nil, "failed to compile cost_expr: " .. err
+end
+local ok, result = pcall(fn)
+if not ok then
+return nil, "failed to evaluate cost_expr: " .. result
+end
+if type(result) ~= "number" then
+return nil, "cost_expr must return a number, got: " .. type(result)
+end
+return math_floor(result + 0.5)
Review Comment:
Fixed — now rejects NaN/inf and clamps negative results to 0. Added test
case (TEST 13) to verify.
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -61,10 +65,19 @@ local schema = {
show_limit_quota_header = {type = "boolean", default = true},
limit_strategy = {
type = "string",
-enum = {"total_tokens", "prompt_tokens", "completion_tokens"},
+enum = {"total_tokens", "prompt_tokens", "completion_tokens",
"expression"},
default = "total_tokens",
description = "The strategy to limit the tokens"
},
+cost_expr = {
+type = "string",
+minLength = 1,
+description = "Lua arithmetic expression for dynamic token cost
calculation. "
+.. "Variables are injected from the LLM API raw usage response
fields. "
+.. "Missing variables default to 0. "
+.. "Only valid when limit_strategy is 'expression'. "
+.. "Example: input_tokens + cache_creation_input_tokens +
output_tokens",
+},
Review Comment:
Having cost_expr set when strategy is not expression is harmless (ignored).
Rejecting it would make config management harder (e.g., switching strategies
would require removing/adding the field). This follows the
be-liberal-in-what-you-accept principle.
##
t/plugin/ai-rate-limiting-expression.t:
##
@@ -0,0 +1,551 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+BEGIN {
+$ENV{TEST_ENABLE_CONTROL_API_V1} = "0";
+}
+
+use t::AP
Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
nic-6443 commented on code in PR #13191:
URL: https://github.com/apache/apisix/pull/13191#discussion_r3057787472
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+math = math,
+abs = math.abs,
+ceil = math.ceil,
+floor = math.floor,
+max = math.max,
+min = math.min,
+}
Review Comment:
Good catch. Fixed — raw usage keys that collide with expr_safe_env names are
now skipped.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
Re: [PR] feat(ai-rate-limiting): add expression-based limit strategy [apisix]
Copilot commented on code in PR #13191:
URL: https://github.com/apache/apisix/pull/13191#discussion_r3057744723
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -136,8 +149,42 @@ local limit_conf_cache = core.lrucache.new({
})
+-- safe math functions allowed in cost expressions
+local expr_safe_env = {
+math = math,
+abs = math.abs,
+ceil = math.ceil,
+floor = math.floor,
+max = math.max,
+min = math.min,
+}
Review Comment:
The runtime env allows upstream-provided usage keys to overwrite
reserved/safe names (e.g., a usage field named `math`/`abs` would shadow the
safe functions), which can break evaluation and potentially bypass rate
limiting. Consider preventing `raw` keys from overwriting names present in
`expr_safe_env` (skip those keys), and avoid exposing the full `math` table if
the intent is to allow only a small whitelist (e.g., provide a minimal `math`
table with only `abs/ceil/floor/max/min`).
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
+-- build environment: safe math + usage variables (missing vars default to
0)
+local env = setmetatable({}, {
+__index = function(_, k)
+local v = expr_safe_env[k]
+if v ~= nil then
+return v
+end
+return 0
+end
+})
+for k, v in pairs(raw) do
+if type(v) == "number" then
Review Comment:
The runtime env allows upstream-provided usage keys to overwrite
reserved/safe names (e.g., a usage field named `math`/`abs` would shadow the
safe functions), which can break evaluation and potentially bypass rate
limiting. Consider preventing `raw` keys from overwriting names present in
`expr_safe_env` (skip those keys), and avoid exposing the full `math` table if
the intent is to allow only a small whitelist (e.g., provide a minimal `math`
table with only `abs/ceil/floor/max/min`).
```suggestion
local safe_math = {
abs = math.abs,
ceil = math.ceil,
floor = math.floor,
max = math.max,
min = math.min,
}
local reserved_names = {
math = true,
abs = true,
ceil = true,
floor = true,
max = true,
min = true,
}
-- build environment: safe math + usage variables (missing vars default
to 0)
local env = setmetatable({}, {
__index = function(_, k)
if k == "math" then
return safe_math
end
local v = safe_math[k]
if v ~= nil then
return v
end
v = expr_safe_env[k]
if v ~= nil then
return v
end
return 0
end
})
for k, v in pairs(raw) do
if type(v) == "number" and not reserved_names[k] and
expr_safe_env[k] == nil then
```
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
+-- build environment: safe math + usage variables (missing vars default to
0)
+local env = setmetatable({}, {
+__index = function(_, k)
+local v = expr_safe_env[k]
+if v ~= nil then
+return v
+end
+return 0
+end
+})
+for k, v in pairs(raw) do
+if type(v) == "number" then
+env[k] = v
+end
+end
+local fn, err = load(fn_code, "cost_expr", "t", env)
+if not fn then
+return nil, "failed to compile cost_expr: " .. err
+end
+local ok, result = pcall(fn)
Review Comment:
The expression is recompiled (`load`) on every request. This adds avoidable
overhead in a hot path. Consider caching the compiled function keyed by
`conf.cost_expr` (e.g., via an `lrucache`) and only rebuilding the per-request
variable environment (or setting a per-call env via `setfenv`/equivalent) while
reusing the compiled chunk.
##
apisix/plugins/ai-rate-limiting.lua:
##
@@ -264,7 +311,51 @@ function _M.check_instance_status(conf, ctx, instance_name)
end
+local function eval_cost_expr(conf_cost_expr, raw)
+local fn_code = "return " .. conf_cost_expr
Review Comment:
The expression is recompiled (`load`) on every request. This adds avoidable
overhead in a hot path. Consider caching the compiled function keyed by
`conf.cost_expr` (e.g., via an `lrucache`) and only rebuilding the per-request
variable environment (or setting a per-call env via `setfenv`/equivalent) while
reusing the compiled chunk.
##
