feat(llm): add auth.allow_override option for llm auth functionality (#…

…13493)

Add `allow_override` option to allow overriding the upstream model auth
parameter or header from the caller's request.

changelog/unreleased/kong/ai-proxy-add-allow-override-opt.yml

@@ -0,0 +1,4 @@
+message: |
+  **AI-proxy-plugin**: Add `allow_auth_override` option to allow overriding the upstream model auth parameter or header from the caller's request.
+scope: Plugin
+type: feature

kong/clustering/compat/removed_fields.lua

@@ -175,6 +175,7 @@ return {
       "model.options.bedrock",
       "auth.aws_access_key_id",
       "auth.aws_secret_access_key",
+      "auth.allow_auth_override",
       "model_name_header",
     },
     ai_prompt_decorator = {

kong/llm/drivers/anthropic.lua

@@ -472,13 +472,18 @@ function _M.configure_request(conf)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_header_name and auth_header_value then
-    kong.service.request.set_header(auth_header_name, auth_header_value)
+    local exist_value = kong.request.get_header(auth_header_name)
+    if exist_value == nil or not conf.auth.allow_auth_override then
+      kong.service.request.set_header(auth_header_name, auth_header_value)
+    end
   end
 
   if auth_param_name and auth_param_value and auth_param_location == "query" then
     local query_table = kong.request.get_query()
-    query_table[auth_param_name] = auth_param_value
-    kong.service.request.set_query(query_table)
+    if query_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      query_table[auth_param_name] = auth_param_value
+      kong.service.request.set_query(query_table)
+    end
   end
 
   -- if auth_param_location is "form", it will have already been set in a pre-request hook

kong/llm/drivers/azure.lua

@@ -131,17 +131,23 @@ function _M.configure_request(conf)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_header_name and auth_header_value then
-    kong.service.request.set_header(auth_header_name, auth_header_value)
+    local exist_value = kong.request.get_header(auth_header_name)
+    if exist_value == nil or not conf.auth.allow_auth_override then
+      kong.service.request.set_header(auth_header_name, auth_header_value)
+    end
   end
 
+
   local query_table = kong.request.get_query()
 
   -- technically min supported version
   query_table["api-version"] = kong.request.get_query_arg("api-version")
                             or (conf.model.options and conf.model.options.azure_api_version)
 
   if auth_param_name and auth_param_value and auth_param_location == "query" then
-    query_table[auth_param_name] = auth_param_value
+    if query_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      query_table[auth_param_name] = auth_param_value
+    end
   end
 
   kong.service.request.set_query(query_table)

kong/llm/drivers/cohere.lua

@@ -480,13 +480,18 @@ function _M.configure_request(conf)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_header_name and auth_header_value then
-    kong.service.request.set_header(auth_header_name, auth_header_value)
+    local exist_value = kong.request.get_header(auth_header_name)
+    if exist_value == nil or not conf.auth.allow_auth_override then
+      kong.service.request.set_header(auth_header_name, auth_header_value)
+    end
   end
 
   if auth_param_name and auth_param_value and auth_param_location == "query" then
     local query_table = kong.request.get_query()
-    query_table[auth_param_name] = auth_param_value
-    kong.service.request.set_query(query_table)
+    if query_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      query_table[auth_param_name] = auth_param_value
+      kong.service.request.set_query(query_table)
+    end
   end
 
   -- if auth_param_location is "form", it will have already been set in a pre-request hook

kong/llm/drivers/llama2.lua

@@ -277,13 +277,18 @@ function _M.configure_request(conf)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_header_name and auth_header_value then
-    kong.service.request.set_header(auth_header_name, auth_header_value)
+    local exist_value = kong.request.get_header(auth_header_name)
+    if exist_value == nil or not conf.auth.allow_auth_override then
+      kong.service.request.set_header(auth_header_name, auth_header_value)
+    end
   end
 
   if auth_param_name and auth_param_value and auth_param_location == "query" then
     local query_table = kong.request.get_query()
-    query_table[auth_param_name] = auth_param_value
-    kong.service.request.set_query(query_table)
+    if query_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      query_table[auth_param_name] = auth_param_value
+      kong.service.request.set_query(query_table)
+    end
   end
 
   -- if auth_param_location is "form", it will have already been set in a pre-request hook

kong/llm/drivers/mistral.lua

@@ -172,13 +172,18 @@ function _M.configure_request(conf)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_header_name and auth_header_value then
-    kong.service.request.set_header(auth_header_name, auth_header_value)
+    local exist_value = kong.request.get_header(auth_header_name)
+    if exist_value == nil or not conf.auth.allow_auth_override then
+      kong.service.request.set_header(auth_header_name, auth_header_value)
+    end
   end
 
   if auth_param_name and auth_param_value and auth_param_location == "query" then
     local query_table = kong.request.get_query()
-    query_table[auth_param_name] = auth_param_value
-    kong.service.request.set_query(query_table)
+    if query_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      query_table[auth_param_name] = auth_param_value
+      kong.service.request.set_query(query_table)
+    end
   end
 
   -- if auth_param_location is "form", it will have already been set in a pre-request hook

kong/llm/drivers/openai.lua

@@ -213,13 +213,18 @@ function _M.configure_request(conf)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_header_name and auth_header_value then
-    kong.service.request.set_header(auth_header_name, auth_header_value)
+    local exist_value = kong.request.get_header(auth_header_name)
+    if exist_value == nil or not conf.auth.allow_auth_override then
+      kong.service.request.set_header(auth_header_name, auth_header_value)
+    end
   end
 
   if auth_param_name and auth_param_value and auth_param_location == "query" then
     local query_table = kong.request.get_query()
-    query_table[auth_param_name] = auth_param_value
-    kong.service.request.set_query(query_table)
+    if query_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      query_table[auth_param_name] = auth_param_value
+      kong.service.request.set_query(query_table)
+    end
   end
 
   -- if auth_param_location is "form", it will have already been set in a global pre-request hook

kong/llm/drivers/shared.lua

@@ -528,7 +528,9 @@ function _M.pre_request(conf, request_table)
   local auth_param_location = conf.auth and conf.auth.param_location
 
   if auth_param_name and auth_param_value and auth_param_location == "body" and request_table then
-    request_table[auth_param_name] = auth_param_value
+    if request_table[auth_param_name] == nil or not conf.auth.allow_auth_override then
+      request_table[auth_param_name] = auth_param_value
+    end
   end
 
   -- retrieve the plugin name

kong/llm/schemas/init.lua

@@ -97,6 +97,11 @@ local auth_schema = {
         required = false,
         encrypted = true,
         referenceable = true }},
