Skip to content

Commit 99a0792

Browse files
shreemaan-abhishekshreemaan-abhishek
authored
chore: rectify business logic/code in ai-proxy (apache#12055)
1 parent 428f43a commit 99a0792

File tree

13 files changed

+28
-149
lines changed

13 files changed

+28
-149
lines changed

‎apisix/init.lua‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -453,6 +453,12 @@ end
453453

454454

455455
function _M.handle_upstream(api_ctx, route, enable_websocket)
456+
-- some plugins(ai-proxy...) request upstream by http client directly
457+
if api_ctx.bypass_nginx_upstream then
458+
common_phase("before_proxy")
459+
return
460+
end
461+
456462
local up_id = route.value.upstream_id
457463

458464
-- used for the traffic-split plugin

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

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -141,22 +141,25 @@ end
141141
function _M.read_response(ctx, res)
142142
local body_reader = res.body_reader
143143
if not body_reader then
144-
core.log.error("AI service sent no response body")
144+
core.log.warn("AI service sent no response body")
145145
return 500
146146
end
147147

148148
local content_type = res.headers["Content-Type"]
149149
core.response.set_header("Content-Type", content_type)
150150

151-
if core.string.find(content_type, "text/event-stream") then
151+
if content_type and core.string.find(content_type, "text/event-stream") then
152152
while true do
153153
local chunk, err = body_reader() -- will read chunk by chunk
154154
if err then
155-
core.log.error("failed to read response chunk: ", err)
156-
break
155+
core.log.warn("failed to read response chunk: ", err)
156+
if core.string.find(err, "timeout") then
157+
return 504
158+
end
159+
return 500
157160
end
158161
if not chunk then
159-
break
162+
return
160163
end
161164

162165
ngx_print(chunk)
@@ -192,6 +195,8 @@ function _M.read_response(ctx, res)
192195

193196
-- usage field is null for non-last events, null is parsed as userdata type
194197
if data and data.usage and type(data.usage) ~= "userdata" then
198+
core.log.info("got token usage from ai service: ",
199+
core.json.delay_encode(data.usage))
195200
ctx.ai_token_usage = {
196201
prompt_tokens = data.usage.prompt_tokens or 0,
197202
completion_tokens = data.usage.completion_tokens or 0,
@@ -202,19 +207,22 @@ function _M.read_response(ctx, res)
202207

203208
::CONTINUE::
204209
end
205-
return
206210
end
207211

208212
local raw_res_body, err = res:read_body()
209213
if not raw_res_body then
210-
core.log.error("failed to read response body: ", err)
214+
core.log.warn("failed to read response body: ", err)
215+
if core.string.find(err, "timeout") then
216+
return 504
217+
end
211218
return 500
212219
end
213220
local res_body, err = core.json.decode(raw_res_body)
214221
if err then
215222
core.log.warn("invalid response body from ai service: ", raw_res_body, " err: ", err,
216223
", it will cause token usage not available")
217224
else
225+
core.log.info("got token usage from ai service: ", core.json.delay_encode(res_body.usage))
218226
ctx.ai_token_usage = {
219227
prompt_tokens = res_body.usage and res_body.usage.prompt_tokens or 0,
220228
completion_tokens = res_body.usage and res_body.usage.completion_tokens or 0,

‎apisix/plugins/ai-proxy-multi.lua‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ function _M.access(conf, ctx)
217217
end
218218
ctx.picked_ai_instance_name = name
219219
ctx.picked_ai_instance = ai_instance
220+
ctx.bypass_nginx_upstream = true
220221
end
221222

222223

‎apisix/plugins/ai-proxy.lua‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ end
4747
function _M.access(conf, ctx)
4848
ctx.picked_ai_instance_name = "ai-proxy"
4949
ctx.picked_ai_instance = conf
50+
ctx.bypass_nginx_upstream = true
5051
end
5152

5253

‎apisix/plugins/ai-proxy/base.lua‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,10 @@ function _M.before_proxy(conf, ctx)
4646

4747
local res, err = ai_driver:request(conf, request_body, extra_opts)
4848
if not res then
49-
core.log.error("failed to send request to AI service: ", err)
49+
core.log.warn("failed to send request to AI service: ", err)
50+
if core.string.find(err, "timeout") then
51+
return 504
52+
end
5053
return internal_server_error
5154
end
5255

‎apisix/plugins/ai-proxy/schema.lua‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -120,7 +120,7 @@ local ai_instance_schema = {
120120
},
121121
},
122122
},
123-
required = {"name", "provider", "auth"}
123+
required = {"name", "provider", "auth", "weight"}
124124
},
125125
}
126126

