feat(ai-proxy): add provider-aware max_tokens override with priority control (#13251)

Author: nic-6443

apisix/plugins/ai-protocols/init.lua

@@ -65,6 +65,7 @@ function _M.get(name)
 end
 
 
+
 --- Find a converter that can bridge from client_protocol to a protocol
 -- supported by the driver. Delegates to the converters registry.
 -- @param client_protocol string The detected client protocol

apisix/plugins/ai-providers/aimlapi.lua

@@ -15,12 +15,23 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         host = "api.aimlapi.com",
         port = 443,
         capabilities = {
-            ["openai-chat"] = { path = "/chat/completions" },
+            ["openai-chat"] = {
+                path = "/chat/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
         },
     }
 )

apisix/plugins/ai-providers/anthropic.lua

@@ -15,13 +15,36 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
+
+local function rewrite_messages_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         host = "api.anthropic.com",
         port = 443,
         capabilities = {
-            ["openai-chat"] = { path = "/v1/chat/completions" },
-            ["anthropic-messages"] = { path = "/v1/messages" },
+            ["openai-chat"] = {
+                path = "/v1/chat/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
+            ["anthropic-messages"] = {
+                path = "/v1/messages",
+                rewrite_request_body = rewrite_messages_request_body,
+            },
         },
     }
 )

apisix/plugins/ai-providers/azure-openai.lua

@@ -15,12 +15,23 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         port = 443,
         remove_model = true,
         capabilities = {
-            ["openai-chat"] = { path = "/completions" },
+            ["openai-chat"] = {
+                path = "/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
         },
     }
 )

apisix/plugins/ai-providers/base.lua

@@ -173,7 +173,7 @@ function _M.build_request(self, ctx, conf, request_body, opts)
                           or opts.target_host or self.host,
     }
 
-    -- Inject model options
+    -- Inject model options (flat overwrite)
     if opts.model_options then
         for opt, val in pairs(opts.model_options) do
             if request_body[opt] ~= nil then
@@ -182,6 +182,15 @@ function _M.build_request(self, ctx, conf, request_body, opts)
             request_body[opt] = val
         end
     end
+
+    -- Apply request body override via provider capability hook
+    if opts.override_request_body then
+        local cap = self.capabilities and self.capabilities[ctx.ai_target_protocol]
+        if cap and cap.rewrite_request_body then
+            cap.rewrite_request_body(request_body, opts.override_request_body,
+                                     opts.request_body_force_override)
+        end
+    end
     params.body = request_body
 
     if self.remove_model then

apisix/plugins/ai-providers/deepseek.lua

@@ -15,12 +15,23 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         host = "api.deepseek.com",
         port = 443,
         capabilities = {
-            ["openai-chat"] = { path = "/chat/completions" },
+            ["openai-chat"] = {
+                path = "/chat/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
         },
     }
 )

apisix/plugins/ai-providers/gemini.lua

@@ -15,12 +15,23 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_completion_tokens == nil then
+            body.max_completion_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         host = "generativelanguage.googleapis.com",
         port = 443,
         capabilities = {
-            ["openai-chat"] = { path = "/v1beta/openai/chat/completions" },
+            ["openai-chat"] = {
+                path = "/v1beta/openai/chat/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
         },
     }
 )

apisix/plugins/ai-providers/openai-compatible.lua

@@ -15,10 +15,33 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
+
+local function rewrite_responses_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_output_tokens == nil then
+            body.max_output_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new({
     capabilities = {
-        ["openai-chat"]       = { path = "/v1/chat/completions" },
-        ["openai-responses"]  = { path = "/v1/responses" },
+        ["openai-chat"]       = {
+            path = "/v1/chat/completions",
+            rewrite_request_body = rewrite_chat_request_body,
+        },
+        ["openai-responses"]  = {
+            path = "/v1/responses",
+            rewrite_request_body = rewrite_responses_request_body,
+        },
         ["openai-embeddings"] = { path = "/v1/embeddings" },
     },
 })

apisix/plugins/ai-providers/openai.lua

@@ -15,13 +15,37 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or (body.max_completion_tokens == nil and body.max_tokens == nil) then
+            body.max_completion_tokens = override.max_tokens
+            body.max_tokens = nil
+        end
+    end
+end
+
+
+local function rewrite_responses_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_output_tokens == nil then
+            body.max_output_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         host = "api.openai.com",
         port = 443,
         capabilities = {
-            ["openai-chat"]       = { path = "/v1/chat/completions" },
-            ["openai-responses"]  = { path = "/v1/responses" },
+            ["openai-chat"]       = {
+                path = "/v1/chat/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
+            ["openai-responses"]  = {
+                path = "/v1/responses",
+                rewrite_request_body = rewrite_responses_request_body,
+            },
             ["openai-embeddings"] = { path = "/v1/embeddings" },
         },
     }

apisix/plugins/ai-providers/openrouter.lua

@@ -15,12 +15,23 @@
 -- limitations under the License.
 --
 
+local function rewrite_chat_request_body(body, override, force)
+    if override.max_tokens then
+        if force or body.max_tokens == nil then
+            body.max_tokens = override.max_tokens
+        end
+    end
+end
+
 return require("apisix.plugins.ai-providers.base").new(
     {
         host = "openrouter.ai",
         port = 443,
         capabilities = {
-            ["openai-chat"] = { path = "/api/v1/chat/completions" },
+            ["openai-chat"] = {
+                path = "/api/v1/chat/completions",
+                rewrite_request_body = rewrite_chat_request_body,
+            },
         },
     }
 )