Skip to content

Commit 3ba27f6

Browse files
shreemaan-abhishekshreemaan-abhishek
authored
chore: header related fixes (#12961)
1 parent 680e4af commit 3ba27f6

File tree

4 files changed

+158
-41
lines changed

4 files changed

+158
-41
lines changed

‎apisix/plugins/ai-drivers/openai-base.lua‎

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,6 +40,7 @@ local math = math
4040
local os = os
4141
local ipairs = ipairs
4242
local setmetatable = setmetatable
43+
local str_lower = string.lower
4344

4445
local HTTP_INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR
4546
local HTTP_GATEWAY_TIMEOUT = ngx.HTTP_GATEWAY_TIMEOUT
@@ -193,6 +194,7 @@ local function read_response(conf, ctx, res, response_filter)
193194
end
194195
ctx.var.llm_prompt_tokens = ctx.ai_token_usage.prompt_tokens or 0
195196
ctx.var.llm_completion_tokens = ctx.ai_token_usage.completion_tokens or 0
197+
196198
if type(res_body.choices) == "table" and #res_body.choices > 0 then
197199
local contents = {}
198200
for _, choice in ipairs(res_body.choices) do
@@ -209,6 +211,31 @@ local function read_response(conf, ctx, res, response_filter)
209211
plugin.lua_response_filter(ctx, headers, raw_res_body)
210212
end
211213

214+
-- We want to forward all client headers to the LLM upstream by copying headers from the client
215+
-- but copying content-length is destructive, similarly some headers like `host`
216+
-- should not be forwarded either
217+
local function construct_forward_headers(ext_opts_headers, ctx)
218+
local blacklist = {
219+
"host",
220+
"content-length"
221+
}
222+
223+
-- make header keys lower case to overwrite downstream headers correctly,
224+
-- because downstream headers are lower case
225+
local opts_headers_lower = {}
226+
for k, v in pairs(ext_opts_headers or {}) do
227+
opts_headers_lower[str_lower(k)] = v
228+
end
229+
local headers = core.table.merge(core.request.headers(ctx), opts_headers_lower)
230+
headers["Content-Type"] = "application/json"
231+
232+
for _, h in ipairs(blacklist) do
233+
headers[h] = nil
234+
end
235+
236+
return headers
237+
end
238+
212239

213240
local gcp_access_token_cache = lrucache.new(1024 * 4)
214241

@@ -297,8 +324,7 @@ function _M.request(self, ctx, conf, request_table, extra_opts)
297324

298325
local path = (parsed_url and parsed_url.path or self.path)
299326

300-
local headers = auth.header or {}
301-
headers["Content-Type"] = "application/json"
327+
local headers = construct_forward_headers(auth.header or {}, ctx)
302328
if token then
303329
headers["Authorization"] = "Bearer " .. token
304330
end

‎t/plugin/ai-proxy-multi3.t‎

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -70,18 +70,6 @@ add_block_preprocessor(sub {
7070
local body, err = ngx.req.get_body_data()
7171
body, err = json.decode(body)
7272
73-
local test_type = ngx.req.get_headers()["test-type"]
74-
if test_type == "options" then
75-
if body.foo == "bar" then
76-
ngx.status = 200
77-
ngx.say("options works")
78-
else
79-
ngx.status = 500
80-
ngx.say("model options feature doesn't work")
81-
end
82-
return
83-
end
84-
8573
local header_auth = ngx.req.get_headers()["authorization"]
8674
local query_auth = ngx.req.get_uri_args()["apikey"]
8775
@@ -255,7 +243,6 @@ passed
255243
}]],
256244
nil,
257245
{
258-
["test-type"] = "options",
259246
["Content-Type"] = "application/json",
260247
}
261248
)
@@ -419,7 +406,6 @@ passed
419406
}]],
420407
nil,
421408
{
422-
["test-type"] = "options",
423409
["Content-Type"] = "application/json",
424410
}
425411
)
@@ -581,7 +567,6 @@ passed
581567
}]],
582568
nil,
583569
{
584-
["test-type"] = "options",
585570
["Content-Type"] = "application/json",
586571
}
587572
)
@@ -769,7 +754,6 @@ passed
769754
}]],
770755
nil,
771756
{
772-
["test-type"] = "options",
773757
["Content-Type"] = "application/json",
774758
}
775759
)
@@ -1022,7 +1006,6 @@ POST /ai
10221006
}]],
10231007
nil,
10241008
{
1025-
["test-type"] = "options",
10261009
["Content-Type"] = "application/json",
10271010
}
10281011
)
@@ -1040,7 +1023,6 @@ POST /ai
10401023
}]],
10411024
nil,
10421025
{
1043-
["test-type"] = "options",
10441026
["Content-Type"] = "application/json",
10451027
}
10461028
)

‎t/plugin/ai-proxy.t‎