@@ -139,7 +139,6 @@ _M.ai_proxy_schema = {
139139
timeout = {
140140
type = "integer",
141141
minimum = 1,
142-
maximum = 60000,
143142
default = 30000,
144143
description = "timeout in milliseconds",
145144
},
@@ -196,7 +195,6 @@ _M.ai_proxy_multi_schema = {
196195
timeout = {
197196
type = "integer",
198197
minimum = 1,
199-
maximum = 60000,
200198
default = 30000,
201199
description = "timeout in milliseconds",
202200
},

‎t/plugin/ai-proxy-multi.balancer.t‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -198,12 +198,6 @@ __DATA__
198198
],
199199
"ssl_verify": false
200200
}
201-
},
202-
"upstream": {
203-
"type": "roundrobin",
204-
"nodes": {
205-
"canbeanything.com": 1
206-
}
207201
}
208202
}]]
209203
)
@@ -308,12 +302,6 @@ deepseek.deepseek.openai.openai.openai.openai.openai.openai.openai.openai
308302
],
309303
"ssl_verify": false
310304
}
311-
},
312-
"upstream": {
313-
"type": "roundrobin",
314-
"nodes": {
315-
"canbeanything.com": 1
316-
}
317305
}
318306
}]]
319307
)

‎t/plugin/ai-proxy-multi.openai-compatible.t‎

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -160,12 +160,6 @@ __DATA__
160160
],
161161
"ssl_verify": false
162162
}
163-
},
164-
"upstream": {
165-
"type": "roundrobin",
166-
"nodes": {
167-
"canbeanything.com": 1
168-
}
169163
}
170164
}]]
171165
)
@@ -227,12 +221,6 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/
227221
],
228222
"ssl_verify": false
229223
}
230-
},
231-
"upstream": {
232-
"type": "roundrobin",
233-
"nodes": {
234-
"canbeanything.com": 1
235-
}
236224
}
237225
}]]
238226
)

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

Lines changed: 0 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -230,12 +230,6 @@ qr/.*property "provider" validation failed: matches none of the enum values*/
230230
],
231231
"ssl_verify": false
232232
}
233-
},
234-
"upstream": {
235-
"type": "roundrobin",
236-
"nodes": {
237-
"canbeanything.com": 1
238-
}
239233
}
240234
}]]
241235
)
@@ -294,12 +288,6 @@ Unauthorized
294288
],
295289
"ssl_verify": false
296290
}
297-
},
298-
"upstream": {
299-
"type": "roundrobin",
300-
"nodes": {
301-
"canbeanything.com": 1
302-
}
303291
}
304292
}]]
305293
)
@@ -417,12 +405,6 @@ request format doesn't match schema: property "messages" is required
417405
],
418406
"ssl_verify": false
419407
}
420-
},
421-
"upstream": {
422-
"type": "roundrobin",
423-
"nodes": {
424-
"canbeanything.com": 1
425-
}
426408
}
427409
}]]
428410
)
@@ -492,12 +474,6 @@ options_works
492474
],
493475
"ssl_verify": false
494476
}
495-
},
496-
"upstream": {
497-
"type": "roundrobin",
498-
"nodes": {
499-
"canbeanything.com": 1
500-
}
501477
}
502478
}]]
503479
)
@@ -567,12 +543,6 @@ path override works
567543
],
568544
"ssl_verify": false
569545
}
570-
},
571-
"upstream": {
572-
"type": "roundrobin",
573-
"nodes": {
574-
"canbeanything.com": 1
575-
}
576546
}
577547
}]]
578548
)

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

Lines changed: 0 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -145,12 +145,6 @@ __DATA__
145145
],
146146
"ssl_verify": false
147147
}
148-
},
149-
"upstream": {
150-
"type": "roundrobin",
151-
"nodes": {
152-
"canbeanything.com": 1
153-
}
154148
}
155149
}]]
156150
)
@@ -209,12 +203,6 @@ Unauthorized
209203
],
210204
"ssl_verify": false
211205
}
212-
},
213-
"upstream": {
214-
"type": "roundrobin",
215-
"nodes": {
216-
"canbeanything.com": 1
217-
}
218206
}
219207
}]]
220208
)
@@ -333,12 +321,6 @@ POST /anything
333321
],
334322
"ssl_verify": false
335323
}
336-
},
337-
"upstream": {
338-
"type": "roundrobin",
339-
"nodes": {
340-
"canbeanything.com": 1
341-
}
342324
}
343325
}]]
344326
)

0 commit comments

Comments
 (0)
⚡