+    { allow_auth_override = {
+        type = "boolean",
+        description = "If enabled, the authorization header or parameter can be overridden in the request by the value configured in the plugin.",
+        required = false,
+        default = true }},
   }
 }
 
@@ -237,6 +242,11 @@ return {
     { logging = logging_schema },
   },
   entity_checks = {
+    { conditional =  { if_field = "model.provider",
+                          if_match = { one_of = { "bedrock", "gemini" } },
+                          then_field = "auth.allow_auth_override",
+                          then_match = { eq = false },
+                          then_err = "bedrock and gemini only support auth.allow_auth_override = false" }},
     { mutually_required = { "auth.header_name", "auth.header_value" }, },
     { mutually_required = { "auth.param_name", "auth.param_value", "auth.param_location" }, },
 

spec/02-integration/09-hybrid_mode/09-config-compat_spec.lua

@@ -495,6 +495,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
               header_value = "value",
               gcp_service_account_json = '{"service": "account"}',
               gcp_use_service_account = true,
+              allow_auth_override = false,
             },
             model = {
               name = "any-model-name",
@@ -526,6 +527,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
         -- gemini fields
         expected.config.auth.gcp_service_account_json = nil
         expected.config.auth.gcp_use_service_account = nil
+        expected.config.auth.allow_auth_override = nil
         expected.config.model.options.gemini = nil
 
         -- bedrock fields
@@ -562,6 +564,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
                 header_value = "value",
                 gcp_service_account_json = '{"service": "account"}',
                 gcp_use_service_account = true,
+                allow_auth_override = false,
               },
               model = {
                 name = "any-model-name",
@@ -625,6 +628,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
                 header_value = "value",
                 gcp_service_account_json = '{"service": "account"}',
                 gcp_use_service_account = true,
+                allow_auth_override = false,
               },
               model = {
                 name = "any-model-name",
@@ -720,6 +724,7 @@ describe("CP/DP config compat transformations #" .. strategy, function()
         -- bedrock fields
         expected.config.auth.aws_access_key_id = nil
         expected.config.auth.aws_secret_access_key = nil
+        expected.config.auth.allow_auth_override = nil
         expected.config.model.options.bedrock = nil
 
         do_assert(uuid(), "3.7.0", expected)

spec/03-plugins/38-ai-proxy/00-config_spec.lua

@@ -289,4 +289,59 @@ describe(PLUGIN_NAME .. ": (schema)", function()
     assert.is_truthy(ok)
   end)
 
+  it("bedrock model can not support ath.allowed_auth_override", function()
+    local config = {
+      route_type = "llm/v1/chat",
+      auth = {
+        param_name = "apikey",
+        param_value = "key",
+        param_location = "query",
+        header_name = "Authorization",
+        header_value = "Bearer token",
+        allow_auth_override = true,
+      },
+      model = {
+        name = "bedrock",
+        provider = "bedrock",
+        options = {
+          max_tokens = 256,
+          temperature = 1.0,
+          upstream_url = "http://nowhere",
+        },
+      },
+    }
+
+    local ok, err = validate(config)
+
+    assert.is_falsy(ok)
+    assert.is_truthy(err)
+  end)
+
+  it("gemini model can not support ath.allowed_auth_override", function()
+    local config = {
+      route_type = "llm/v1/chat",
+      auth = {
+        param_name = "apikey",
+        param_value = "key",
+        param_location = "query",
+        header_name = "Authorization",
+        header_value = "Bearer token",
+        allow_auth_override = true,
+      },
+      model = {
+        name = "gemini",
+        provider = "gemini",
+        options = {
+          max_tokens = 256,
+          temperature = 1.0,
+          upstream_url = "http://nowhere",
+        },
+      },
+    }
+
+    local ok, err = validate(config)
+
+    assert.is_falsy(ok)
+    assert.is_truthy(err)
+  end)
 end)