Lines changed: 130 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -62,14 +62,19 @@ add_block_preprocessor(sub {
6262
if test_type == "options" then
6363
if body.foo == "bar" then
6464
ngx.status = 200
65-
ngx.say("options works")
65+
ngx.print("options works")
6666
else
6767
ngx.status = 500
6868
ngx.say("model options feature doesn't work")
6969
end
7070
return
7171
end
7272
73+
if test_type == "header_forwarding" then
74+
ngx.say(json.encode(ngx.req.get_headers()))
75+
return
76+
end
77+
7378
local header_auth = ngx.req.get_headers()["authorization"]
7479
local query_auth = ngx.req.get_uri_args()["apikey"]
7580
@@ -168,7 +173,7 @@ add_block_preprocessor(sub {
168173
169174
location /random {
170175
content_by_lua_block {
171-
ngx.say("path override works")
176+
ngx.print("path override works")
172177
}
173178
}
174179
}
@@ -448,8 +453,8 @@ unsupported content-type: application/x-www-form-urlencoded, only application/js
448453
}
449454
}
450455
--- error_code: 200
451-
--- response_body_chomp
452-
options_works
456+
--- response_body
457+
options works
453458
454459
455460
@@ -510,7 +515,7 @@ options_works
510515
511516
}
512517
}
513-
--- response_body_chomp
518+
--- response_body
514519
path override works
515520
516521
@@ -698,3 +703,123 @@ qr/.*text-embedding-ada-002*/
698703
GET /t
699704
--- response_body
700705
ok
706+
707+
708+
709+
=== TEST 18: set route with right auth header
710+
--- config
711+
location /t {
712+
content_by_lua_block {
713+
local t = require("lib.test_admin").test
714+
local code, body = t('/apisix/admin/routes/1',
715+
ngx.HTTP_PUT,
716+
[[{
717+
"uri": "/anything",
718+
"plugins": {
719+
"ai-proxy": {
720+
"provider": "openai",
721+
"auth": {
722+
"header": {
723+
"Authorization": "Bearer token"
724+
}
725+
},
726+
"options": {
727+
"model": "gpt-35-turbo-instruct",
728+
"max_tokens": 512,
729+
"temperature": 1.0
730+
},
731+
"override": {
732+
"endpoint": "http://localhost:6724"
733+
},
734+
"ssl_verify": false
735+
}
736+
}
737+
}]]
738+
)
739+
740+
if code >= 300 then
741+
ngx.status = code
742+
end
743+
ngx.say(body)
744+
}
745+
}
746+
--- response_body
747+
passed
748+
749+
750+
751+
=== TEST 19: send request
752+
--- request
753+
POST /anything
754+
{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] }
755+
--- more_headers
756+
test-type: header_forwarding
757+
--- error_code: 200
758+
--- response_body eval
759+
qr/"test-type":"header_forwarding"/
760+
761+
762+
763+
=== TEST 20: set route with right auth header
764+
--- config
765+
location /t {
766+
content_by_lua_block {
767+
local t = require("lib.test_admin").test
768+
local code, body = t('/apisix/admin/routes/1',
769+
ngx.HTTP_PUT,
770+
[[{
771+
"uri": "/anything",
772+
"plugins": {
773+
"ai-proxy": {
774+
"provider": "openai",
775+
"auth": {
776+
"header": {
777+
"Authorization": "Bearer token"
778+
}
779+
},
780+
"options": {
781+
"model": "gpt-35-turbo-instruct",
782+
"max_tokens": 512,
783+
"temperature": 1.0
784+
},
785+
"override": {
786+
"endpoint": "http://localhost:6724"
787+
},
788+
"ssl_verify": false
789+
},
790+
"request-id": {
791+
}
792+
}
793+
}]]
794+
)
795+
796+
if code >= 300 then
797+
ngx.status = code
798+
end
799+
ngx.say(body)
800+
}
801+
}
802+
--- response_body
803+
passed
804+
805+
806+
807+
=== TEST 21: send request
808+
--- request
809+
POST /anything
810+
{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] }
811+
--- more_headers
812+
test-type: header_forwarding
813+
--- error_code: 200
814+
--- response_body eval
815+
qr/"x-request-id":"[\d\w-]+"/
816+
817+
818+
819+
=== TEST 22: send request with Authorization header
820+
--- request
821+
POST /anything
822+
{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] }
823+
--- more_headers
824+
Authorization: Bearer wrong token
825+
--- error_code: 200

‎t/plugin/ai-rate-limiting.t‎

Lines changed: 0 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -78,18 +78,6 @@ add_block_preprocessor(sub {
7878
local body, err = ngx.req.get_body_data()
7979
body, err = json.decode(body)
8080
81-
local test_type = ngx.req.get_headers()["test-type"]
82-
if test_type == "options" then
83-
if body.foo == "bar" then
84-
ngx.status = 200
85-
ngx.say("options works")
86-
else
87-
ngx.status = 500
88-
ngx.say("model options feature doesn't work")
89-
end
90-
return
91-
end
92-
9381
local header_auth = ngx.req.get_headers()["authorization"]
9482
local query_auth = ngx.req.get_uri_args()["apikey"]
9583
@@ -656,7 +644,6 @@ passed
656644
}]],
657645
nil,
658646
{
659-
["test-type"] = "options",
660647
["Content-Type"] = "application/json",
661648
}
662649
)
@@ -675,7 +662,6 @@ passed
675662
}]],
676663
nil,
677664
{
678-
["test-type"] = "options",
679665
["Content-Type"] = "application/json",
680666
}
681667
)
@@ -694,7 +680,6 @@ passed
694680
}]],
695681
nil,
696682
{
697-
["test-type"] = "options",
698683
["Content-Type"] = "application/json",
699684
}
700685
)
@@ -793,7 +778,6 @@ passed
793778
}]],
794779
nil,
795780
{
796-
["test-type"] = "options",
797781
["Content-Type"] = "application/json",
798782
}
799783
)

0 commit comments

Comments
 (0)