This is an automated email from the ASF dual-hosted git repository.
shreemaanabhishek pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/apisix.git
The following commit(s) were added to refs/heads/master by this push:
new 99a079258 chore: rectify business logic/code in ai-proxy (#12055)
99a079258 is described below
commit 99a079258d39879bb7fe371d8265d9a4056938f1
Author: Shreemaan Abhishek <[email protected]>
AuthorDate: Mon Mar 17 19:29:37 2025 +0545
chore: rectify business logic/code in ai-proxy (#12055)
---
apisix/init.lua | 6 ++++++
apisix/plugins/ai-drivers/openai-base.lua | 22 ++++++++++++++-------
apisix/plugins/ai-proxy-multi.lua | 1 +
apisix/plugins/ai-proxy.lua | 1 +
apisix/plugins/ai-proxy/base.lua | 5 ++++-
apisix/plugins/ai-proxy/schema.lua | 4 +---
t/plugin/ai-proxy-multi.balancer.t | 12 ------------
t/plugin/ai-proxy-multi.openai-compatible.t | 12 ------------
t/plugin/ai-proxy-multi.t | 30 -----------------------------
t/plugin/ai-proxy-multi2.t | 18 -----------------
t/plugin/ai-proxy.openai-compatible.t | 18 -----------------
t/plugin/ai-proxy.t | 30 -----------------------------
t/plugin/ai-proxy2.t | 18 -----------------
13 files changed, 28 insertions(+), 149 deletions(-)
diff --git a/apisix/init.lua b/apisix/init.lua
index 4b2ad17e1..426f7fd71 100644
--- a/apisix/init.lua
+++ b/apisix/init.lua
@@ -453,6 +453,12 @@ end
function _M.handle_upstream(api_ctx, route, enable_websocket)
+ -- some plugins(ai-proxy...) request upstream by http client directly
+ if api_ctx.bypass_nginx_upstream then
+ common_phase("before_proxy")
+ return
+ end
+
local up_id = route.value.upstream_id
-- used for the traffic-split plugin
diff --git a/apisix/plugins/ai-drivers/openai-base.lua
b/apisix/plugins/ai-drivers/openai-base.lua
index 4f0b38afe..a34261202 100644
--- a/apisix/plugins/ai-drivers/openai-base.lua
+++ b/apisix/plugins/ai-drivers/openai-base.lua
@@ -141,22 +141,25 @@ end
function _M.read_response(ctx, res)
local body_reader = res.body_reader
if not body_reader then
- core.log.error("AI service sent no response body")
+ core.log.warn("AI service sent no response body")
return 500
end
local content_type = res.headers["Content-Type"]
core.response.set_header("Content-Type", content_type)
- if core.string.find(content_type, "text/event-stream") then
+ if content_type and core.string.find(content_type, "text/event-stream")
then
while true do
local chunk, err = body_reader() -- will read chunk by chunk
if err then
- core.log.error("failed to read response chunk: ", err)
- break
+ core.log.warn("failed to read response chunk: ", err)
+ if core.string.find(err, "timeout") then
+ return 504
+ end
+ return 500
end
if not chunk then
- break
+ return
end
ngx_print(chunk)
@@ -192,6 +195,8 @@ function _M.read_response(ctx, res)
-- usage field is null for non-last events, null is parsed as
userdata type
if data and data.usage and type(data.usage) ~= "userdata" then
+ core.log.info("got token usage from ai service: ",
+ core.json.delay_encode(data.usage))
ctx.ai_token_usage = {
prompt_tokens = data.usage.prompt_tokens or 0,
completion_tokens = data.usage.completion_tokens or 0,
@@ -202,12 +207,14 @@ function _M.read_response(ctx, res)
::CONTINUE::
end
- return
end
local raw_res_body, err = res:read_body()
if not raw_res_body then
- core.log.error("failed to read response body: ", err)
+ core.log.warn("failed to read response body: ", err)
+ if core.string.find(err, "timeout") then
+ return 504
+ end
return 500
end
local res_body, err = core.json.decode(raw_res_body)
@@ -215,6 +222,7 @@ function _M.read_response(ctx, res)
core.log.warn("invalid response body from ai service: ", raw_res_body,
" err: ", err,
", it will cause token usage not available")
else
+ core.log.info("got token usage from ai service: ",
core.json.delay_encode(res_body.usage))
ctx.ai_token_usage = {
prompt_tokens = res_body.usage and res_body.usage.prompt_tokens or
0,
completion_tokens = res_body.usage and
res_body.usage.completion_tokens or 0,
diff --git a/apisix/plugins/ai-proxy-multi.lua
b/apisix/plugins/ai-proxy-multi.lua
index 6065dd3ea..7ac8bb206 100644
--- a/apisix/plugins/ai-proxy-multi.lua
+++ b/apisix/plugins/ai-proxy-multi.lua
@@ -217,6 +217,7 @@ function _M.access(conf, ctx)
end
ctx.picked_ai_instance_name = name
ctx.picked_ai_instance = ai_instance
+ ctx.bypass_nginx_upstream = true
end
diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua
index 2301a65e6..fa7f5f24b 100644
--- a/apisix/plugins/ai-proxy.lua
+++ b/apisix/plugins/ai-proxy.lua
@@ -47,6 +47,7 @@ end
function _M.access(conf, ctx)
ctx.picked_ai_instance_name = "ai-proxy"
ctx.picked_ai_instance = conf
+ ctx.bypass_nginx_upstream = true
end
diff --git a/apisix/plugins/ai-proxy/base.lua b/apisix/plugins/ai-proxy/base.lua
index d8f1a8944..73c683c0a 100644
--- a/apisix/plugins/ai-proxy/base.lua
+++ b/apisix/plugins/ai-proxy/base.lua
@@ -46,7 +46,10 @@ function _M.before_proxy(conf, ctx)
local res, err = ai_driver:request(conf, request_body, extra_opts)
if not res then
- core.log.error("failed to send request to AI service: ", err)
+ core.log.warn("failed to send request to AI service: ", err)
+ if core.string.find(err, "timeout") then
+ return 504
+ end
return internal_server_error
end
diff --git a/apisix/plugins/ai-proxy/schema.lua
b/apisix/plugins/ai-proxy/schema.lua
index a2c41b97b..1b9d07b1c 100644
--- a/apisix/plugins/ai-proxy/schema.lua
+++ b/apisix/plugins/ai-proxy/schema.lua
@@ -120,7 +120,7 @@ local ai_instance_schema = {
},
},
},
- required = {"name", "provider", "auth"}
+ required = {"name", "provider", "auth", "weight"}
},
}
@@ -139,7 +139,6 @@ _M.ai_proxy_schema = {
timeout = {
type = "integer",
minimum = 1,
- maximum = 60000,
default = 30000,
description = "timeout in milliseconds",
},
@@ -196,7 +195,6 @@ _M.ai_proxy_multi_schema = {
timeout = {
type = "integer",
minimum = 1,
- maximum = 60000,
default = 30000,
description = "timeout in milliseconds",
},
diff --git a/t/plugin/ai-proxy-multi.balancer.t
b/t/plugin/ai-proxy-multi.balancer.t
index 09076e4a8..48ab5d8ca 100644
--- a/t/plugin/ai-proxy-multi.balancer.t
+++ b/t/plugin/ai-proxy-multi.balancer.t
@@ -198,12 +198,6 @@ __DATA__
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -308,12 +302,6 @@
deepseek.deepseek.openai.openai.openai.openai.openai.openai.openai.openai
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
diff --git a/t/plugin/ai-proxy-multi.openai-compatible.t
b/t/plugin/ai-proxy-multi.openai-compatible.t
index 923c12d37..d5be5d204 100644
--- a/t/plugin/ai-proxy-multi.openai-compatible.t
+++ b/t/plugin/ai-proxy-multi.openai-compatible.t
@@ -160,12 +160,6 @@ __DATA__
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -227,12 +221,6 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
diff --git a/t/plugin/ai-proxy-multi.t b/t/plugin/ai-proxy-multi.t
index 0da04c9de..a70392e60 100644
--- a/t/plugin/ai-proxy-multi.t
+++ b/t/plugin/ai-proxy-multi.t
@@ -230,12 +230,6 @@ qr/.*property "provider" validation failed: matches none
of the enum values*/
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -294,12 +288,6 @@ Unauthorized
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -417,12 +405,6 @@ request format doesn't match schema: property "messages"
is required
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -492,12 +474,6 @@ options_works
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -567,12 +543,6 @@ path override works
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
diff --git a/t/plugin/ai-proxy-multi2.t b/t/plugin/ai-proxy-multi2.t
index c54e7a67e..536c98c7d 100644
--- a/t/plugin/ai-proxy-multi2.t
+++ b/t/plugin/ai-proxy-multi2.t
@@ -145,12 +145,6 @@ __DATA__
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -209,12 +203,6 @@ Unauthorized
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -333,12 +321,6 @@ POST /anything
],
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
diff --git a/t/plugin/ai-proxy.openai-compatible.t
b/t/plugin/ai-proxy.openai-compatible.t
index 84ae175da..a5147648c 100644
--- a/t/plugin/ai-proxy.openai-compatible.t
+++ b/t/plugin/ai-proxy.openai-compatible.t
@@ -148,12 +148,6 @@ __DATA__
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -208,12 +202,6 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -277,12 +265,6 @@ path override works
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t
index 08220fc3c..c5696a2fb 100644
--- a/t/plugin/ai-proxy.t
+++ b/t/plugin/ai-proxy.t
@@ -206,12 +206,6 @@ qr/.*property "provider" validation failed: matches none
of the enum values.*/
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -264,12 +258,6 @@ Unauthorized
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -381,12 +369,6 @@ request format doesn't match schema: property "messages"
is required
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -450,12 +432,6 @@ options_works
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -519,12 +495,6 @@ path override works
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
diff --git a/t/plugin/ai-proxy2.t b/t/plugin/ai-proxy2.t
index 942f449cd..43cdd3008 100644
--- a/t/plugin/ai-proxy2.t
+++ b/t/plugin/ai-proxy2.t
@@ -133,12 +133,6 @@ __DATA__
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -191,12 +185,6 @@ Unauthorized
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
@@ -301,12 +289,6 @@ POST /anything
},
"ssl_verify": false
}
- },
- "upstream": {
- "type": "roundrobin",
- "nodes": {
- "canbeanything.com": 1
- }
}
}]]
)
