From c9741da3d293f8ded06e87838cf4c205350e4b69 Mon Sep 17 00:00:00 2001 From: Jun Ouyang Date: Tue, 13 Aug 2024 16:06:09 +0800 Subject: [PATCH] chore(llm): fix code --- kong/llm/drivers/azure.lua | 10 ++++++++-- kong/llm/drivers/cohere.lua | 11 ++++++++--- kong/llm/drivers/gemini.lua | 12 +++++++++--- kong/llm/drivers/llama2.lua | 11 ++++++++--- kong/llm/drivers/mistral.lua | 11 ++++++++--- kong/llm/drivers/openai.lua | 11 ++++++++--- 6 files changed, 49 insertions(+), 17 deletions(-) diff --git a/kong/llm/drivers/azure.lua b/kong/llm/drivers/azure.lua index 343904ffad249..28487d6538d63 100644 --- a/kong/llm/drivers/azure.lua +++ b/kong/llm/drivers/azure.lua @@ -131,9 +131,13 @@ 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.can_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 @@ -141,7 +145,9 @@ function _M.configure_request(conf) 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.can_override then + query_table[auth_param_name] = auth_param_value + end end kong.service.request.set_query(query_table) diff --git a/kong/llm/drivers/cohere.lua b/kong/llm/drivers/cohere.lua index d25764164086d..b16ee6cc96210 100644 --- a/kong/llm/drivers/cohere.lua +++ b/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.can_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.can_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 diff --git a/kong/llm/drivers/gemini.lua b/kong/llm/drivers/gemini.lua index d386961997f44..b5fceac08a8a9 100644 --- a/kong/llm/drivers/gemini.lua +++ b/kong/llm/drivers/gemini.lua @@ -390,14 +390,20 @@ function _M.configure_request(conf, identity_interface) -- DBO restrictions makes sure that only one of these auth blocks runs in one plugin config 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.can_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.can_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 -- if we're passed a GCP SDK, for cloud identity / SSO, use it appropriately diff --git a/kong/llm/drivers/llama2.lua b/kong/llm/drivers/llama2.lua index 0526453f8a52e..cc9859c0a8dc8 100644 --- a/kong/llm/drivers/llama2.lua +++ b/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.can_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.can_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 diff --git a/kong/llm/drivers/mistral.lua b/kong/llm/drivers/mistral.lua index d1d2303b6919d..a16c1d9306662 100644 --- a/kong/llm/drivers/mistral.lua +++ b/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.can_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.can_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 diff --git a/kong/llm/drivers/openai.lua b/kong/llm/drivers/openai.lua index 52df1910586ea..1433d6f17ae91 100644 --- a/kong/llm/drivers/openai.lua +++ b/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.can_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.can_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