-
Notifications
You must be signed in to change notification settings - Fork 2.5k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
feat: ai-proxy plugin #11499
feat: ai-proxy plugin #11499
Changes from 79 commits
b98a48f
8188ae4
7b83b3a
97cafa5
35b1787
e18caef
28f06ae
82f9692
0577e8e
e5f00f7
c307b04
ef4cf84
4bf6bd2
42adfd1
0af00ae
aff56a0
f25f21a
d2d253e
58ca8a7
f146f20
2317aa8
3ac0fe5
6e31cfe
e302360
83f2197
bcc21cb
10a07c1
7220c08
a4afb30
6248005
e88683c
7d9c075
9823570
94d00f4
b24e439
5ca70f3
284ad76
6baa7d1
bdab563
2d0a7a1
530448f
ed11fa4
cba307a
3febd29
e566a37
4c13cef
28c9c4d
780561d
4ffdd85
0521f89
bd8309b
6f5d158
d4f3a5a
1bfeac9
25975c0
7c77ed6
fa46abe
45e4f98
99af867
11aef59
4eb0ff4
92ebd9d
c0dc59e
99cf3a4
243b5f5
217b5af
6661a0e
1c00e2c
37cdd7b
1062cc2
7b23623
4278fd5
d915292
8c3bcb3
bcca3f2
929cbb1
e149170
d346d38
3429e03
7c08290
977cf68
073b3f8
5831108
ab4c37e
cdd37db
e0e3e15
f6768e4
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -726,6 +726,7 @@ function _M.http_access_phase() | |
end | ||
|
||
|
||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. If there are no substantive changes, please do not change the code style. |
||
function _M.dubbo_access_phase() | ||
ngx.ctx = fetch_ctx() | ||
end | ||
|
Original file line number | Diff line number | Diff line change | ||||||
---|---|---|---|---|---|---|---|---|
@@ -0,0 +1,136 @@ | ||||||||
-- | ||||||||
-- Licensed to the Apache Software Foundation (ASF) under one or more | ||||||||
-- contributor license agreements. See the NOTICE file distributed with | ||||||||
-- this work for additional information regarding copyright ownership. | ||||||||
-- The ASF licenses this file to You under the Apache License, Version 2.0 | ||||||||
-- (the "License"); you may not use this file except in compliance with | ||||||||
-- the License. You may obtain a copy of the License at | ||||||||
-- | ||||||||
-- http://www.apache.org/licenses/LICENSE-2.0 | ||||||||
-- | ||||||||
-- Unless required by applicable law or agreed to in writing, software | ||||||||
-- distributed under the License is distributed on an "AS IS" BASIS, | ||||||||
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||||||||
-- See the License for the specific language governing permissions and | ||||||||
-- limitations under the License. | ||||||||
-- | ||||||||
local core = require("apisix.core") | ||||||||
local schema = require("apisix.plugins.ai-proxy.schema") | ||||||||
local require = require | ||||||||
local pcall = pcall | ||||||||
local ngx_req = ngx.req | ||||||||
local ngx_print = ngx.print | ||||||||
local ngx_flush = ngx.flush | ||||||||
|
||||||||
local plugin_name = "ai-proxy" | ||||||||
local _M = { | ||||||||
version = 0.5, | ||||||||
priority = 1004, | ||||||||
name = plugin_name, | ||||||||
schema = schema, | ||||||||
} | ||||||||
|
||||||||
|
||||||||
function _M.check_schema(conf) | ||||||||
local ai_driver = pcall(require, "apisix.plugins.ai-proxy.drivers." .. conf.model.provider) | ||||||||
if not ai_driver then | ||||||||
return false, "provider: " .. conf.model.provider .. " is not supported." | ||||||||
end | ||||||||
return core.schema.check(schema.plugin_schema, conf) | ||||||||
end | ||||||||
|
||||||||
|
||||||||
local CONTENT_TYPE_JSON = "application/json" | ||||||||
bzp2010 marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
|
||||||||
|
||||||||
local function keepalive_or_close(conf, httpc) | ||||||||
if conf.set_keepalive then | ||||||||
httpc:set_keepalive(10000, 100) | ||||||||
bzp2010 marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
return | ||||||||
end | ||||||||
httpc:close() | ||||||||
end | ||||||||
|
||||||||
|
||||||||
function _M.access(conf, ctx) | ||||||||
local ct = core.request.header(ctx, "Content-Type") or CONTENT_TYPE_JSON | ||||||||
if not core.string.has_prefix(ct, CONTENT_TYPE_JSON) then | ||||||||
return 400, "unsupported content-type: " .. ct | ||||||||
end | ||||||||
|
||||||||
local request_table, err = core.request.get_request_body_table() | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. new naming style? request_table, count_num, flag_boolean? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't understand what you mean. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Although adding a suffix makes it more readable, it is not the style used in the project. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. oh I got it now 😂 . |
||||||||
if not request_table then | ||||||||
return 400, err | ||||||||
end | ||||||||
|
||||||||
local ok, err = core.schema.check(schema.chat_request_schema, request_table) | ||||||||
if not ok then | ||||||||
return 400, "request format doesn't match schema: " .. err | ||||||||
end | ||||||||
|
||||||||
if conf.model.name then | ||||||||
request_table.model = conf.model.name | ||||||||
end | ||||||||
|
||||||||
if core.table.try_read_attr(conf, "model", "options", "stream") then | ||||||||
request_table.stream = true | ||||||||
end | ||||||||
|
||||||||
local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) | ||||||||
shreemaan-abhishek marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||
local res, err, httpc = ai_driver.request(conf, request_table, ctx) | ||||||||
if not res then | ||||||||
core.log.error("failed to send request to AI service: ", err) | ||||||||
return 500 | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. use the predefined Ngx constants?
|
||||||||
end | ||||||||
|
||||||||
local body_reader = res.body_reader | ||||||||
if not body_reader then | ||||||||
core.log.error("LLM sent no response body") | ||||||||
return 500 | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ditto |
||||||||
end | ||||||||
|
||||||||
if conf.passthrough then | ||||||||
ngx_req.init_body() | ||||||||
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 | ||||||||
end | ||||||||
if not chunk then | ||||||||
break | ||||||||
end | ||||||||
ngx_req.append_body(chunk) | ||||||||
end | ||||||||
ngx_req.finish_body() | ||||||||
keepalive_or_close(conf, httpc) | ||||||||
return | ||||||||
end | ||||||||
|
||||||||
if core.table.try_read_attr(conf, "model", "options", "stream") then | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||||
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 | ||||||||
end | ||||||||
if not chunk then | ||||||||
break | ||||||||
end | ||||||||
ngx_print(chunk) | ||||||||
ngx_flush(true) | ||||||||
end | ||||||||
keepalive_or_close(conf, httpc) | ||||||||
return | ||||||||
else | ||||||||
local res_body, err = res:read_body() | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. also using body_reader, ngx_print * N + ngx_flush * 1? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I don't understand what you mean. There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The purpose is to save the intermediate buffering process and directly But, $ restydoc -s ngx.print
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. got it, I think we don't need There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Hi @shreemaan-abhishek Is it resolved? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. yes |
||||||||
if not res_body then | ||||||||
core.log.error("failed to read response body: ", err) | ||||||||
return 500 | ||||||||
end | ||||||||
keepalive_or_close(conf, httpc) | ||||||||
return res.status, res_body | ||||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
also There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. yes, thank you. |
||||||||
end | ||||||||
end | ||||||||
|
||||||||
return _M | ||||||||
bzp2010 marked this conversation as resolved.
Show resolved
Hide resolved
|
Original file line number | Diff line number | Diff line change | ||||
---|---|---|---|---|---|---|
@@ -0,0 +1,92 @@ | ||||||
-- | ||||||
-- Licensed to the Apache Software Foundation (ASF) under one or more | ||||||
-- contributor license agreements. See the NOTICE file distributed with | ||||||
-- this work for additional information regarding copyright ownership. | ||||||
-- The ASF licenses this file to You under the Apache License, Version 2.0 | ||||||
-- (the "License"); you may not use this file except in compliance with | ||||||
-- the License. You may obtain a copy of the License at | ||||||
-- | ||||||
-- http://www.apache.org/licenses/LICENSE-2.0 | ||||||
-- | ||||||
-- Unless required by applicable law or agreed to in writing, software | ||||||
-- distributed under the License is distributed on an "AS IS" BASIS, | ||||||
-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||||||
-- See the License for the specific language governing permissions and | ||||||
-- limitations under the License. | ||||||
-- | ||||||
local _M = {} | ||||||
|
||||||
local core = require("apisix.core") | ||||||
local http = require("resty.http") | ||||||
local url = require("socket.url") | ||||||
|
||||||
local pairs = pairs | ||||||
local type = type | ||||||
|
||||||
-- globals | ||||||
local DEFAULT_HOST = "api.openai.com" | ||||||
local DEFAULT_PORT = 443 | ||||||
|
||||||
|
||||||
function _M.request(conf, request_table, ctx) | ||||||
local httpc, err = http.new() | ||||||
if not httpc then | ||||||
return nil, "failed to create http client to send request to LLM server: " .. err | ||||||
end | ||||||
httpc:set_timeout(conf.timeout) | ||||||
|
||||||
local endpoint = core.table.try_read_attr(conf, "override", "endpoint") | ||||||
local parsed_url | ||||||
if endpoint then | ||||||
parsed_url = url.parse(endpoint) | ||||||
end | ||||||
|
||||||
local ok, err = httpc:connect({ | ||||||
scheme = parsed_url.scheme or "https", | ||||||
host = parsed_url.host or DEFAULT_HOST, | ||||||
port = parsed_url.port or DEFAULT_PORT, | ||||||
ssl_verify = conf.ssl_verify, | ||||||
ssl_server_name = parsed_url.host or DEFAULT_HOST, | ||||||
pool_size = conf.keepalive and conf.keepalive_pool, | ||||||
}) | ||||||
|
||||||
if not ok then | ||||||
return nil, "failed to connect to LLM server: " .. err | ||||||
end | ||||||
|
||||||
local query_params = "" | ||||||
if conf.auth.query and type(conf.auth.query) == "table" then | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
query_params = core.string.encode_args(conf.auth.query) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. FYI, from lua-resty-http,
|
||||||
if query_params and query_params ~= "" then | ||||||
query_params = "?" .. query_params | ||||||
end | ||||||
end | ||||||
|
||||||
local path = (parsed_url.path or "/v1/chat/completions") .. query_params | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
|
||||||
local headers = (conf.auth.header or {}) | ||||||
headers["Content-Type"] = "application/json" | ||||||
local params = { | ||||||
method = "POST", | ||||||
headers = headers, | ||||||
keepalive = conf.keepalive, | ||||||
ssl_verify = conf.ssl_verify, | ||||||
path = path, | ||||||
} | ||||||
|
||||||
if conf.model.options then | ||||||
for opt, val in pairs(conf.model.options) do | ||||||
request_table[opt] = val | ||||||
end | ||||||
end | ||||||
params.body = core.json.encode(request_table) | ||||||
|
||||||
local res, err = httpc:request(params) | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Replace with httpc:request_uri to avoid invoking close or keepalive elsewhere? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
You need to support streaming chunk outside, so |
||||||
if not res then | ||||||
return 500, "failed to send request to LLM server: " .. err | ||||||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||
end | ||||||
|
||||||
return res, nil, httpc | ||||||
end | ||||||
|
||||||
return _M |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Decode body according content-type or rename the method, e.g. get_json_request_body_table?