From b98a48fc55513dced7ed0b0982a26c184abe9cf0 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:02:46 +0545 Subject: [PATCH 01/85] feat: ai proxy plugin --- apisix/balancer.lua | 2 + apisix/cli/config.lua | 1 + apisix/cli/ngx_tpl.lua | 90 +++++++++++ apisix/constants.lua | 2 + apisix/init.lua | 64 +++++++- apisix/plugins/ai-proxy.lua | 82 ++++++++++ apisix/plugins/ai-proxy/drivers/openai.lua | 52 +++++++ apisix/plugins/ai-proxy/schema.lua | 166 +++++++++++++++++++++ 8 files changed, 457 insertions(+), 2 deletions(-) create mode 100644 apisix/plugins/ai-proxy.lua create mode 100644 apisix/plugins/ai-proxy/drivers/openai.lua create mode 100644 apisix/plugins/ai-proxy/schema.lua diff --git a/apisix/balancer.lua b/apisix/balancer.lua index 0fe2e6539922..76cba643a1e8 100644 --- a/apisix/balancer.lua +++ b/apisix/balancer.lua @@ -285,6 +285,8 @@ do local default_keepalive_pool function set_current_peer(server, ctx) + core.log.warn("dibag peer: ", core.json.encode(server)) + core.log.warn(debug.traceback("dibag")) local up_conf = ctx.upstream_conf local keepalive_pool = up_conf.keepalive_pool diff --git a/apisix/cli/config.lua b/apisix/cli/config.lua index 94843621a74b..2e6cf7eed001 100644 --- a/apisix/cli/config.lua +++ b/apisix/cli/config.lua @@ -216,6 +216,7 @@ local _M = { "proxy-mirror", "proxy-rewrite", "workflow", + "ai-proxy", "api-breaker", "limit-conn", "limit-count", diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index 4b7ff4102bc1..bc3a24a470e6 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -809,6 +809,96 @@ http { } } + location /subrequest { + internal; + + access_by_lua_block {;} + header_filter_by_lua_block {;} + body_filter_by_lua_block {;} + log_by_lua_block {;} + + proxy_http_version 1.1; + proxy_set_header Host $upstream_host; + proxy_set_header Upgrade $upstream_upgrade; + proxy_set_header Connection $upstream_connection; + proxy_set_header X-Real-IP $remote_addr; + proxy_pass_header Date; + + + set $var_x_forwarded_proto $scheme; + set $var_x_forwarded_host $host; + set $var_x_forwarded_port $server_port; + + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $var_x_forwarded_proto; + proxy_set_header X-Forwarded-Host $var_x_forwarded_host; + proxy_set_header X-Forwarded-Port $var_x_forwarded_port; + + proxy_pass_header Server; + proxy_pass_header Date; + proxy_ssl_name $upstream_host; + proxy_ssl_server_name on; + proxy_pass $upstream_scheme://apisix_backend$upstream_uri; + } + + location @disable_proxy_buffering { + # http server location configuration snippet starts + {% if http_server_location_configuration_snippet then %} + {* http_server_location_configuration_snippet *} + {% end %} + # http server location configuration snippet ends + + proxy_http_version 1.1; + proxy_set_header Host $upstream_host; + proxy_set_header Upgrade $upstream_upgrade; + proxy_set_header Connection $upstream_connection; + proxy_set_header X-Real-IP $remote_addr; + proxy_pass_header Date; + + ### the following x-forwarded-* headers is to send to upstream server + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $var_x_forwarded_proto; + proxy_set_header X-Forwarded-Host $var_x_forwarded_host; + proxy_set_header X-Forwarded-Port $var_x_forwarded_port; + + {% if enabled_plugins["proxy-cache"] then %} + ### the following configuration is to cache response content from upstream server + proxy_cache $upstream_cache_zone; + proxy_cache_valid any {% if proxy_cache.cache_ttl then %} {* proxy_cache.cache_ttl *} {% else %} 10s {% end %}; + proxy_cache_min_uses 1; + proxy_cache_methods GET HEAD POST; + proxy_cache_lock_timeout 5s; + proxy_cache_use_stale off; + proxy_cache_key $upstream_cache_key; + proxy_no_cache $upstream_no_cache; + proxy_cache_bypass $upstream_cache_bypass; + + {% end %} + + proxy_pass $upstream_scheme://apisix_backend$upstream_uri; + + {% if enabled_plugins["proxy-mirror"] then %} + mirror /proxy_mirror; + {% end %} + + header_filter_by_lua_block { + apisix.http_header_filter_phase() + } + + body_filter_by_lua_block { + apisix.http_body_filter_phase() + } + + log_by_lua_block { + apisix.http_log_phase() + } + + proxy_buffering off; + access_by_lua_block { + apisix.disable_proxy_buffering_access_phase() + } + } + location @grpc_pass { access_by_lua_block { diff --git a/apisix/constants.lua b/apisix/constants.lua index 0b3ec160b53d..2995618cd433 100644 --- a/apisix/constants.lua +++ b/apisix/constants.lua @@ -43,4 +43,6 @@ return { ["/stream_routes"] = true, ["/plugin_metadata"] = true, }, + CHAT = "llm/chat", + COMPLETION = "llm/completion", } diff --git a/apisix/init.lua b/apisix/init.lua index 103a8c1d7584..ba49e6b9c993 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -59,6 +59,7 @@ local tonumber = tonumber local type = type local pairs = pairs local ngx_re_match = ngx.re.match +local balancer = require("ngx.balancer") local control_api_router local is_http = false @@ -79,6 +80,23 @@ local ver_header = "APISIX/" .. core.version.VERSION local has_mod, apisix_ngx_client = pcall(require, "resty.apisix.client") local _M = {version = 0.4} +local HTTP_METHODS = { + GET = ngx.HTTP_GET, + HEAD = ngx.HTTP_HEAD, + PUT = ngx.HTTP_PUT, + POST = ngx.HTTP_POST, + DELETE = ngx.HTTP_DELETE, + OPTIONS = ngx.HTTP_OPTIONS, + MKCOL = ngx.HTTP_MKCOL, + COPY = ngx.HTTP_COPY, + MOVE = ngx.HTTP_MOVE, + PROPFIND = ngx.HTTP_PROPFIND, + PROPPATCH = ngx.HTTP_PROPPATCH, + LOCK = ngx.HTTP_LOCK, + UNLOCK = ngx.HTTP_UNLOCK, + PATCH = ngx.HTTP_PATCH, + TRACE = ngx.HTTP_TRACE, + } function _M.http_init(args) @@ -722,7 +740,41 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - _M.handle_upstream(api_ctx, route, enable_websocket) + ngx.req.read_body() + local options = { + always_forward_body = true, + share_all_vars = true, + method = HTTP_METHODS[ngx.req.get_method()], + ctx = ngx.ctx, + } + + local res, err = ngx.location.capture("/subrequest", options) + if not res then + core.log.error("dibaggg: ", err) + return core.response.exit(599) + end + core.log.warn("dibag sub: ", core.json.encode(res, true)) + if res.truncated and options.method ~= ngx.HTTP_HEAD then + return core.response.exit(503) + end + + api_ctx.subreq_status = res.status + api_ctx.subreq_headers = res.header + api_ctx.subreq_body = res.body + + if not api_ctx.custom_upstream_ip then + _M.handle_upstream(api_ctx, route, enable_websocket) + end + + if ngx.ctx.disable_proxy_buffering then + stash_ngx_ctx() + return ngx.exec("@disable_proxy_buffering") + end +end + + +function _M.disable_proxy_buffering_access_phase() + ngx.ctx = fetch_ctx() end @@ -893,7 +945,15 @@ function _M.http_balancer_phase() return core.response.exit(500) end - load_balancer.run(api_ctx.matched_route, api_ctx, common_phase) + if api_ctx.custom_upstream_ip then + local ok, err = balancer.set_current_peer(api_ctx.custom_upstream_ip, api_ctx.custom_upstream_port) + if not ok then + core.log.error("failed to overwrite upstream for ai_proxy: ", err) + return core.response.exit(500) + end + else + load_balancer.run(api_ctx.matched_route, api_ctx, common_phase) + end end diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua new file mode 100644 index 000000000000..fc82a46dae6d --- /dev/null +++ b/apisix/plugins/ai-proxy.lua @@ -0,0 +1,82 @@ +local core = require("apisix.core") +local schema = require("apisix.plugins.ai-proxy.schema") +local constants = require("apisix.constants") + +local ngx_req = ngx.req +local ngx = ngx + +local plugin_name = "ai-proxy" +local _M = { + version = 0.5, + priority = 1002, + name = plugin_name, + schema = schema, +} + + +function _M.check_schema(conf) + -- TODO: check custom URL correctness + return core.schema.check(schema.plugin_schema, conf) +end + + +local CONTENT_TYPE_JSON = "application/json" + + +local function get_request_table() + local req_body, err = core.request.get_body() -- TODO: max size + if not req_body then + return nil, "failed to get request body: " .. err + end + req_body, err = req_body:gsub("\\\"", "\"") -- remove escaping in JSON + if not req_body then + return nil, "failed to remove escaping from body: " .. req_body .. ". err: " .. err + end + return core.json.decode(req_body) +end + +function _M.access(conf, ctx) + local route_type = conf.route_type + ctx.ai_proxy = {} + + local content_type = core.request.header(ctx, "Content-Type") or CONTENT_TYPE_JSON + if content_type ~= CONTENT_TYPE_JSON then + return 400, "unsupported content-type: " .. content_type + end + + local request_table, err = get_request_table() + if not request_table then + return 400, err + end + + local req_schema = schema.chat_request_schema + if route_type == constants.COMPLETION then + req_schema = schema.chat_completion_request_schema + end + local ok, err = core.schema.check(req_schema, request_table) + if not ok then + return 400, "request format doesn't match schema: " .. err + end + + if conf.model.options and conf.model.options.response_streaming then + request_table.stream = true + ngx.ctx.disable_proxy_buffering = true + end + + if conf.model.name then + request_table.model = conf.model.name + end + + if route_type ~= "preserve" then + ngx_req.set_body_data(core.json.encode(request_table)) + end + + local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) + local ok, err = ai_driver.configure_request(conf, ctx) + if not ok then + core.log.error("failed to configure request for AI service: ", err) + return 500 + end +end + +return _M diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua new file mode 100644 index 000000000000..fb8549b9d3c0 --- /dev/null +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -0,0 +1,52 @@ +local _M = {} + +local core = require("apisix.core") + +-- globals +local DEFAULT_HOST = "api.openai.com" +local DEFAULT_PORT = 443 + +local path_mapper = { + ["llm/completions"] = "/v1/completions", + ["llm/chat"] = "/v1/chat/completions", +} + + +function _M.configure_request(conf, ctx) + local ip, err = core.resolver.parse_domain(conf.model.options.upstream_host or DEFAULT_HOST) + if not ip then + core.log.error("failed to resolve ai_proxy upstream host: ", err) + return core.response.exit(500) + end + ctx.custom_upstream_ip = ip + ctx.custom_upstream_port = conf.model.options.upstream_port or DEFAULT_PORT + + local ups_path = (conf.model.options and conf.model.options.upstream_path) or path_mapper[conf.route_type].path + + ngx.var.upstream_uri = ups_path + ngx.var.upstream_scheme = "https" -- TODO: allow override for tests + ngx.var.upstream_host = conf.model.options.upstream_host or DEFAULT_HOST -- TODO: sanity checks. encapsulate to a func + ctx.custom_balancer_host = conf.model.options.upstream_host or DEFAULT_HOST + ctx.custom_balancer_port = conf.model.options.port or DEFAULT_PORT + + local auth_header_name = conf.auth and conf.auth.header_name + local auth_header_value = conf.auth and conf.auth.header_value + local auth_param_name = conf.auth and conf.auth.param_name + local auth_param_value = conf.auth and conf.auth.param_value + local auth_param_location = conf.auth and conf.auth.param_location + + -- TODO: simplify auth structure + if auth_header_name and auth_header_value then + core.request.set_header(ctx, auth_header_name, auth_header_value) + end + + if auth_param_name and auth_param_value and auth_param_location == "query" then -- TODO: test uris + local query_table = core.request.get_uri_args(ctx) + query_table[auth_param_name] = auth_param_value + core.request.set_uri_args(query_table) + end + + return true, nil +end + +return _M diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua new file mode 100644 index 000000000000..35763d402fef --- /dev/null +++ b/apisix/plugins/ai-proxy/schema.lua @@ -0,0 +1,166 @@ +local _M = {} + +local auth_schema = { + type = "object", + properties = { + header_name = { + type = "string", + description = + "Name of the header carrying Authorization or API key.", + }, + header_value = { + type = "string", + description = + "Full auth-header value.", + encrypted = true, -- TODO + }, + param_name = { + type = "string", + description = "Name of the param carrying Authorization or API key.", + }, + param_value = { + type = "string", + description = "full parameter value for 'param_name'.", + encrypted = true, -- TODO + }, + param_location = { + type = "string", + description = + "location of the auth param: query string, or the POST form/JSON body.", + oneOf = { "query", "body" }, + }, + oneOf = { + { required = { "header_name", "header_value" } }, + { required = { "param_name", "param_location", "param_value" } } + } + } +} + +local model_options_schema = { + description = "Key/value settings for the model", + type = "object", + properties = { + max_tokens = { + type = "integer", + description = "Defines the max_tokens, if using chat or completion models.", + default = 256 + + }, + input_cost = { + type = "number", + description = "Defines the cost per 1M tokens in your prompt.", + minimum = 0 + + }, + output_cost = { + type = "number", + description = "Defines the cost per 1M tokens in the output of the AI.", + minimum = 0 + + }, + temperature = { + type = "number", + description = "Defines the matching temperature, if using chat or completion models.", + minimum = 0.0, + maximum = 5.0, + + }, + top_p = { + type = "number", + description = "Defines the top-p probability mass, if supported.", + minimum = 0, + maximum = 1, + + }, + upstream_host = { + type = "string", + description = "To be specified to override the host of the AI provider", + }, + upstream_port = { + type = "integer", + description = "To be specified to override the AI provider port", + + }, + upstream_path = { + type = "string", + description = "To be specified to override the URL to the AI provider endpoints", + }, + response_streaming = { + description = "Stream response by SSE", + type = "boolean", + default = false, + } + } +} + +local model_schema = { + type = "object", + properties = { + provider = { + type = "string", + description = "AI provider request format - kapisix translates " + .. "requests to and from the specified backend compatible formats.", + oneOf = { "openai" }, -- add more providers later + + }, + name = { + type = "string", + description = "Model name to execute.", + }, + options = model_options_schema, + }, + required = {"provider"} +} + +_M.plugin_schema = { + type = "object", + properties = { + route_type = { + type = "string", + description = "The model's operation implementation, for this provider. " .. + "Set to `preserve` to pass through without transformation.", + enum = { "llm/chat", "llm/completions", "passthrough" } + }, + auth = auth_schema, + model = model_schema, + }, + required = {"route_type", "model"} +} + +_M.chat_request_schema = { + type = "object", + properties = { + messages = { + type = "array", + minItems = 1, + items = { + properties = { + role = { + type = "string", + enum = {"system", "user", "assistant"} + }, + content = { + type = "string", + minLength = "1", + }, + }, + additionalProperties = false, + required = {"role", "content"}, + }, + } + }, + required = {"messages"} +} + +_M.chat_completion_request_schema = { + type = "object", + properties = { + prompt = { + type = "string", + minLength = 1 + } + }, + required = {"prompt"} +} + +return _M From 8188ae4e66e0d50bca129964bbbd4ba927e58be5 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:06:31 +0545 Subject: [PATCH 02/85] remove subrequest --- apisix/balancer.lua | 2 -- apisix/cli/ngx_tpl.lua | 32 -------------------------------- apisix/init.lua | 39 --------------------------------------- 3 files changed, 73 deletions(-) diff --git a/apisix/balancer.lua b/apisix/balancer.lua index 76cba643a1e8..0fe2e6539922 100644 --- a/apisix/balancer.lua +++ b/apisix/balancer.lua @@ -285,8 +285,6 @@ do local default_keepalive_pool function set_current_peer(server, ctx) - core.log.warn("dibag peer: ", core.json.encode(server)) - core.log.warn(debug.traceback("dibag")) local up_conf = ctx.upstream_conf local keepalive_pool = up_conf.keepalive_pool diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index bc3a24a470e6..46f2ed857b22 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -808,38 +808,6 @@ http { apisix.http_log_phase() } } - - location /subrequest { - internal; - - access_by_lua_block {;} - header_filter_by_lua_block {;} - body_filter_by_lua_block {;} - log_by_lua_block {;} - - proxy_http_version 1.1; - proxy_set_header Host $upstream_host; - proxy_set_header Upgrade $upstream_upgrade; - proxy_set_header Connection $upstream_connection; - proxy_set_header X-Real-IP $remote_addr; - proxy_pass_header Date; - - - set $var_x_forwarded_proto $scheme; - set $var_x_forwarded_host $host; - set $var_x_forwarded_port $server_port; - - proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto $var_x_forwarded_proto; - proxy_set_header X-Forwarded-Host $var_x_forwarded_host; - proxy_set_header X-Forwarded-Port $var_x_forwarded_port; - - proxy_pass_header Server; - proxy_pass_header Date; - proxy_ssl_name $upstream_host; - proxy_ssl_server_name on; - proxy_pass $upstream_scheme://apisix_backend$upstream_uri; - } location @disable_proxy_buffering { # http server location configuration snippet starts diff --git a/apisix/init.lua b/apisix/init.lua index ba49e6b9c993..87f197ff1d2e 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -80,23 +80,6 @@ local ver_header = "APISIX/" .. core.version.VERSION local has_mod, apisix_ngx_client = pcall(require, "resty.apisix.client") local _M = {version = 0.4} -local HTTP_METHODS = { - GET = ngx.HTTP_GET, - HEAD = ngx.HTTP_HEAD, - PUT = ngx.HTTP_PUT, - POST = ngx.HTTP_POST, - DELETE = ngx.HTTP_DELETE, - OPTIONS = ngx.HTTP_OPTIONS, - MKCOL = ngx.HTTP_MKCOL, - COPY = ngx.HTTP_COPY, - MOVE = ngx.HTTP_MOVE, - PROPFIND = ngx.HTTP_PROPFIND, - PROPPATCH = ngx.HTTP_PROPPATCH, - LOCK = ngx.HTTP_LOCK, - UNLOCK = ngx.HTTP_UNLOCK, - PATCH = ngx.HTTP_PATCH, - TRACE = ngx.HTTP_TRACE, - } function _M.http_init(args) @@ -740,28 +723,6 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - ngx.req.read_body() - local options = { - always_forward_body = true, - share_all_vars = true, - method = HTTP_METHODS[ngx.req.get_method()], - ctx = ngx.ctx, - } - - local res, err = ngx.location.capture("/subrequest", options) - if not res then - core.log.error("dibaggg: ", err) - return core.response.exit(599) - end - core.log.warn("dibag sub: ", core.json.encode(res, true)) - if res.truncated and options.method ~= ngx.HTTP_HEAD then - return core.response.exit(503) - end - - api_ctx.subreq_status = res.status - api_ctx.subreq_headers = res.header - api_ctx.subreq_body = res.body - if not api_ctx.custom_upstream_ip then _M.handle_upstream(api_ctx, route, enable_websocket) end From 7b83b3aed52df36afc5549b8281b9bb2e655adb1 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:45:25 +0545 Subject: [PATCH 03/85] fix diff test --- Makefile | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/Makefile b/Makefile index 21a2389633b3..545a21e4f29f 100644 --- a/Makefile +++ b/Makefile @@ -374,6 +374,12 @@ install: runtime $(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/utils $(ENV_INSTALL) apisix/utils/*.lua $(ENV_INST_LUADIR)/apisix/utils/ + $(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-proxy + $(ENV_INSTALL) apisix/plugins/ai-proxy/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-proxy + + $(ENV_INSTALL) -d $(ENV_INST_LUADIR)/apisix/plugins/ai-proxy/drivers + $(ENV_INSTALL) apisix/plugins/ai-proxy/drivers/*.lua $(ENV_INST_LUADIR)/apisix/plugins/ai-proxy/drivers + $(ENV_INSTALL) bin/apisix $(ENV_INST_BINDIR)/apisix From 97cafa51bf2ed3a948dee331844122660e1ea587 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:46:49 +0545 Subject: [PATCH 04/85] long line fix --- apisix/plugins/ai-proxy/drivers/openai.lua | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index fb8549b9d3c0..ea87a30fc286 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -21,11 +21,12 @@ function _M.configure_request(conf, ctx) ctx.custom_upstream_ip = ip ctx.custom_upstream_port = conf.model.options.upstream_port or DEFAULT_PORT - local ups_path = (conf.model.options and conf.model.options.upstream_path) or path_mapper[conf.route_type].path - + local ups_path = (conf.model.options and conf.model.options.upstream_path) + or path_mapper[conf.route_type] ngx.var.upstream_uri = ups_path ngx.var.upstream_scheme = "https" -- TODO: allow override for tests - ngx.var.upstream_host = conf.model.options.upstream_host or DEFAULT_HOST -- TODO: sanity checks. encapsulate to a func + ngx.var.upstream_host = conf.model.options.upstream_host + or DEFAULT_HOST -- TODO: sanity checks. encapsulate to a func ctx.custom_balancer_host = conf.model.options.upstream_host or DEFAULT_HOST ctx.custom_balancer_port = conf.model.options.port or DEFAULT_PORT @@ -40,7 +41,8 @@ function _M.configure_request(conf, ctx) core.request.set_header(ctx, auth_header_name, auth_header_value) end - if auth_param_name and auth_param_value and auth_param_location == "query" then -- TODO: test uris + -- TODO: test uris + if auth_param_name and auth_param_value and auth_param_location == "query" then local query_table = core.request.get_uri_args(ctx) query_table[auth_param_name] = auth_param_value core.request.set_uri_args(query_table) From 35b1787d924c3b2afebe676111d52c88182e74ec Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:48:01 +0545 Subject: [PATCH 05/85] completions typo in consts --- apisix/constants.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/constants.lua b/apisix/constants.lua index 2995618cd433..83566686c37d 100644 --- a/apisix/constants.lua +++ b/apisix/constants.lua @@ -44,5 +44,5 @@ return { ["/plugin_metadata"] = true, }, CHAT = "llm/chat", - COMPLETION = "llm/completion", + COMPLETION = "llm/completions", } From e18caef44ec31121ded8735794394d452fc0ca2a Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:48:26 +0545 Subject: [PATCH 06/85] license --- apisix/plugins/ai-proxy.lua | 16 ++++++++++++++++ apisix/plugins/ai-proxy/drivers/openai.lua | 16 ++++++++++++++++ apisix/plugins/ai-proxy/schema.lua | 16 ++++++++++++++++ 3 files changed, 48 insertions(+) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index fc82a46dae6d..2b978cd0f87d 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -1,3 +1,19 @@ +-- +-- 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 constants = require("apisix.constants") diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index ea87a30fc286..311555c90bdd 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -1,3 +1,19 @@ +-- +-- 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") diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 35763d402fef..e9e6d6ff126e 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -1,3 +1,19 @@ +-- +-- 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 auth_schema = { From 28f06ae37d36877e9fe700fe0a0a4a20f6eb7c49 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 16 Aug 2024 10:53:16 +0545 Subject: [PATCH 07/85] plugins.t fix --- apisix/cli/config.lua | 2 +- apisix/plugins/ai-proxy.lua | 2 +- t/admin/plugins.t | 1 + 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/apisix/cli/config.lua b/apisix/cli/config.lua index 2e6cf7eed001..a1aca88c5cfb 100644 --- a/apisix/cli/config.lua +++ b/apisix/cli/config.lua @@ -216,8 +216,8 @@ local _M = { "proxy-mirror", "proxy-rewrite", "workflow", - "ai-proxy", "api-breaker", + "ai-proxy", "limit-conn", "limit-count", "limit-req", diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 2b978cd0f87d..206024275065 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -24,7 +24,7 @@ local ngx = ngx local plugin_name = "ai-proxy" local _M = { version = 0.5, - priority = 1002, + priority = 1004, name = plugin_name, schema = schema, } diff --git a/t/admin/plugins.t b/t/admin/plugins.t index 911205f48cb4..664272969924 100644 --- a/t/admin/plugins.t +++ b/t/admin/plugins.t @@ -97,6 +97,7 @@ proxy-mirror proxy-rewrite workflow api-breaker +ai-proxy limit-conn limit-count limit-req From 82f96928f20152169ecfa8154f6a812612e36b6d Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 08:36:41 +0545 Subject: [PATCH 08/85] handle empty req body problem --- apisix/plugins/ai-proxy.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 206024275065..6e280b8038df 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -42,7 +42,7 @@ local CONTENT_TYPE_JSON = "application/json" local function get_request_table() local req_body, err = core.request.get_body() -- TODO: max size if not req_body then - return nil, "failed to get request body: " .. err + return nil, "failed to get request body: " .. (err or "request body is empty") end req_body, err = req_body:gsub("\\\"", "\"") -- remove escaping in JSON if not req_body then From 0577e8e57ee7d403214919b1a05182dfc2e575fe Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 10:05:17 +0545 Subject: [PATCH 09/85] auth schema fix --- apisix/plugins/ai-proxy/drivers/openai.lua | 23 ++++---------- apisix/plugins/ai-proxy/schema.lua | 35 +++++++--------------- 2 files changed, 17 insertions(+), 41 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 311555c90bdd..ce4c892a96f6 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -45,23 +45,12 @@ function _M.configure_request(conf, ctx) or DEFAULT_HOST -- TODO: sanity checks. encapsulate to a func ctx.custom_balancer_host = conf.model.options.upstream_host or DEFAULT_HOST ctx.custom_balancer_port = conf.model.options.port or DEFAULT_PORT - - local auth_header_name = conf.auth and conf.auth.header_name - local auth_header_value = conf.auth and conf.auth.header_value - local auth_param_name = conf.auth and conf.auth.param_name - local auth_param_value = conf.auth and conf.auth.param_value - local auth_param_location = conf.auth and conf.auth.param_location - - -- TODO: simplify auth structure - if auth_header_name and auth_header_value then - core.request.set_header(ctx, auth_header_name, auth_header_value) - end - - -- TODO: test uris - if auth_param_name and auth_param_value and auth_param_location == "query" then - local query_table = core.request.get_uri_args(ctx) - query_table[auth_param_name] = auth_param_value - core.request.set_uri_args(query_table) + if conf.auth.source == "header" then + core.request.set_header(ctx, conf.auth.name, conf.auth.value) + else + local args = core.request.get_uri_args(ctx) + args[conf.auth.name] = conf.auth.value + core.request.set_uri_args(ctx, args) end return true, nil diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index e9e6d6ff126e..da53980b6586 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -19,37 +19,24 @@ local _M = {} local auth_schema = { type = "object", properties = { - header_name = { + source = { type = "string", - description = - "Name of the header carrying Authorization or API key.", + enum = {"header", "param"} }, - header_value = { - type = "string", - description = - "Full auth-header value.", - encrypted = true, -- TODO - }, - param_name = { + name = { type = "string", - description = "Name of the param carrying Authorization or API key.", + description = "Name of the param/header carrying Authorization or API key.", + minLength = 1, }, - param_value = { + value = { type = "string", - description = "full parameter value for 'param_name'.", + description = "Full auth-header/param value.", + minLength = 1, encrypted = true, -- TODO }, - param_location = { - type = "string", - description = - "location of the auth param: query string, or the POST form/JSON body.", - oneOf = { "query", "body" }, - }, - oneOf = { - { required = { "header_name", "header_value" } }, - { required = { "param_name", "param_location", "param_value" } } - } - } + }, + required = { "source", "name", "value" }, + additionalProperties = false, } local model_options_schema = { From e5f00f7344a2cbc901044dfc52d266edd47e3cb3 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 10:05:44 +0545 Subject: [PATCH 10/85] scheme and method --- apisix/plugins/ai-proxy/drivers/openai.lua | 3 ++- t/APISIX.pm | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index ce4c892a96f6..b34f7231ad24 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -40,7 +40,8 @@ function _M.configure_request(conf, ctx) local ups_path = (conf.model.options and conf.model.options.upstream_path) or path_mapper[conf.route_type] ngx.var.upstream_uri = ups_path - ngx.var.upstream_scheme = "https" -- TODO: allow override for tests + ngx.var.upstream_scheme = test_scheme or "https" + ngx.req.set_method(ngx.HTTP_POST) ngx.var.upstream_host = conf.model.options.upstream_host or DEFAULT_HOST -- TODO: sanity checks. encapsulate to a func ctx.custom_balancer_host = conf.model.options.upstream_host or DEFAULT_HOST diff --git a/t/APISIX.pm b/t/APISIX.pm index 50f7cfaecab6..f11d6fd60dbf 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -266,6 +266,7 @@ env APISIX_PROFILE; env PATH; # for searching external plugin runner's binary env TEST_NGINX_HTML_DIR; env OPENSSL_BIN; +env AI_PROXY_TEST_SCHEME; _EOC_ From c307b042b24a5294825f9215be136ad14abe73f5 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 10:06:19 +0545 Subject: [PATCH 11/85] auth and model.name required --- apisix/plugins/ai-proxy/schema.lua | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index da53980b6586..54a5e7fcf926 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -112,7 +112,7 @@ local model_schema = { }, options = model_options_schema, }, - required = {"provider"} + required = {"provider", "name"} } _M.plugin_schema = { @@ -120,14 +120,12 @@ _M.plugin_schema = { properties = { route_type = { type = "string", - description = "The model's operation implementation, for this provider. " .. - "Set to `preserve` to pass through without transformation.", enum = { "llm/chat", "llm/completions", "passthrough" } }, auth = auth_schema, model = model_schema, }, - required = {"route_type", "model"} + required = {"route_type", "model", "auth"} } _M.chat_request_schema = { From ef4cf845e76a8739e7db35dd8fc27f76a77d966a Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 10:06:35 +0545 Subject: [PATCH 12/85] scheme in lua code forgot to commit --- apisix/plugins/ai-proxy/drivers/openai.lua | 1 + 1 file changed, 1 insertion(+) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index b34f7231ad24..15014049d717 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -17,6 +17,7 @@ local _M = {} local core = require("apisix.core") +local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME") -- globals local DEFAULT_HOST = "api.openai.com" From 4bf6bd2e790a26c44a3f752585c3e4feabf29893 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 10:06:43 +0545 Subject: [PATCH 13/85] tests --- t/assets/ai-proxy-response.json | 15 ++ t/plugin/ai-proxy-chat.t | 307 ++++++++++++++++++++++++++++++++ t/plugin/ai-proxy-completions.t | 299 +++++++++++++++++++++++++++++++ 3 files changed, 621 insertions(+) create mode 100644 t/assets/ai-proxy-response.json create mode 100644 t/plugin/ai-proxy-chat.t create mode 100644 t/plugin/ai-proxy-completions.t diff --git a/t/assets/ai-proxy-response.json b/t/assets/ai-proxy-response.json new file mode 100644 index 000000000000..94665e5eaea9 --- /dev/null +++ b/t/assets/ai-proxy-response.json @@ -0,0 +1,15 @@ +{ + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { "content": "1 + 1 = 2.", "role": "assistant" } + } + ], + "created": 1723780938, + "id": "chatcmpl-9wiSIg5LYrrpxwsr2PubSQnbtod1P", + "model": "gpt-4o-2024-05-13", + "object": "chat.completion", + "system_fingerprint": "fp_abc28019ad", + "usage": { "completion_tokens": 8, "prompt_tokens": 23, "total_tokens": 31 } +} diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t new file mode 100644 index 000000000000..23b760f49d6e --- /dev/null +++ b/t/plugin/ai-proxy-chat.t @@ -0,0 +1,307 @@ +# +# 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. +# +BEGIN { + $ENV{AI_PROXY_TEST_SCHEME} = "http"; +} + +use t::APISIX 'no_plan'; + +log_level("info"); +repeat_each(1); +no_long_string(); +no_root_location(); + + +my $resp_file = 't/assets/ai-proxy-response.json'; +open(my $fh, '<', $resp_file) or die "Could not open file '$resp_file' $!"; +my $resp = do { local $/; <$fh> }; +close($fh); + +print "Hello, World!\n"; +print $resp; + + +add_block_preprocessor(sub { + my ($block) = @_; + + if (!defined $block->request) { + $block->set_value("request", "GET /t"); + } + + my $http_config = $block->http_config // <<_EOC_; + server { + server_name openai; + listen 6724; + + default_type 'application/json'; + + location /v1/chat/completions { + content_by_lua_block { + local json = require("cjson.safe") + + if ngx.req.get_method() ~= "POST" then + ngx.status = 400 + ngx.say("Unsupported reqeust method: ", ngx.req.get_method()) + end + ngx.req.read_body() + local body, err = ngx.req.get_body_data() + body, err = json.decode(body) + + local header_auth = ngx.req.get_headers()["authorization"] + local query_auth = ngx.req.get_uri_args()["apikey"] + + if header_auth ~= "Bearer token" and query_auth ~= "apikey" then + ngx.status = 401 + ngx.say("Unauthorized") + return + end + + if header_auth == "Bearer token" or query_auth == "apikey" then + ngx.req.read_body() + local body, err = ngx.req.get_body_data() + local esc = body:gsub('"\\\""', '\"') + body, err = json.decode(esc) + + if body.messages and #body.messages > 1 then + ngx.status = 200 + ngx.say([[$resp]]) + return + else + ngx.status = 400 + ngx.say([[{ "error": "bad request"}]]) + return + end + else + ngx.status = 401 + end + } + } + + } +_EOC_ + + $block->set_value("http_config", $http_config); +}); + +run_tests(); + +__DATA__ + +=== TEST 1: minimal viable configuration +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-proxy") + local ok, err = plugin.check_schema({ + route_type = "llm/chat", + model = { + provider = "openai", + name = "gpt-4", + }, + auth = { + source = "header", + value = "some value", + name = "some name", + } + }) + + if not ok then + ngx.say(err) + else + ngx.say("passed") + end + } + } +--- response_body +passed + + + +=== TEST 2: set route with wrong auth header +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "source": "header", + "name": "Authorization", + "value": "Bearer wrongtoken" + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0, + "upstream_host": "localhost", + "upstream_port": 6724 + } + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 3: send request +--- request +POST /anything +{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } +--- error_code: 401 +--- response_body +Unauthorized + + + +=== TEST 4: set route with right auth header +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "source": "header", + "name": "Authorization", + "value": "Bearer token" + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0, + "upstream_host": "localhost", + "upstream_port": 6724 + } + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 5: send request +--- request +POST /anything +{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } +--- more_headers +Authorization: Bearer token +--- error_code: 200 +--- response_body eval +qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ + + + +=== TEST 6: send request with empty body +--- request +POST /anything +--- more_headers +Authorization: Bearer token +--- error_code: 400 +--- response_body_chomp +failed to get request body: request body is empty + + + +=== TEST 7: send request with wrong method (GET) should work +--- request +GET /anything +{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } +--- more_headers +Authorization: Bearer token +--- error_code: 200 +--- response_body eval +qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ + + + +=== TEST 8: wrong JSON in request body should give error +--- request +GET /anything +{}"messages": [ { "role": "system", "cont +--- error_code: 400 +--- response_body chomp +Expected the end but found T_STRING at character 3 + + + +=== TEST 9: content-type should be JSON +--- request +POST /anything +prompt%3Dwhat%2520is%25201%2520%252B%25201 +--- more_headers +Content-Type: application/x-www-form-urlencoded +--- error_code: 400 +--- response_body chomp +unsupported content-type: application/x-www-form-urlencoded + + + +=== TEST 10: request schema validity check +--- request +POST /anything +{ "messages-missing": [ { "role": "system", "content": "xyz" } ] } +--- more_headers +Authorization: Bearer token +--- error_code: 400 +--- response_body chomp +request format doesn't match schema: property "messages" is required diff --git a/t/plugin/ai-proxy-completions.t b/t/plugin/ai-proxy-completions.t new file mode 100644 index 000000000000..1389602a939f --- /dev/null +++ b/t/plugin/ai-proxy-completions.t @@ -0,0 +1,299 @@ +# +# 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. +# +BEGIN { + $ENV{AI_PROXY_TEST_SCHEME} = "http"; +} + +use t::APISIX 'no_plan'; + +log_level("info"); +repeat_each(1); +no_long_string(); +no_root_location(); + + +my $resp_file = 't/assets/ai-proxy-response.json'; +open(my $fh, '<', $resp_file) or die "Could not open file '$resp_file' $!"; +my $resp = do { local $/; <$fh> }; +close($fh); + +print "Hello, World!\n"; +print $resp; + + +add_block_preprocessor(sub { + my ($block) = @_; + + if (!defined $block->request) { + $block->set_value("request", "GET /t"); + } + + my $http_config = $block->http_config // <<_EOC_; + server { + server_name openai; + listen 6724; + + default_type 'application/json'; + + location /v1/completions { + content_by_lua_block { + local json = require("cjson.safe") + + if ngx.req.get_method() ~= "POST" then + ngx.status = 400 + ngx.say("Unsupported reqeust method: ", ngx.req.get_method()) + end + ngx.req.read_body() + local body, err = ngx.req.get_body_data() + body, err = json.decode(body) + + local header_auth = ngx.req.get_headers()["authorization"] + local query_auth = ngx.req.get_uri_args()["apikey"] + + if header_auth ~= "Bearer token" and query_auth ~= "apikey" then + ngx.status = 401 + ngx.say("Unauthorized") + return + end + + if header_auth == "Bearer token" or query_auth == "apikey" then + ngx.req.read_body() + local body, err = ngx.req.get_body_data() + local esc = body:gsub('"\\\""', '\"') + body, err = json.decode(esc) + + if body.prompt then + ngx.status = 200 + ngx.say([[$resp]]) + return + end + else + ngx.status = 401 + end + } + } + + } +_EOC_ + + $block->set_value("http_config", $http_config); +}); + +run_tests(); + +__DATA__ + +=== TEST 1: minimal viable configuration +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-proxy") + local ok, err = plugin.check_schema({ + route_type = "llm/chat", + model = { + provider = "openai", + name = "gpt-4", + }, + auth = { + source = "header", + value = "some value", + name = "some name", + } + }) + + if not ok then + ngx.say(err) + else + ngx.say("passed") + end + } + } +--- response_body +passed + + + +=== TEST 2: set route with wrong auth header +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/completions", + "auth": { + "source": "header", + "name": "Authorization", + "value": "Bearer wrongtoken" + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0, + "upstream_host": "localhost", + "upstream_port": 6724 + } + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 3: send request +--- request +POST /anything +{"prompt": "what is 1 + 1"} +--- error_code: 401 +--- response_body +Unauthorized + + + +=== TEST 4: set route with correct auth header +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/completions", + "auth": { + "source": "header", + "name": "Authorization", + "value": "Bearer token" + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0, + "upstream_host": "localhost", + "upstream_port": 6724 + } + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 5: send request +--- request +POST /anything +{"prompt": "what is 1 + 1"} +--- error_code: 200 +--- response_body eval +qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ + + + +=== TEST 6: send request with empty body +--- request +POST /anything +--- error_code: 400 +--- response_body_chomp +failed to get request body: request body is empty + + + +=== TEST 7: send request with wrong method (GET) should work +--- request +GET /anything +{"prompt": "what is 1 + 1"} +--- more_headers +Authorization: Bearer token +--- error_code: 200 +--- response_body eval +qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ + + + +=== TEST 8: wrong JSON in request body should give error +--- request +GET /anything +{"prompt": "what is 1 + 1" +--- error_code: 400 +--- response_body chomp +Expected comma or object end but found T_END at character 27 + + + +=== TEST 9: content-type should be JSON +--- request +POST /anything +prompt%3Dwhat%2520is%25201%2520%252B%25201 +--- more_headers +Content-Type: application/x-www-form-urlencoded +--- error_code: 400 +--- response_body chomp +unsupported content-type: application/x-www-form-urlencoded + + + +=== TEST 10: request schema validity check +--- request +POST /anything +{"prompt-is-missing": "what is 1 + 1"} +--- more_headers +Authorization: Bearer token +--- error_code: 400 +--- response_body chomp +request format doesn't match schema: property "prompt" is required From 42adfd148a29a6d86c46a4a3433b5893ad0a0da9 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Sun, 18 Aug 2024 11:13:27 +0545 Subject: [PATCH 14/85] lint --- apisix/cli/ngx_tpl.lua | 2 +- apisix/init.lua | 3 ++- apisix/plugins/ai-proxy.lua | 3 +-- apisix/plugins/ai-proxy/drivers/openai.lua | 3 +-- t/plugin/ai-proxy-chat.t | 4 ++-- t/plugin/ai-proxy-completions.t | 4 ++-- 6 files changed, 9 insertions(+), 10 deletions(-) diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index 46f2ed857b22..ffc8ac55a915 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -808,7 +808,7 @@ http { apisix.http_log_phase() } } - + location @disable_proxy_buffering { # http server location configuration snippet starts {% if http_server_location_configuration_snippet then %} diff --git a/apisix/init.lua b/apisix/init.lua index 87f197ff1d2e..3c912694d375 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -907,7 +907,8 @@ function _M.http_balancer_phase() end if api_ctx.custom_upstream_ip then - local ok, err = balancer.set_current_peer(api_ctx.custom_upstream_ip, api_ctx.custom_upstream_port) + local ok, err = balancer.set_current_peer(api_ctx.custom_upstream_ip, + api_ctx.custom_upstream_port) if not ok then core.log.error("failed to overwrite upstream for ai_proxy: ", err) return core.response.exit(500) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 6e280b8038df..8d0339bc3fbf 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -31,7 +31,6 @@ local _M = { function _M.check_schema(conf) - -- TODO: check custom URL correctness return core.schema.check(schema.plugin_schema, conf) end @@ -40,7 +39,7 @@ local CONTENT_TYPE_JSON = "application/json" local function get_request_table() - local req_body, err = core.request.get_body() -- TODO: max size + local req_body, err = core.request.get_body() if not req_body then return nil, "failed to get request body: " .. (err or "request body is empty") end diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 15014049d717..5e450e7fe504 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -43,8 +43,7 @@ function _M.configure_request(conf, ctx) ngx.var.upstream_uri = ups_path ngx.var.upstream_scheme = test_scheme or "https" ngx.req.set_method(ngx.HTTP_POST) - ngx.var.upstream_host = conf.model.options.upstream_host - or DEFAULT_HOST -- TODO: sanity checks. encapsulate to a func + ngx.var.upstream_host = conf.model.options.upstream_host or DEFAULT_HOST ctx.custom_balancer_host = conf.model.options.upstream_host or DEFAULT_HOST ctx.custom_balancer_port = conf.model.options.port or DEFAULT_PORT if conf.auth.source == "header" then diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 23b760f49d6e..32566b6bd098 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -55,7 +55,7 @@ add_block_preprocessor(sub { if ngx.req.get_method() ~= "POST" then ngx.status = 400 - ngx.say("Unsupported reqeust method: ", ngx.req.get_method()) + ngx.say("Unsupported request method: ", ngx.req.get_method()) end ngx.req.read_body() local body, err = ngx.req.get_body_data() @@ -63,7 +63,7 @@ add_block_preprocessor(sub { local header_auth = ngx.req.get_headers()["authorization"] local query_auth = ngx.req.get_uri_args()["apikey"] - + if header_auth ~= "Bearer token" and query_auth ~= "apikey" then ngx.status = 401 ngx.say("Unauthorized") diff --git a/t/plugin/ai-proxy-completions.t b/t/plugin/ai-proxy-completions.t index 1389602a939f..9f89e62104c5 100644 --- a/t/plugin/ai-proxy-completions.t +++ b/t/plugin/ai-proxy-completions.t @@ -55,7 +55,7 @@ add_block_preprocessor(sub { if ngx.req.get_method() ~= "POST" then ngx.status = 400 - ngx.say("Unsupported reqeust method: ", ngx.req.get_method()) + ngx.say("Unsupported request method: ", ngx.req.get_method()) end ngx.req.read_body() local body, err = ngx.req.get_body_data() @@ -63,7 +63,7 @@ add_block_preprocessor(sub { local header_auth = ngx.req.get_headers()["authorization"] local query_auth = ngx.req.get_uri_args()["apikey"] - + if header_auth ~= "Bearer token" and query_auth ~= "apikey" then ngx.status = 401 ngx.say("Unauthorized") From 0af00aeac17e60d3a1874177af1fd38e2bc5307a Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 20 Aug 2024 11:53:41 +0545 Subject: [PATCH 15/85] add docs --- docs/en/latest/plugins/ai-proxy.md | 144 +++++++++++++++++++++++++++++ 1 file changed, 144 insertions(+) create mode 100644 docs/en/latest/plugins/ai-proxy.md diff --git a/docs/en/latest/plugins/ai-proxy.md b/docs/en/latest/plugins/ai-proxy.md new file mode 100644 index 000000000000..6fe34c744b1f --- /dev/null +++ b/docs/en/latest/plugins/ai-proxy.md @@ -0,0 +1,144 @@ +--- +title: ai-proxy +keywords: + - Apache APISIX + - API Gateway + - Plugin + - ai-proxy +description: This document contains information about the Apache APISIX ai-proxy Plugin. +--- + + + +## Description + +The `ai-proxy` plugin simplifies access to AI providers and models by defining a standard request format +that allows configuring key fields in plugin configuration to embed into the request. + +Proxying requests to OpenAI is supported for now, other AI models will be supported soon. + +## Request Format + +### OpenAI + +- Chat API + +| Name | Type | Required | Description | +| ------------------ | ------ | -------- | --------------------------------------------------- | +| `messages` | Array | Yes | An array of message objects | +| `messages.role` | String | Yes | Role of the message (`system`, `user`, `assistant`) | +| `messages.content` | String | Yes | Content of the message | + +- Completion API + +| Name | Type | Required | Description | +| -------- | ------ | -------- | --------------------------------- | +| `prompt` | String | Yes | Prompt to be sent to the upstream | + +## Plugin Attributes + +| Field | Type | Description | Required | +| ---------------------------------- | ------- | --------------------------------------------------------------------------------------------- | -------- | +| `route_type` | String | Specifies the type of route (`llm/chat`, `llm/completions`, `passthrough`) | Yes | +| `auth` | Object | Authentication configuration | Yes | +| `auth.source` | String | Source of the authentication (`header`, `param`) | Yes | +| `auth.name` | String | Name of the param/header carrying Authorization or API key. Minimum length: 1 | Yes | +| `auth.value` | String | Full auth-header/param value. Minimum length: 1. Encrypted. | Yes | +| `model` | Object | Model configuration | Yes | +| `model.provider` | String | AI provider request format. Translates requests to/from specified backend compatible formats. | Yes | +| `model.name` | String | Model name to execute. | Yes | +| `model.options` | Object | Key/value settings for the model | No | +| `model.options.max_tokens` | Integer | Defines the max_tokens, if using chat or completion models. Default: 256 | No | +| `model.options.temperature` | Number | Defines the matching temperature, if using chat or completion models. Range: 0.0 - 5.0 | No | +| `model.options.top_p` | Number | Defines the top-p probability mass, if supported. Range: 0.0 - 1.0 | No | +| `model.options.upstream_host` | String | To be specified to override the host of the AI provider | No | +| `model.options.upstream_port` | Integer | To be specified to override the AI provider port | No | +| `model.options.upstream_path` | String | To be specified to override the URL to the AI provider endpoints | No | +| `model.options.response_streaming` | Boolean | Stream response by SSE. Default: false | No | + +## Example usage + +Create a route with the `ai-proxy` plugin like so: + +```shell +curl "http://127.0.0.1:9180/apisix/admin/routes/1" -X PUT \ + -H "X-API-KEY: ${ADMIN_API_KEY}" \ + -d '{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "header_name": "Authorization", + "header_value": "Bearer " + }, + "model": { + "provider": "openai", + "name": "gpt-4", + "options": { + "max_tokens": 512, + "temperature": 1.0 + } + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "somerandom.com:443": 1 + } + } + }' +``` + +The upstream node can be any arbitrary value because it won't be contacted. + +Now send a request: + +```shell +curl http://127.0.0.1:9080/anything -i -XPOST -H 'Content-Type: application/json' -d '{ + "messages": [ + { "role": "system", "content": "You are a mathematician" }, + { "role": "user", "a": 1, "content": "What is 1+1?" } + ] + }' +``` + +You will recieve a response like this: +```json +{ + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The sum of \\(1 + 1\\) is \\(2\\).", + "role": "assistant" + } + } + ], + "created": 1723777034, + "id": "chatcmpl-9whRKFodKl5sGhOgHIjWltdeB8sr7", + "model": "gpt-4o-2024-05-13", + "object": "chat.completion", + "system_fingerprint": "fp_abc28019ad", + "usage": { "completion_tokens": 15, "prompt_tokens": 23, "total_tokens": 38 } +} +``` From aff56a0188ad6fcd723efdf6abc93b53bbab5327 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 20 Aug 2024 12:38:12 +0545 Subject: [PATCH 16/85] options merger test --- apisix/plugins/ai-proxy.lua | 15 ++-- apisix/plugins/ai-proxy/drivers/openai.lua | 7 +- t/plugin/ai-proxy-chat.t | 82 ++++++++++++++++++++++ 3 files changed, 98 insertions(+), 6 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 8d0339bc3fbf..975eb3987826 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -82,16 +82,21 @@ function _M.access(conf, ctx) request_table.model = conf.model.name end - if route_type ~= "preserve" then - ngx_req.set_body_data(core.json.encode(request_table)) - end - local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) - local ok, err = ai_driver.configure_request(conf, ctx) + local ok, err = ai_driver.configure_request(conf, request_table, ctx) if not ok then core.log.error("failed to configure request for AI service: ", err) return 500 end + + if route_type ~= "passthrough" then + local final_body, err = core.json.encode(request_table) + if not final_body then + core.log.error("failed to encode request body to JSON: ", err) + return 500 + end + ngx_req.set_body_data(final_body) + end end return _M diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 5e450e7fe504..42f61b101c73 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -29,7 +29,7 @@ local path_mapper = { } -function _M.configure_request(conf, ctx) +function _M.configure_request(conf, request_table, ctx) local ip, err = core.resolver.parse_domain(conf.model.options.upstream_host or DEFAULT_HOST) if not ip then core.log.error("failed to resolve ai_proxy upstream host: ", err) @@ -54,6 +54,11 @@ function _M.configure_request(conf, ctx) core.request.set_uri_args(ctx, args) end + if conf.model.options then + for opt, val in pairs(conf.model.options) do + request_table[opt] = val + end + end return true, nil end diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 32566b6bd098..4428a503aaef 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -61,6 +61,18 @@ add_block_preprocessor(sub { local body, err = ngx.req.get_body_data() body, err = json.decode(body) + local test_type = ngx.req.get_headers()["test-type"] + if test_type == "options" then + if body.foo == "bar" then + ngx.status = 200 + ngx.say("options works") + else + ngx.status = 500 + ngx.say("model options feature doesn't work") + end + return + end + local header_auth = ngx.req.get_headers()["authorization"] local query_auth = ngx.req.get_uri_args()["apikey"] @@ -305,3 +317,73 @@ Authorization: Bearer token --- error_code: 400 --- response_body chomp request format doesn't match schema: property "messages" is required + + + +=== TEST 4: model options being merged to request body +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "source": "header", + "name": "Authorization", + "value": "Bearer token" + }, + "model": { + "provider": "openai", + "name": "some-model", + "options": { + "foo": "bar", + "temperature": 1.0, + "upstream_host": "localhost", + "upstream_port": 6724 + } + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + ngx.say(body) + return + end + + local code, body, actual_body = t("/anything", + ngx.HTTP_POST, + [[{ + "messages": [ + { "role": "system", "content": "You are a mathematician" }, + { "role": "user", "content": "What is 1+1?" } + ] + }]], + nil, + { + ["test-type"] = "options", + ["Content-Type"] = "application/json", + } + ) + + ngx.status = code + ngx.say(actual_body) + + } + } +--- error_code: 200 +--- response_body_chomp +options_works From f25f21ae50c906f9c0faf99d548954498e4478f0 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 20 Aug 2024 12:44:19 +0545 Subject: [PATCH 17/85] fix encryption mode comment --- apisix/plugins/ai-proxy/schema.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 54a5e7fcf926..a657a3e448bb 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -32,7 +32,7 @@ local auth_schema = { type = "string", description = "Full auth-header/param value.", minLength = 1, - encrypted = true, -- TODO + -- TODO encrypted = true, }, }, required = { "source", "name", "value" }, From d2d253e7f502ecd7183d1f199a5bc169022fa405 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 20 Aug 2024 12:46:17 +0545 Subject: [PATCH 18/85] fix(lint): local ngx --- apisix/plugins/ai-proxy/drivers/openai.lua | 1 + 1 file changed, 1 insertion(+) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 42f61b101c73..ad4a52599056 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -18,6 +18,7 @@ local _M = {} local core = require("apisix.core") local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME") +local ngx = ngx -- globals local DEFAULT_HOST = "api.openai.com" From 58ca8a7b40941789a895cecf01782c3871864bbf Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 21 Aug 2024 09:35:39 +0545 Subject: [PATCH 19/85] fix lint --- apisix/plugins/ai-proxy.lua | 1 + apisix/plugins/ai-proxy/drivers/openai.lua | 1 + docs/en/latest/plugins/ai-proxy.md | 3 ++- 3 files changed, 4 insertions(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 975eb3987826..1001e39c7d40 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -17,6 +17,7 @@ local core = require("apisix.core") local schema = require("apisix.plugins.ai-proxy.schema") local constants = require("apisix.constants") +local require = require local ngx_req = ngx.req local ngx = ngx diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index ad4a52599056..f8565c8d36a8 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -19,6 +19,7 @@ local _M = {} local core = require("apisix.core") local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME") local ngx = ngx +local pairs = pairs -- globals local DEFAULT_HOST = "api.openai.com" diff --git a/docs/en/latest/plugins/ai-proxy.md b/docs/en/latest/plugins/ai-proxy.md index 6fe34c744b1f..8813882e03be 100644 --- a/docs/en/latest/plugins/ai-proxy.md +++ b/docs/en/latest/plugins/ai-proxy.md @@ -121,7 +121,8 @@ curl http://127.0.0.1:9080/anything -i -XPOST -H 'Content-Type: application/jso }' ``` -You will recieve a response like this: +You will receive a response like this: + ```json { "choices": [ From f146f2068ba67c37f6246d5038ada3beaface574 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 21 Aug 2024 09:43:24 +0545 Subject: [PATCH 20/85] index to json --- docs/en/latest/config.json | 3 ++- t/plugin/ai-proxy-chat.t | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/docs/en/latest/config.json b/docs/en/latest/config.json index cd6aeb94b444..01aeb3c3353a 100644 --- a/docs/en/latest/config.json +++ b/docs/en/latest/config.json @@ -94,7 +94,8 @@ "plugins/fault-injection", "plugins/mocking", "plugins/degraphql", - "plugins/body-transformer" + "plugins/body-transformer", + "plugins/ai-proxy" ] }, { diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 4428a503aaef..b05f39348470 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -378,7 +378,7 @@ request format doesn't match schema: property "messages" is required ["Content-Type"] = "application/json", } ) - + ngx.status = code ngx.say(actual_body) From 2317aa8df8da0c69e3fd86ede6502c712f6fdaff Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 21 Aug 2024 09:55:12 +0545 Subject: [PATCH 21/85] reindex --- t/plugin/ai-proxy-chat.t | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index b05f39348470..53c10cdff38e 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -320,7 +320,7 @@ request format doesn't match schema: property "messages" is required -=== TEST 4: model options being merged to request body +=== TEST 11: model options being merged to request body --- config location /t { content_by_lua_block { From 3ac0fe50435e797c6818307f133f9702701b4ae1 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 12:28:45 +0545 Subject: [PATCH 22/85] unsupported provider --- apisix/plugins/ai-proxy.lua | 4 +++ t/plugin/ai-proxy-chat.t | 50 +++++++++++++++++++++++++++++-------- 2 files changed, 44 insertions(+), 10 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 1001e39c7d40..5800aba3548c 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -32,6 +32,10 @@ local _M = { 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 diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 53c10cdff38e..055caf24b4d2 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -143,7 +143,37 @@ passed -=== TEST 2: set route with wrong auth header +=== TEST 2: unsupported provider +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-proxy") + local ok, err = plugin.check_schema({ + route_type = "llm/chat", + model = { + provider = "some-unique", + name = "gpt-4", + }, + auth = { + source = "header", + value = "some value", + name = "some name", + } + }) + + if not ok then + ngx.say(err) + else + ngx.say("passed") + end + } + } +--- response_body eval +qr/.*provider: some-unique is not supported.*/ + + + +=== TEST 3: set route with wrong auth header --- config location /t { content_by_lua_block { @@ -192,7 +222,7 @@ passed -=== TEST 3: send request +=== TEST 4: send request --- request POST /anything { "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } @@ -202,7 +232,7 @@ Unauthorized -=== TEST 4: set route with right auth header +=== TEST 5: set route with right auth header --- config location /t { content_by_lua_block { @@ -251,7 +281,7 @@ passed -=== TEST 5: send request +=== TEST 6: send request --- request POST /anything { "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } @@ -263,7 +293,7 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ -=== TEST 6: send request with empty body +=== TEST 7: send request with empty body --- request POST /anything --- more_headers @@ -274,7 +304,7 @@ failed to get request body: request body is empty -=== TEST 7: send request with wrong method (GET) should work +=== TEST 8: send request with wrong method (GET) should work --- request GET /anything { "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } @@ -286,7 +316,7 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ -=== TEST 8: wrong JSON in request body should give error +=== TEST 9: wrong JSON in request body should give error --- request GET /anything {}"messages": [ { "role": "system", "cont @@ -296,7 +326,7 @@ Expected the end but found T_STRING at character 3 -=== TEST 9: content-type should be JSON +=== TEST 10: content-type should be JSON --- request POST /anything prompt%3Dwhat%2520is%25201%2520%252B%25201 @@ -308,7 +338,7 @@ unsupported content-type: application/x-www-form-urlencoded -=== TEST 10: request schema validity check +=== TEST 11: request schema validity check --- request POST /anything { "messages-missing": [ { "role": "system", "content": "xyz" } ] } @@ -320,7 +350,7 @@ request format doesn't match schema: property "messages" is required -=== TEST 11: model options being merged to request body +=== TEST 12: model options being merged to request body --- config location /t { content_by_lua_block { From 6e31cfe0c0f76944168a5f9fa9886c700f7e4d4e Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 12:29:35 +0545 Subject: [PATCH 23/85] remove , nil --- apisix/plugins/ai-proxy/drivers/openai.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index f8565c8d36a8..a4aae6b1cd65 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -61,7 +61,7 @@ function _M.configure_request(conf, request_table, ctx) request_table[opt] = val end end - return true, nil + return true end return _M From e302360466aff188c590a3406c44cbbdc0760ae3 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 12:44:17 +0545 Subject: [PATCH 24/85] move to core.request --- apisix/core/request.lua | 21 +++++++++++++++++++++ apisix/plugins/ai-proxy.lua | 14 +------------- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/apisix/core/request.lua b/apisix/core/request.lua index c5278b6b8072..e39dca7b1d7d 100644 --- a/apisix/core/request.lua +++ b/apisix/core/request.lua @@ -21,6 +21,7 @@ local lfs = require("lfs") local log = require("apisix.core.log") +local json = require("apisix.core.json") local io = require("apisix.core.io") local req_add_header if ngx.config.subsystem == "http" then @@ -334,6 +335,26 @@ function _M.get_body(max_size, ctx) end +function _M.get_request_body_table() + local body, err = _M.get_body() + if not body then + return nil, { message = "could not get body: " .. err } + end + + body, err = body:gsub("\\\"", "\"") -- remove escaping in JSON + if not body then + return nil, { message = "failed to remove escaping from body. err: " .. err} + end + + local body_tab, err = json.decode(body) + if not body_tab then + return nil, { message = "could not get parse JSON request body: " .. err } + end + + return body_tab +end + + function _M.get_scheme(ctx) if not ctx then ctx = ngx.ctx.api_ctx diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 5800aba3548c..4325a4f228fe 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -43,18 +43,6 @@ end local CONTENT_TYPE_JSON = "application/json" -local function get_request_table() - local req_body, err = core.request.get_body() - if not req_body then - return nil, "failed to get request body: " .. (err or "request body is empty") - end - req_body, err = req_body:gsub("\\\"", "\"") -- remove escaping in JSON - if not req_body then - return nil, "failed to remove escaping from body: " .. req_body .. ". err: " .. err - end - return core.json.decode(req_body) -end - function _M.access(conf, ctx) local route_type = conf.route_type ctx.ai_proxy = {} @@ -64,7 +52,7 @@ function _M.access(conf, ctx) return 400, "unsupported content-type: " .. content_type end - local request_table, err = get_request_table() + local request_table, err = core.request.get_request_body_table() if not request_table then return 400, err end From 83f219708ca29bfb28e696e7d00166d63cb45897 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 17:03:15 +0545 Subject: [PATCH 25/85] update empty body test --- apisix/core/request.lua | 2 +- t/plugin/ai-proxy-chat.t | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/apisix/core/request.lua b/apisix/core/request.lua index e39dca7b1d7d..801ce583ffd3 100644 --- a/apisix/core/request.lua +++ b/apisix/core/request.lua @@ -338,7 +338,7 @@ end function _M.get_request_body_table() local body, err = _M.get_body() if not body then - return nil, { message = "could not get body: " .. err } + return nil, { message = "could not get body: " .. (err or "request body is empty") } end body, err = body:gsub("\\\"", "\"") -- remove escaping in JSON diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 055caf24b4d2..6b585a558f96 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -321,8 +321,8 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ GET /anything {}"messages": [ { "role": "system", "cont --- error_code: 400 ---- response_body chomp -Expected the end but found T_STRING at character 3 +--- response_body +{"message":"could not get parse JSON request body: Expected the end but found T_STRING at character 3"} From bcc21cb5410fe87093dbb30613e0d38610f0cdee Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 17:06:27 +0545 Subject: [PATCH 26/85] fix way to set upstream --- apisix/core/resolver.lua | 20 +++++ apisix/init.lua | 11 +-- apisix/plugins/ai-proxy/drivers/openai.lua | 28 ++++--- apisix/plugins/traffic-split.lua | 91 ++-------------------- apisix/upstream.lua | 56 +++++++++++++ 5 files changed, 103 insertions(+), 103 deletions(-) diff --git a/apisix/core/resolver.lua b/apisix/core/resolver.lua index 3568a9762063..790d8945360a 100644 --- a/apisix/core/resolver.lua +++ b/apisix/core/resolver.lua @@ -24,6 +24,7 @@ local log = require("apisix.core.log") local utils = require("apisix.core.utils") local dns_utils = require("resty.dns.utils") local config_local = require("apisix.core.config_local") +local ipmatcher = require("resty.ipmatcher") local HOSTS_IP_MATCH_CACHE = {} @@ -93,4 +94,23 @@ function _M.parse_domain(host) end +function _M.parse_domain_for_node(node) + local host = node.domain or node.host + if not ipmatcher.parse_ipv4(host) + and not ipmatcher.parse_ipv6(host) + then + node.domain = host + + local ip, err = _M.parse_domain(host) + if ip then + node.host = ip + end + + if err then + log.error("dns resolver domain: ", host, " error: ", err) + end + end +end + + return _M diff --git a/apisix/init.lua b/apisix/init.lua index 3c912694d375..4d46caa8d0ca 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -906,16 +906,7 @@ function _M.http_balancer_phase() return core.response.exit(500) end - if api_ctx.custom_upstream_ip then - local ok, err = balancer.set_current_peer(api_ctx.custom_upstream_ip, - api_ctx.custom_upstream_port) - if not ok then - core.log.error("failed to overwrite upstream for ai_proxy: ", err) - return core.response.exit(500) - end - else - load_balancer.run(api_ctx.matched_route, api_ctx, common_phase) - end + load_balancer.run(api_ctx.matched_route, api_ctx, common_phase) end diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index a4aae6b1cd65..ccf44d120abb 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -18,6 +18,7 @@ local _M = {} local core = require("apisix.core") local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME") +local upstream = require("apisix.upstream") local ngx = ngx local pairs = pairs @@ -32,22 +33,29 @@ local path_mapper = { function _M.configure_request(conf, request_table, ctx) - local ip, err = core.resolver.parse_domain(conf.model.options.upstream_host or DEFAULT_HOST) - if not ip then - core.log.error("failed to resolve ai_proxy upstream host: ", err) - return core.response.exit(500) + local ups_host = DEFAULT_HOST + if conf.override and conf.override.upstream_host and conf.override.upstream_host ~= "" then + ups_host = conf.override.upstream_host end - ctx.custom_upstream_ip = ip - ctx.custom_upstream_port = conf.model.options.upstream_port or DEFAULT_PORT + local ups_port = DEFAULT_PORT + if conf.override and conf.override.upstream_port and conf.override.upstream_host ~= "" then + ups_port = conf.override.upstream_host + end + local upstream_addr = ups_host .. ":" .. ups_port + local upstream_node = { + nodes = { + [upstream_addr] = 1 + }, + pass_host = "node", + scheme = test_scheme or "https", + vid = "openai", + } + upstream.set_upstream(upstream_node, ctx) local ups_path = (conf.model.options and conf.model.options.upstream_path) or path_mapper[conf.route_type] ngx.var.upstream_uri = ups_path - ngx.var.upstream_scheme = test_scheme or "https" ngx.req.set_method(ngx.HTTP_POST) - ngx.var.upstream_host = conf.model.options.upstream_host or DEFAULT_HOST - ctx.custom_balancer_host = conf.model.options.upstream_host or DEFAULT_HOST - ctx.custom_balancer_port = conf.model.options.port or DEFAULT_PORT if conf.auth.source == "header" then core.request.set_header(ctx, conf.auth.name, conf.auth.value) else diff --git a/apisix/plugins/traffic-split.lua b/apisix/plugins/traffic-split.lua index f546225c8c95..12d557716c6c 100644 --- a/apisix/plugins/traffic-split.lua +++ b/apisix/plugins/traffic-split.lua @@ -23,7 +23,6 @@ local expr = require("resty.expr.v1") local pairs = pairs local ipairs = ipairs local type = type -local table_insert = table.insert local tostring = tostring local lrucache = core.lrucache.new({ @@ -124,80 +123,6 @@ function _M.check_schema(conf) end -local function parse_domain_for_node(node) - local host = node.domain or node.host - if not ipmatcher.parse_ipv4(host) - and not ipmatcher.parse_ipv6(host) - then - node.domain = host - - local ip, err = core.resolver.parse_domain(host) - if ip then - node.host = ip - end - - if err then - core.log.error("dns resolver domain: ", host, " error: ", err) - end - end -end - - -local function set_upstream(upstream_info, ctx) - local nodes = upstream_info.nodes - local new_nodes = {} - if core.table.isarray(nodes) then - for _, node in ipairs(nodes) do - parse_domain_for_node(node) - table_insert(new_nodes, node) - end - else - for addr, weight in pairs(nodes) do - local node = {} - local port, host - host, port = core.utils.parse_addr(addr) - node.host = host - parse_domain_for_node(node) - node.port = port - node.weight = weight - table_insert(new_nodes, node) - end - end - - local up_conf = { - name = upstream_info.name, - type = upstream_info.type, - hash_on = upstream_info.hash_on, - pass_host = upstream_info.pass_host, - upstream_host = upstream_info.upstream_host, - key = upstream_info.key, - nodes = new_nodes, - timeout = upstream_info.timeout, - scheme = upstream_info.scheme - } - - local ok, err = upstream.check_schema(up_conf) - if not ok then - core.log.error("failed to validate generated upstream: ", err) - return 500, err - end - - local matched_route = ctx.matched_route - up_conf.parent = matched_route - local upstream_key = up_conf.type .. "#route_" .. - matched_route.value.id .. "_" .. upstream_info.vid - if upstream_info.node_tid then - upstream_key = upstream_key .. "_" .. upstream_info.node_tid - end - core.log.info("upstream_key: ", upstream_key) - upstream.set(ctx, upstream_key, ctx.conf_version, up_conf) - if upstream_info.scheme == "https" then - upstream.set_scheme(ctx, up_conf) - end - return -end - - local function new_rr_obj(weighted_upstreams) local server_list = {} for i, upstream_obj in ipairs(weighted_upstreams) do @@ -287,18 +212,18 @@ function _M.access(conf, ctx) return 500 end - local upstream = rr_up:find() - if upstream and type(upstream) == "table" then - core.log.info("upstream: ", core.json.encode(upstream)) - return set_upstream(upstream, ctx) - elseif upstream and upstream ~= "plugin#upstream#is#empty" then - ctx.upstream_id = upstream - core.log.info("upstream_id: ", upstream) + local rr_ups = rr_up:find() + if rr_ups and type(rr_ups) == "table" then + core.log.info("upstream: ", core.json.encode(rr_ups)) + return upstream.set_upstream(rr_ups, ctx) + elseif rr_ups and rr_ups ~= "plugin#upstream#is#empty" then + ctx.upstream_id = rr_ups + core.log.info("upstream_id: ", rr_ups) return end ctx.upstream_id = nil - core.log.info("route_up: ", upstream) + core.log.info("route_up: ", rr_ups) return end diff --git a/apisix/upstream.lua b/apisix/upstream.lua index eb5e467daaca..591a99a375b5 100644 --- a/apisix/upstream.lua +++ b/apisix/upstream.lua @@ -20,11 +20,13 @@ local discovery = require("apisix.discovery.init").discovery local upstream_util = require("apisix.utils.upstream") local apisix_ssl = require("apisix.ssl") local events = require("apisix.events") +local resolver = require("apisix.core.resolver") local error = error local tostring = tostring local ipairs = ipairs local pairs = pairs local pcall = pcall +local table_insert = table.insert local ngx_var = ngx.var local is_http = ngx.config.subsystem == "http" local upstreams @@ -252,6 +254,60 @@ local function fill_node_info(up_conf, scheme, is_stream) end +function _M.set_upstream(upstream_info, ctx) + local nodes = upstream_info.nodes + local new_nodes = {} + if core.table.isarray(nodes) then + for _, node in ipairs(nodes) do + core.utils.parse_domain_for_node(node) + table_insert(new_nodes, node) + end + else + for addr, weight in pairs(nodes) do + local node = {} + local port, host + host, port = core.utils.parse_addr(addr) + node.host = host + resolver.parse_domain_for_node(node) + node.port = port + node.weight = weight + table_insert(new_nodes, node) + end + end + + local up_conf = { + name = upstream_info.name, + type = upstream_info.type, + hash_on = upstream_info.hash_on, + pass_host = upstream_info.pass_host, + upstream_host = upstream_info.upstream_host, + key = upstream_info.key, + nodes = new_nodes, + timeout = upstream_info.timeout, + scheme = upstream_info.scheme + } + + local ok, err = _M.check_schema(up_conf) + if not ok then + core.log.error("failed to validate generated upstream: ", err) + return 500, err + end + + local matched_route = ctx.matched_route + up_conf.parent = matched_route + local upstream_key = up_conf.type .. "#route_" .. + matched_route.value.id .. "_" .. upstream_info.vid + if upstream_info.node_tid then + upstream_key = upstream_key .. "_" .. upstream_info.node_tid + end + core.log.info("upstream_key: ", upstream_key) + _M.set(ctx, upstream_key, ctx.conf_version, up_conf) + if upstream_info.scheme == "https" then + _M.set_scheme(ctx, up_conf) + end +end + + function _M.set_by_route(route, api_ctx) if api_ctx.upstream_conf then -- upstream_conf has been set by traffic-split plugin From 10a07c150dbd739963599a5e24ac411b14674151 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 17:13:16 +0545 Subject: [PATCH 27/85] add log --- apisix/plugins/ai-proxy.lua | 1 + 1 file changed, 1 insertion(+) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 4325a4f228fe..4cbeabd500f1 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -84,6 +84,7 @@ function _M.access(conf, ctx) if route_type ~= "passthrough" then local final_body, err = core.json.encode(request_table) + core.log.info("final parsed body: ", final_body) if not final_body then core.log.error("failed to encode request body to JSON: ", err) return 500 From 7220c08936536c104cc0d048d00d51d5de0a7ab9 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 17:30:50 +0545 Subject: [PATCH 28/85] response_streaming -> stream --- apisix/plugins/ai-proxy.lua | 2 +- apisix/plugins/ai-proxy/schema.lua | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 4cbeabd500f1..0a506d2704d2 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -66,7 +66,7 @@ function _M.access(conf, ctx) return 400, "request format doesn't match schema: " .. err end - if conf.model.options and conf.model.options.response_streaming then + if conf.model.options and conf.model.options.stream then request_table.stream = true ngx.ctx.disable_proxy_buffering = true end diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index a657a3e448bb..dc92dd16e918 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -88,7 +88,7 @@ local model_options_schema = { type = "string", description = "To be specified to override the URL to the AI provider endpoints", }, - response_streaming = { + stream = { description = "Stream response by SSE", type = "boolean", default = false, From a4afb30554e08eb9c395ed727aec587feaeff7ca Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 20:22:43 +0545 Subject: [PATCH 29/85] refactor override schema --- apisix/plugins/ai-proxy/drivers/openai.lua | 11 ++- apisix/plugins/ai-proxy/schema.lua | 17 ++++ t/plugin/ai-proxy-chat.t | 101 +++++++++++++++++++-- 3 files changed, 115 insertions(+), 14 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index ccf44d120abb..85f5f3676a6a 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -34,14 +34,15 @@ local path_mapper = { function _M.configure_request(conf, request_table, ctx) local ups_host = DEFAULT_HOST - if conf.override and conf.override.upstream_host and conf.override.upstream_host ~= "" then - ups_host = conf.override.upstream_host + if conf.override and conf.override.host and conf.override.host ~= "" then + ups_host = conf.override.host end local ups_port = DEFAULT_PORT - if conf.override and conf.override.upstream_port and conf.override.upstream_host ~= "" then - ups_port = conf.override.upstream_host + if conf.override and conf.override.port and conf.override.host ~= "" then + ups_port = conf.override.port end local upstream_addr = ups_host .. ":" .. ups_port + core.log.info("modified upstream address: ", upstream_addr) local upstream_node = { nodes = { [upstream_addr] = 1 @@ -52,7 +53,7 @@ function _M.configure_request(conf, request_table, ctx) } upstream.set_upstream(upstream_node, ctx) - local ups_path = (conf.model.options and conf.model.options.upstream_path) + local ups_path = (conf.override and conf.override.path) or path_mapper[conf.route_type] ngx.var.upstream_uri = ups_path ngx.req.set_method(ngx.HTTP_POST) diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index dc92dd16e918..88b340d8a51c 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -111,6 +111,23 @@ local model_schema = { description = "Model name to execute.", }, options = model_options_schema, + override = { + type = "object", + properties = { + host = { + type = "string", + description = "To be specified to override the host of the AI provider", + }, + port = { + type = "integer", + description = "To be specified to override the AI provider port", + }, + path = { + type = "string", + description = "To be specified to override the URL to the AI provider endpoints", + }, + } + } }, required = {"provider", "name"} } diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 6b585a558f96..682caf0ad83c 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -103,6 +103,11 @@ add_block_preprocessor(sub { } } + location /random { + content_by_lua_block { + ngx.say("path override works") + } + } } _EOC_ @@ -195,10 +200,12 @@ qr/.*provider: some-unique is not supported.*/ "name": "gpt-35-turbo-instruct", "options": { "max_tokens": 512, - "temperature": 1.0, - "upstream_host": "localhost", - "upstream_port": 6724 + "temperature": 1.0 } + }, + "override": { + "host": "localhost", + "port": 6724 } } }, @@ -254,10 +261,12 @@ Unauthorized "name": "gpt-35-turbo-instruct", "options": { "max_tokens": 512, - "temperature": 1.0, - "upstream_host": "localhost", - "upstream_port": 6724 + "temperature": 1.0 } + }, + "override": { + "host": "localhost", + "port": 6724 } } }, @@ -372,10 +381,12 @@ request format doesn't match schema: property "messages" is required "name": "some-model", "options": { "foo": "bar", - "temperature": 1.0, - "upstream_host": "localhost", - "upstream_port": 6724 + "temperature": 1.0 } + }, + "override": { + "host": "localhost", + "port": 6724 } } }, @@ -417,3 +428,75 @@ request format doesn't match schema: property "messages" is required --- error_code: 200 --- response_body_chomp options_works + + + +=== TEST 12: override path +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "source": "header", + "name": "Authorization", + "value": "Bearer token" + }, + "model": { + "provider": "openai", + "name": "some-model", + "options": { + "foo": "bar", + "temperature": 1.0 + } + }, + "override": { + "host": "localhost", + "port": 6724, + "path": "/random" + } + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + ngx.say(body) + return + end + + local code, body, actual_body = t("/anything", + ngx.HTTP_POST, + [[{ + "messages": [ + { "role": "system", "content": "You are a mathematician" }, + { "role": "user", "content": "What is 1+1?" } + ] + }]], + nil, + { + ["test-type"] = "path", + ["Content-Type"] = "application/json", + } + ) + + ngx.status = code + ngx.say(actual_body) + + } + } +--- response_body_chomp +path override works From 624800560d68f59e33d6c26dd09db8af8c28046f Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 20:24:45 +0545 Subject: [PATCH 30/85] content type update --- apisix/plugins/ai-proxy.lua | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 0a506d2704d2..ac51e6238802 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -47,9 +47,9 @@ function _M.access(conf, ctx) local route_type = conf.route_type ctx.ai_proxy = {} - local content_type = core.request.header(ctx, "Content-Type") or CONTENT_TYPE_JSON - if content_type ~= CONTENT_TYPE_JSON then - return 400, "unsupported content-type: " .. content_type + local ct = core.request.header(ctx, "Content-Type") or CONTENT_TYPE_JSON + if ct ~= CONTENT_TYPE_JSON or ct ~= CONTENT_TYPE_JSON .. ";charset=utf8" then + return 400, "unsupported content-type: " .. ct end local request_table, err = core.request.get_request_body_table() From e88683ced3fa8226f096000079aa7605f26d9f8b Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 20:54:45 +0545 Subject: [PATCH 31/85] remove completions --- apisix/plugins/ai-proxy.lua | 6 +- apisix/plugins/ai-proxy/drivers/openai.lua | 1 - apisix/plugins/ai-proxy/schema.lua | 16 +- t/plugin/ai-proxy-completions.t | 299 --------------------- 4 files changed, 3 insertions(+), 319 deletions(-) delete mode 100644 t/plugin/ai-proxy-completions.t diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index ac51e6238802..0d0d94d9cb89 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -57,11 +57,7 @@ function _M.access(conf, ctx) return 400, err end - local req_schema = schema.chat_request_schema - if route_type == constants.COMPLETION then - req_schema = schema.chat_completion_request_schema - end - local ok, err = core.schema.check(req_schema, request_table) + 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 diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 85f5f3676a6a..91159872db4c 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -27,7 +27,6 @@ local DEFAULT_HOST = "api.openai.com" local DEFAULT_PORT = 443 local path_mapper = { - ["llm/completions"] = "/v1/completions", ["llm/chat"] = "/v1/chat/completions", } diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 88b340d8a51c..f4ca58a0754a 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -101,8 +101,7 @@ local model_schema = { properties = { provider = { type = "string", - description = "AI provider request format - kapisix translates " - .. "requests to and from the specified backend compatible formats.", + description = "Name of the AI service provider.", oneOf = { "openai" }, -- add more providers later }, @@ -137,7 +136,7 @@ _M.plugin_schema = { properties = { route_type = { type = "string", - enum = { "llm/chat", "llm/completions", "passthrough" } + enum = { "llm/chat", "passthrough" } }, auth = auth_schema, model = model_schema, @@ -170,15 +169,4 @@ _M.chat_request_schema = { required = {"messages"} } -_M.chat_completion_request_schema = { - type = "object", - properties = { - prompt = { - type = "string", - minLength = 1 - } - }, - required = {"prompt"} -} - return _M diff --git a/t/plugin/ai-proxy-completions.t b/t/plugin/ai-proxy-completions.t deleted file mode 100644 index 9f89e62104c5..000000000000 --- a/t/plugin/ai-proxy-completions.t +++ /dev/null @@ -1,299 +0,0 @@ -# -# 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. -# -BEGIN { - $ENV{AI_PROXY_TEST_SCHEME} = "http"; -} - -use t::APISIX 'no_plan'; - -log_level("info"); -repeat_each(1); -no_long_string(); -no_root_location(); - - -my $resp_file = 't/assets/ai-proxy-response.json'; -open(my $fh, '<', $resp_file) or die "Could not open file '$resp_file' $!"; -my $resp = do { local $/; <$fh> }; -close($fh); - -print "Hello, World!\n"; -print $resp; - - -add_block_preprocessor(sub { - my ($block) = @_; - - if (!defined $block->request) { - $block->set_value("request", "GET /t"); - } - - my $http_config = $block->http_config // <<_EOC_; - server { - server_name openai; - listen 6724; - - default_type 'application/json'; - - location /v1/completions { - content_by_lua_block { - local json = require("cjson.safe") - - if ngx.req.get_method() ~= "POST" then - ngx.status = 400 - ngx.say("Unsupported request method: ", ngx.req.get_method()) - end - ngx.req.read_body() - local body, err = ngx.req.get_body_data() - body, err = json.decode(body) - - local header_auth = ngx.req.get_headers()["authorization"] - local query_auth = ngx.req.get_uri_args()["apikey"] - - if header_auth ~= "Bearer token" and query_auth ~= "apikey" then - ngx.status = 401 - ngx.say("Unauthorized") - return - end - - if header_auth == "Bearer token" or query_auth == "apikey" then - ngx.req.read_body() - local body, err = ngx.req.get_body_data() - local esc = body:gsub('"\\\""', '\"') - body, err = json.decode(esc) - - if body.prompt then - ngx.status = 200 - ngx.say([[$resp]]) - return - end - else - ngx.status = 401 - end - } - } - - } -_EOC_ - - $block->set_value("http_config", $http_config); -}); - -run_tests(); - -__DATA__ - -=== TEST 1: minimal viable configuration ---- config - location /t { - content_by_lua_block { - local plugin = require("apisix.plugins.ai-proxy") - local ok, err = plugin.check_schema({ - route_type = "llm/chat", - model = { - provider = "openai", - name = "gpt-4", - }, - auth = { - source = "header", - value = "some value", - name = "some name", - } - }) - - if not ok then - ngx.say(err) - else - ngx.say("passed") - end - } - } ---- response_body -passed - - - -=== TEST 2: set route with wrong auth header ---- config - location /t { - content_by_lua_block { - local t = require("lib.test_admin").test - local code, body = t('/apisix/admin/routes/1', - ngx.HTTP_PUT, - [[{ - "uri": "/anything", - "plugins": { - "ai-proxy": { - "route_type": "llm/completions", - "auth": { - "source": "header", - "name": "Authorization", - "value": "Bearer wrongtoken" - }, - "model": { - "provider": "openai", - "name": "gpt-35-turbo-instruct", - "options": { - "max_tokens": 512, - "temperature": 1.0, - "upstream_host": "localhost", - "upstream_port": 6724 - } - } - } - }, - "upstream": { - "type": "roundrobin", - "nodes": { - "canbeanything.com": 1 - } - } - }]] - ) - - if code >= 300 then - ngx.status = code - end - ngx.say(body) - } - } ---- response_body -passed - - - -=== TEST 3: send request ---- request -POST /anything -{"prompt": "what is 1 + 1"} ---- error_code: 401 ---- response_body -Unauthorized - - - -=== TEST 4: set route with correct auth header ---- config - location /t { - content_by_lua_block { - local t = require("lib.test_admin").test - local code, body = t('/apisix/admin/routes/1', - ngx.HTTP_PUT, - [[{ - "uri": "/anything", - "plugins": { - "ai-proxy": { - "route_type": "llm/completions", - "auth": { - "source": "header", - "name": "Authorization", - "value": "Bearer token" - }, - "model": { - "provider": "openai", - "name": "gpt-35-turbo-instruct", - "options": { - "max_tokens": 512, - "temperature": 1.0, - "upstream_host": "localhost", - "upstream_port": 6724 - } - } - } - }, - "upstream": { - "type": "roundrobin", - "nodes": { - "canbeanything.com": 1 - } - } - }]] - ) - - if code >= 300 then - ngx.status = code - end - ngx.say(body) - } - } ---- response_body -passed - - - -=== TEST 5: send request ---- request -POST /anything -{"prompt": "what is 1 + 1"} ---- error_code: 200 ---- response_body eval -qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ - - - -=== TEST 6: send request with empty body ---- request -POST /anything ---- error_code: 400 ---- response_body_chomp -failed to get request body: request body is empty - - - -=== TEST 7: send request with wrong method (GET) should work ---- request -GET /anything -{"prompt": "what is 1 + 1"} ---- more_headers -Authorization: Bearer token ---- error_code: 200 ---- response_body eval -qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ - - - -=== TEST 8: wrong JSON in request body should give error ---- request -GET /anything -{"prompt": "what is 1 + 1" ---- error_code: 400 ---- response_body chomp -Expected comma or object end but found T_END at character 27 - - - -=== TEST 9: content-type should be JSON ---- request -POST /anything -prompt%3Dwhat%2520is%25201%2520%252B%25201 ---- more_headers -Content-Type: application/x-www-form-urlencoded ---- error_code: 400 ---- response_body chomp -unsupported content-type: application/x-www-form-urlencoded - - - -=== TEST 10: request schema validity check ---- request -POST /anything -{"prompt-is-missing": "what is 1 + 1"} ---- more_headers -Authorization: Bearer token ---- error_code: 400 ---- response_body chomp -request format doesn't match schema: property "prompt" is required From 7d9c0757f26590f51793600d9cd178931139fc10 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 21:51:42 +0545 Subject: [PATCH 32/85] source -> type --- apisix/plugins/ai-proxy/drivers/openai.lua | 2 +- apisix/plugins/ai-proxy/schema.lua | 4 ++-- t/plugin/ai-proxy-chat.t | 12 ++++++------ 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 91159872db4c..8c71fc5e0f15 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -56,7 +56,7 @@ function _M.configure_request(conf, request_table, ctx) or path_mapper[conf.route_type] ngx.var.upstream_uri = ups_path ngx.req.set_method(ngx.HTTP_POST) - if conf.auth.source == "header" then + if conf.auth.type == "header" then core.request.set_header(ctx, conf.auth.name, conf.auth.value) else local args = core.request.get_uri_args(ctx) diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index f4ca58a0754a..f41995deb974 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -19,7 +19,7 @@ local _M = {} local auth_schema = { type = "object", properties = { - source = { + type = { type = "string", enum = {"header", "param"} }, @@ -35,7 +35,7 @@ local auth_schema = { -- TODO encrypted = true, }, }, - required = { "source", "name", "value" }, + required = { "type", "name", "value" }, additionalProperties = false, } diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy-chat.t index 682caf0ad83c..81b5b30c2e48 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy-chat.t @@ -130,7 +130,7 @@ __DATA__ name = "gpt-4", }, auth = { - source = "header", + type = "header", value = "some value", name = "some name", } @@ -160,7 +160,7 @@ passed name = "gpt-4", }, auth = { - source = "header", + type = "header", value = "some value", name = "some name", } @@ -191,7 +191,7 @@ qr/.*provider: some-unique is not supported.*/ "ai-proxy": { "route_type": "llm/chat", "auth": { - "source": "header", + "type": "header", "name": "Authorization", "value": "Bearer wrongtoken" }, @@ -252,7 +252,7 @@ Unauthorized "ai-proxy": { "route_type": "llm/chat", "auth": { - "source": "header", + "type": "header", "name": "Authorization", "value": "Bearer token" }, @@ -372,7 +372,7 @@ request format doesn't match schema: property "messages" is required "ai-proxy": { "route_type": "llm/chat", "auth": { - "source": "header", + "type": "header", "name": "Authorization", "value": "Bearer token" }, @@ -444,7 +444,7 @@ options_works "ai-proxy": { "route_type": "llm/chat", "auth": { - "source": "header", + "type": "header", "name": "Authorization", "value": "Bearer token" }, From 9823570f18bb5c2828cc988ad114ca6e0d43c8ba Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 22 Aug 2024 22:09:46 +0545 Subject: [PATCH 33/85] fix lint --- apisix/init.lua | 1 - apisix/plugins/ai-proxy.lua | 1 - apisix/plugins/ai-proxy/schema.lua | 2 +- apisix/plugins/traffic-split.lua | 2 -- 4 files changed, 1 insertion(+), 5 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index 4d46caa8d0ca..b94b9759221a 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -59,7 +59,6 @@ local tonumber = tonumber local type = type local pairs = pairs local ngx_re_match = ngx.re.match -local balancer = require("ngx.balancer") local control_api_router local is_http = false diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 0d0d94d9cb89..5e8f7f880cb5 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -16,7 +16,6 @@ -- local core = require("apisix.core") local schema = require("apisix.plugins.ai-proxy.schema") -local constants = require("apisix.constants") local require = require local ngx_req = ngx.req diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index f41995deb974..3889bab97bfa 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -123,7 +123,7 @@ local model_schema = { }, path = { type = "string", - description = "To be specified to override the URL to the AI provider endpoints", + description = "Overrieds the request path to the AI provider endpoints", }, } } diff --git a/apisix/plugins/traffic-split.lua b/apisix/plugins/traffic-split.lua index 12d557716c6c..7b7e05357d21 100644 --- a/apisix/plugins/traffic-split.lua +++ b/apisix/plugins/traffic-split.lua @@ -18,9 +18,7 @@ local core = require("apisix.core") local upstream = require("apisix.upstream") local schema_def = require("apisix.schema_def") local roundrobin = require("resty.roundrobin") -local ipmatcher = require("resty.ipmatcher") local expr = require("resty.expr.v1") -local pairs = pairs local ipairs = ipairs local type = type local tostring = tostring From 94d00f47117f9ce58c2cd4e5d78e254f8c9a3941 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 08:04:27 +0545 Subject: [PATCH 34/85] or -> and --- apisix/plugins/ai-proxy.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 5e8f7f880cb5..08febb3b7477 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -47,7 +47,7 @@ function _M.access(conf, ctx) ctx.ai_proxy = {} local ct = core.request.header(ctx, "Content-Type") or CONTENT_TYPE_JSON - if ct ~= CONTENT_TYPE_JSON or ct ~= CONTENT_TYPE_JSON .. ";charset=utf8" then + if ct ~= CONTENT_TYPE_JSON and ct ~= CONTENT_TYPE_JSON .. ";charset=utf8" then return 400, "unsupported content-type: " .. ct end From b24e439bf0fc88325037ab365a55edfdd41d55b8 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 08:04:43 +0545 Subject: [PATCH 35/85] core.utils -> resolver --- apisix/upstream.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/upstream.lua b/apisix/upstream.lua index 591a99a375b5..d3a7c34cd129 100644 --- a/apisix/upstream.lua +++ b/apisix/upstream.lua @@ -259,7 +259,7 @@ function _M.set_upstream(upstream_info, ctx) local new_nodes = {} if core.table.isarray(nodes) then for _, node in ipairs(nodes) do - core.utils.parse_domain_for_node(node) + resolver.parse_domain_for_node(node) table_insert(new_nodes, node) end else From 5ca70f369d956e06451fb37f3be26a3022415426 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 08:12:46 +0545 Subject: [PATCH 36/85] global pcall --- apisix/plugins/ai-proxy.lua | 1 + 1 file changed, 1 insertion(+) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 08febb3b7477..ad860d6e28e2 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -17,6 +17,7 @@ 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 = ngx From 284ad76d2c882df8f498c96ffe6548ba65ed2fa9 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 08:32:01 +0545 Subject: [PATCH 37/85] rname test file --- t/plugin/{ai-proxy-chat.t => ai-proxy.t} | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename t/plugin/{ai-proxy-chat.t => ai-proxy.t} (99%) diff --git a/t/plugin/ai-proxy-chat.t b/t/plugin/ai-proxy.t similarity index 99% rename from t/plugin/ai-proxy-chat.t rename to t/plugin/ai-proxy.t index 81b5b30c2e48..c44f4398c177 100644 --- a/t/plugin/ai-proxy-chat.t +++ b/t/plugin/ai-proxy.t @@ -431,7 +431,7 @@ options_works -=== TEST 12: override path +=== TEST 13: override path --- config location /t { content_by_lua_block { From 6baa7d16e56a3816b0713109767df6dfb4c320ee Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 13:31:34 +0545 Subject: [PATCH 38/85] use has_prefix --- apisix/plugins/ai-proxy.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index ad860d6e28e2..2d36c7cf7b6d 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -48,7 +48,7 @@ function _M.access(conf, ctx) ctx.ai_proxy = {} local ct = core.request.header(ctx, "Content-Type") or CONTENT_TYPE_JSON - if ct ~= CONTENT_TYPE_JSON and ct ~= CONTENT_TYPE_JSON .. ";charset=utf8" then + if not core.string.has_prefix(ct, CONTENT_TYPE_JSON) then return 400, "unsupported content-type: " .. ct end From bdab563786f9c0e1159ff711ca84992acec83eab Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 13:32:53 +0545 Subject: [PATCH 39/85] fix upstream handling --- apisix/init.lua | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index b94b9759221a..0b0a30033d3b 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -722,9 +722,7 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - if not api_ctx.custom_upstream_ip then - _M.handle_upstream(api_ctx, route, enable_websocket) - end + _M.handle_upstream(api_ctx, route, enable_websocket) if ngx.ctx.disable_proxy_buffering then stash_ngx_ctx() From 2d0a7a1c9a819e44d6f74372f2f4341bcee74f52 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 23 Aug 2024 13:39:43 +0545 Subject: [PATCH 40/85] dont modify tfsp --- apisix/plugins/traffic-split.lua | 93 +++++++++++++++++++++++++++++--- 1 file changed, 85 insertions(+), 8 deletions(-) diff --git a/apisix/plugins/traffic-split.lua b/apisix/plugins/traffic-split.lua index 7b7e05357d21..f546225c8c95 100644 --- a/apisix/plugins/traffic-split.lua +++ b/apisix/plugins/traffic-split.lua @@ -18,9 +18,12 @@ local core = require("apisix.core") local upstream = require("apisix.upstream") local schema_def = require("apisix.schema_def") local roundrobin = require("resty.roundrobin") +local ipmatcher = require("resty.ipmatcher") local expr = require("resty.expr.v1") +local pairs = pairs local ipairs = ipairs local type = type +local table_insert = table.insert local tostring = tostring local lrucache = core.lrucache.new({ @@ -121,6 +124,80 @@ function _M.check_schema(conf) end +local function parse_domain_for_node(node) + local host = node.domain or node.host + if not ipmatcher.parse_ipv4(host) + and not ipmatcher.parse_ipv6(host) + then + node.domain = host + + local ip, err = core.resolver.parse_domain(host) + if ip then + node.host = ip + end + + if err then + core.log.error("dns resolver domain: ", host, " error: ", err) + end + end +end + + +local function set_upstream(upstream_info, ctx) + local nodes = upstream_info.nodes + local new_nodes = {} + if core.table.isarray(nodes) then + for _, node in ipairs(nodes) do + parse_domain_for_node(node) + table_insert(new_nodes, node) + end + else + for addr, weight in pairs(nodes) do + local node = {} + local port, host + host, port = core.utils.parse_addr(addr) + node.host = host + parse_domain_for_node(node) + node.port = port + node.weight = weight + table_insert(new_nodes, node) + end + end + + local up_conf = { + name = upstream_info.name, + type = upstream_info.type, + hash_on = upstream_info.hash_on, + pass_host = upstream_info.pass_host, + upstream_host = upstream_info.upstream_host, + key = upstream_info.key, + nodes = new_nodes, + timeout = upstream_info.timeout, + scheme = upstream_info.scheme + } + + local ok, err = upstream.check_schema(up_conf) + if not ok then + core.log.error("failed to validate generated upstream: ", err) + return 500, err + end + + local matched_route = ctx.matched_route + up_conf.parent = matched_route + local upstream_key = up_conf.type .. "#route_" .. + matched_route.value.id .. "_" .. upstream_info.vid + if upstream_info.node_tid then + upstream_key = upstream_key .. "_" .. upstream_info.node_tid + end + core.log.info("upstream_key: ", upstream_key) + upstream.set(ctx, upstream_key, ctx.conf_version, up_conf) + if upstream_info.scheme == "https" then + upstream.set_scheme(ctx, up_conf) + end + return +end + + local function new_rr_obj(weighted_upstreams) local server_list = {} for i, upstream_obj in ipairs(weighted_upstreams) do @@ -210,18 +287,18 @@ function _M.access(conf, ctx) return 500 end - local rr_ups = rr_up:find() - if rr_ups and type(rr_ups) == "table" then - core.log.info("upstream: ", core.json.encode(rr_ups)) - return upstream.set_upstream(rr_ups, ctx) - elseif rr_ups and rr_ups ~= "plugin#upstream#is#empty" then - ctx.upstream_id = rr_ups - core.log.info("upstream_id: ", rr_ups) + local upstream = rr_up:find() + if upstream and type(upstream) == "table" then + core.log.info("upstream: ", core.json.encode(upstream)) + return set_upstream(upstream, ctx) + elseif upstream and upstream ~= "plugin#upstream#is#empty" then + ctx.upstream_id = upstream + core.log.info("upstream_id: ", upstream) return end ctx.upstream_id = nil - core.log.info("route_up: ", rr_ups) + core.log.info("route_up: ", upstream) return end From 530448fd8d9d76886ece6217e4da5d3e840f307a Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 27 Aug 2024 14:46:01 +0545 Subject: [PATCH 41/85] subrequest --- apisix/cli/ngx_tpl.lua | 21 +++++++++++++++++ apisix/init.lua | 45 +++++++++++++++++++++++++++++++++++-- apisix/plugins/ai-proxy.lua | 4 +++- t/APISIX.pm | 21 +++++++++++++++++ 4 files changed, 88 insertions(+), 3 deletions(-) diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index ffc8ac55a915..027ac863b57f 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -809,6 +809,27 @@ http { } } + location /subrequest { + internal; + + proxy_http_version 1.1; + proxy_set_header Host $upstream_host; + proxy_set_header Upgrade $upstream_upgrade; + proxy_set_header Connection $upstream_connection; + proxy_set_header X-Real-IP $remote_addr; + + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $scheme; + proxy_set_header X-Forwarded-Host $host; + proxy_set_header X-Forwarded-Port $server_port; + + proxy_pass_header Server; + proxy_pass_header Date; + proxy_ssl_name $upstream_host; + proxy_ssl_server_name on; + proxy_pass $upstream_scheme://apisix_backend$upstream_uri; + } + location @disable_proxy_buffering { # http server location configuration snippet starts {% if http_server_location_configuration_snippet then %} diff --git a/apisix/init.lua b/apisix/init.lua index 0b0a30033d3b..e962c52b513e 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -458,6 +458,45 @@ local function common_phase(phase_name) return plugin.run_plugin(phase_name, nil, api_ctx) end +local methods_map = { + GET = ngx.HTTP_GET, + PUT = ngx.HTTP_PUT, + POST = ngx.HTTP_POST, + PATCH = ngx.HTTP_PATCH, + DELETE = ngx.HTTP_DELETE, + OPTIONS = ngx.HTTP_OPTIONS, + TRACE = ngx.HTTP_TRACE, + HEAD = ngx.HTTP_HEAD, +} + +local function subrequest(api_ctx) + ngx.req.read_body() + local options = { + always_forward_body = true, + share_all_vars = true, + method = methods_map[ngx.req.get_method()], + ctx = ngx.ctx, + } + + local res = ngx.location.capture("/subrequest", options) + if not res or res.truncated then + return core.response.exit(502) + end + + if res.truncated and options.method ~= ngx.HTTP_HEAD then + return core.response.exit(503) + end + + api_ctx.subreq_status = res.status + api_ctx.subreq_headers = res.header + api_ctx.subreq_body = res.body + + for key, value in pairs(res.header) do + core.response.set_header(key, value) + end + core.log.info("finishing subrequest") + core.response.exit(res.status, res.body) +end function _M.handle_upstream(api_ctx, route, enable_websocket) @@ -723,8 +762,10 @@ function _M.http_access_phase() end _M.handle_upstream(api_ctx, route, enable_websocket) - - if ngx.ctx.disable_proxy_buffering then + if api_ctx.subrequest then + subrequest(api_ctx) + end + if api_ctx.disable_proxy_buffering then stash_ngx_ctx() return ngx.exec("@disable_proxy_buffering") end diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 2d36c7cf7b6d..12e61aef0080 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -64,7 +64,9 @@ function _M.access(conf, ctx) if conf.model.options and conf.model.options.stream then request_table.stream = true - ngx.ctx.disable_proxy_buffering = true + ctx.disable_proxy_buffering = true + else + ctx.subrequest = true end if conf.model.name then diff --git a/t/APISIX.pm b/t/APISIX.pm index f11d6fd60dbf..d27c9fad335f 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -778,6 +778,27 @@ _EOC_ } } + location /subrequest { + internal; + + proxy_http_version 1.1; + proxy_set_header Host \$upstream_host; + proxy_set_header Upgrade \$upstream_upgrade; + proxy_set_header Connection \$upstream_connection; + proxy_set_header X-Real-IP \$remote_addr; + + proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto \$scheme; + proxy_set_header X-Forwarded-Host \$host; + proxy_set_header X-Forwarded-Port \$server_port; + + proxy_pass_header Server; + proxy_pass_header Date; + proxy_ssl_name \$upstream_host; + proxy_ssl_server_name on; + proxy_pass \$upstream_scheme://apisix_backend\$upstream_uri; + } + location / { set \$upstream_mirror_host ''; set \$upstream_mirror_uri ''; From ed11fa439d78232a11a8e16cdc2180f9e821296e Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 27 Aug 2024 14:49:53 +0545 Subject: [PATCH 42/85] subrequest log check --- t/plugin/ai-proxy.t | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index c44f4398c177..ce362c91d05b 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -299,7 +299,8 @@ Authorization: Bearer token --- error_code: 200 --- response_body eval qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ - +--- error_log +finishing subrequest === TEST 7: send request with empty body From cba307af1426931dd628479bd843abc55bb36193 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 27 Aug 2024 15:17:27 +0545 Subject: [PATCH 43/85] unused var --- apisix/plugins/ai-proxy.lua | 1 - 1 file changed, 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 12e61aef0080..5e3017d0fdfb 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -20,7 +20,6 @@ local require = require local pcall = pcall local ngx_req = ngx.req -local ngx = ngx local plugin_name = "ai-proxy" local _M = { From 3febd29a1c9af494b00e0c4df5dfc78ec3425d45 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 28 Aug 2024 08:08:42 +0545 Subject: [PATCH 44/85] http version check --- apisix/init.lua | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/apisix/init.lua b/apisix/init.lua index e962c52b513e..d630300ff551 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -763,8 +763,15 @@ function _M.http_access_phase() _M.handle_upstream(api_ctx, route, enable_websocket) if api_ctx.subrequest then - subrequest(api_ctx) + local version = ngx.req.http_version() + if version < 2 then + subrequest(api_ctx) + else + api_ctx.subrequest = nil + core.log.error("cannot perform subrequest in HTTP version: ", version) + end end + if api_ctx.disable_proxy_buffering then stash_ngx_ctx() return ngx.exec("@disable_proxy_buffering") From e566a37618fe41cb64898ab80ffb07787610b771 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 30 Aug 2024 09:39:25 +0545 Subject: [PATCH 45/85] reindex --- t/plugin/ai-proxy.t | 1 + 1 file changed, 1 insertion(+) diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index ce362c91d05b..18fae6cc83cc 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -303,6 +303,7 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ finishing subrequest + === TEST 7: send request with empty body --- request POST /anything From 28c9c4d902d0ba5c33a7a0f43bb74519a15df6ac Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 3 Sep 2024 11:09:07 +0545 Subject: [PATCH 46/85] Revert "subrequest" This reverts commit 530448fd8d9d76886ece6217e4da5d3e840f307a. --- apisix/cli/ngx_tpl.lua | 21 ---------------- apisix/init.lua | 50 +------------------------------------ apisix/plugins/ai-proxy.lua | 4 +-- t/APISIX.pm | 21 ---------------- 4 files changed, 2 insertions(+), 94 deletions(-) diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index 027ac863b57f..ffc8ac55a915 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -809,27 +809,6 @@ http { } } - location /subrequest { - internal; - - proxy_http_version 1.1; - proxy_set_header Host $upstream_host; - proxy_set_header Upgrade $upstream_upgrade; - proxy_set_header Connection $upstream_connection; - proxy_set_header X-Real-IP $remote_addr; - - proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto $scheme; - proxy_set_header X-Forwarded-Host $host; - proxy_set_header X-Forwarded-Port $server_port; - - proxy_pass_header Server; - proxy_pass_header Date; - proxy_ssl_name $upstream_host; - proxy_ssl_server_name on; - proxy_pass $upstream_scheme://apisix_backend$upstream_uri; - } - location @disable_proxy_buffering { # http server location configuration snippet starts {% if http_server_location_configuration_snippet then %} diff --git a/apisix/init.lua b/apisix/init.lua index d630300ff551..0b0a30033d3b 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -458,45 +458,6 @@ local function common_phase(phase_name) return plugin.run_plugin(phase_name, nil, api_ctx) end -local methods_map = { - GET = ngx.HTTP_GET, - PUT = ngx.HTTP_PUT, - POST = ngx.HTTP_POST, - PATCH = ngx.HTTP_PATCH, - DELETE = ngx.HTTP_DELETE, - OPTIONS = ngx.HTTP_OPTIONS, - TRACE = ngx.HTTP_TRACE, - HEAD = ngx.HTTP_HEAD, -} - -local function subrequest(api_ctx) - ngx.req.read_body() - local options = { - always_forward_body = true, - share_all_vars = true, - method = methods_map[ngx.req.get_method()], - ctx = ngx.ctx, - } - - local res = ngx.location.capture("/subrequest", options) - if not res or res.truncated then - return core.response.exit(502) - end - - if res.truncated and options.method ~= ngx.HTTP_HEAD then - return core.response.exit(503) - end - - api_ctx.subreq_status = res.status - api_ctx.subreq_headers = res.header - api_ctx.subreq_body = res.body - - for key, value in pairs(res.header) do - core.response.set_header(key, value) - end - core.log.info("finishing subrequest") - core.response.exit(res.status, res.body) -end function _M.handle_upstream(api_ctx, route, enable_websocket) @@ -762,17 +723,8 @@ function _M.http_access_phase() end _M.handle_upstream(api_ctx, route, enable_websocket) - if api_ctx.subrequest then - local version = ngx.req.http_version() - if version < 2 then - subrequest(api_ctx) - else - api_ctx.subrequest = nil - core.log.error("cannot perform subrequest in HTTP version: ", version) - end - end - if api_ctx.disable_proxy_buffering then + if ngx.ctx.disable_proxy_buffering then stash_ngx_ctx() return ngx.exec("@disable_proxy_buffering") end diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 5e3017d0fdfb..3684354efd6b 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -63,9 +63,7 @@ function _M.access(conf, ctx) if conf.model.options and conf.model.options.stream then request_table.stream = true - ctx.disable_proxy_buffering = true - else - ctx.subrequest = true + ngx.ctx.disable_proxy_buffering = true end if conf.model.name then diff --git a/t/APISIX.pm b/t/APISIX.pm index d27c9fad335f..f11d6fd60dbf 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -778,27 +778,6 @@ _EOC_ } } - location /subrequest { - internal; - - proxy_http_version 1.1; - proxy_set_header Host \$upstream_host; - proxy_set_header Upgrade \$upstream_upgrade; - proxy_set_header Connection \$upstream_connection; - proxy_set_header X-Real-IP \$remote_addr; - - proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto \$scheme; - proxy_set_header X-Forwarded-Host \$host; - proxy_set_header X-Forwarded-Port \$server_port; - - proxy_pass_header Server; - proxy_pass_header Date; - proxy_ssl_name \$upstream_host; - proxy_ssl_server_name on; - proxy_pass \$upstream_scheme://apisix_backend\$upstream_uri; - } - location / { set \$upstream_mirror_host ''; set \$upstream_mirror_uri ''; From 780561df9f2a14104bfdac3a5d9c39298cabdd41 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 3 Sep 2024 11:09:26 +0545 Subject: [PATCH 47/85] Revert "subrequest log check" This reverts commit ed11fa439d78232a11a8e16cdc2180f9e821296e. --- t/plugin/ai-proxy.t | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 18fae6cc83cc..0e5eec4dee34 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -299,8 +299,7 @@ Authorization: Bearer token --- error_code: 200 --- response_body eval qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ ---- error_log -finishing subrequest + From 4ffdd8547a6d1f67513221a7345c8779186b580a Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 3 Sep 2024 21:04:16 +0545 Subject: [PATCH 48/85] use httpc to request LLM --- apisix/init.lua | 5 ++ apisix/plugins/ai-proxy.lua | 41 ++++++++++---- apisix/plugins/ai-proxy/drivers/openai.lua | 66 +++++++++++++--------- apisix/plugins/ai-proxy/schema.lua | 25 ++++---- t/plugin/ai-proxy.t | 12 ++-- 5 files changed, 93 insertions(+), 56 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index 0b0a30033d3b..1606f9bbdc1b 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -733,6 +733,11 @@ end function _M.disable_proxy_buffering_access_phase() ngx.ctx = fetch_ctx() + local api_ctx = ngx.ctx.api_ctx + local plugins = plugin.filter(api_ctx, api_ctx.matched_route) + + -- plugins to be run after proxy_buffering is disabled + plugin.run_plugin("delayed_access", plugins, api_ctx) end diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 3684354efd6b..b5b4bf4a6160 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -43,9 +43,6 @@ local CONTENT_TYPE_JSON = "application/json" function _M.access(conf, ctx) - local route_type = conf.route_type - ctx.ai_proxy = {} - 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 @@ -69,22 +66,42 @@ function _M.access(conf, ctx) if conf.model.name then request_table.model = conf.model.name end + ctx.request_table = request_table +end + +function _M.delayed_access(conf, ctx) + local request_table = ctx.request_table local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) - local ok, err = ai_driver.configure_request(conf, request_table, ctx) - if not ok then - core.log.error("failed to configure request for AI service: ", err) + 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 end - if route_type ~= "passthrough" then - local final_body, err = core.json.encode(request_table) - core.log.info("final parsed body: ", final_body) - if not final_body then - core.log.error("failed to encode request body to JSON: ", err) + + if core.table.try_read_attr(conf, "model", "options", "stream") then + local chunk, err + local content_length = 0 + while true do + chunk, err = res.body_reader() -- will read chunk by chunk + if err then + core.log.error("failed to read response chunk: ", err) + return 500 + end + content_length = content_length + (#chunk or 0) + ngx.print(chunk) + ngx.flush(true) + end + core.response.set_header("Content-Length", content_length) + httpc:set_keepalive(10000, 100) + else + local res_body, err = res:read_body() + if not res_body then + core.log.error("failed to read response body: ", err) return 500 end - ngx_req.set_body_data(final_body) + return res.status, res_body end end diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 8c71fc5e0f15..e1382cf5a196 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -17,6 +17,7 @@ local _M = {} local core = require("apisix.core") +local http = require("resty.http") local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME") local upstream = require("apisix.upstream") local ngx = ngx @@ -30,38 +31,41 @@ local path_mapper = { ["llm/chat"] = "/v1/chat/completions", } - -function _M.configure_request(conf, request_table, ctx) - local ups_host = DEFAULT_HOST - if conf.override and conf.override.host and conf.override.host ~= "" then - ups_host = conf.override.host +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 - local ups_port = DEFAULT_PORT - if conf.override and conf.override.port and conf.override.host ~= "" then - ups_port = conf.override.port + httpc:set_timeout(conf.timeout) + + local custom_host = core.table.try_read_attr(conf, "override", "host") + local custom_port = core.table.try_read_attr(conf, "override", "port") + + local ok, err = httpc:connect({ + scheme = test_scheme or "https", + host = custom_host or DEFAULT_HOST, + port = custom_port or DEFAULT_PORT, + ssl_verify = conf.ssl_verify, + ssl_server_name = custom_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 upstream_addr = ups_host .. ":" .. ups_port - core.log.info("modified upstream address: ", upstream_addr) - local upstream_node = { - nodes = { - [upstream_addr] = 1 + + local params = { + method = "POST", + headers = { + ["Content-Type"] = "application/json", }, - pass_host = "node", - scheme = test_scheme or "https", - vid = "openai", + keepalive = conf.keepalive, + ssl_verify = conf.ssl_verify, + path = "/v1/chat/completions", } - upstream.set_upstream(upstream_node, ctx) - local ups_path = (conf.override and conf.override.path) - or path_mapper[conf.route_type] - ngx.var.upstream_uri = ups_path - ngx.req.set_method(ngx.HTTP_POST) if conf.auth.type == "header" then - core.request.set_header(ctx, conf.auth.name, conf.auth.value) - else - local args = core.request.get_uri_args(ctx) - args[conf.auth.name] = conf.auth.value - core.request.set_uri_args(ctx, args) + params.headers[conf.auth.name] = conf.auth.value end if conf.model.options then @@ -69,7 +73,15 @@ function _M.configure_request(conf, request_table, ctx) request_table[opt] = val end end - return true + params.body = core.json.encode(request_table) + + local res, err = httpc:request(params) + if not res then + return 500, "failed to send request to LLM server: " .. err + end + + -- TOOD: keepalive maintainance + return res, nil, httpc end return _M diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 3889bab97bfa..fc87f47d3f3d 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -75,19 +75,6 @@ local model_options_schema = { maximum = 1, }, - upstream_host = { - type = "string", - description = "To be specified to override the host of the AI provider", - }, - upstream_port = { - type = "integer", - description = "To be specified to override the AI provider port", - - }, - upstream_path = { - type = "string", - description = "To be specified to override the URL to the AI provider endpoints", - }, stream = { description = "Stream response by SSE", type = "boolean", @@ -140,6 +127,18 @@ _M.plugin_schema = { }, auth = auth_schema, model = model_schema, + passthrough = { type = "boolean", default = false }, + timeout = { + type = "integer", + minimum = 1, + maximum = 60000, + default = 3000, + description = "timeout in milliseconds", + }, + keepalive = {type = "boolean", default = true}, + keepalive_timeout = {type = "integer", minimum = 1000, default = 60000}, + keepalive_pool = {type = "integer", minimum = 1, default = 30}, + ssl_verify = {type = "boolean", default = true }, }, required = {"route_type", "model", "auth"} } diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 0e5eec4dee34..8ff29fcfa9b1 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -206,7 +206,8 @@ qr/.*provider: some-unique is not supported.*/ "override": { "host": "localhost", "port": 6724 - } + }, + "ssl_verify": false } }, "upstream": { @@ -267,7 +268,8 @@ Unauthorized "override": { "host": "localhost", "port": 6724 - } + }, + "ssl_verify": false } }, "upstream": { @@ -388,7 +390,8 @@ request format doesn't match schema: property "messages" is required "override": { "host": "localhost", "port": 6724 - } + }, + "ssl_verify": false } }, "upstream": { @@ -461,7 +464,8 @@ options_works "host": "localhost", "port": 6724, "path": "/random" - } + }, + "ssl_verify": false } }, "upstream": { From 0521f896f7c60db7fe9e242c5687c8c1ba17507a Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 3 Sep 2024 21:44:10 +0545 Subject: [PATCH 49/85] pass through test --- apisix/plugins/ai-proxy.lua | 77 ++++++++++++++++---------- t/plugin/ai-proxy.t | 107 +++++++++++++++++++++++++++++++++--- 2 files changed, 147 insertions(+), 37 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index b5b4bf4a6160..e0bb8a77a6b9 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -42,35 +42,7 @@ end local CONTENT_TYPE_JSON = "application/json" -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() - 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.options and conf.model.options.stream then - request_table.stream = true - ngx.ctx.disable_proxy_buffering = true - end - - if conf.model.name then - request_table.model = conf.model.name - end - ctx.request_table = request_table -end - - -function _M.delayed_access(conf, ctx) +local function send_request(conf, ctx) local request_table = ctx.request_table local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) local res, err, httpc = ai_driver.request(conf, request_table, ctx) @@ -79,6 +51,17 @@ function _M.delayed_access(conf, ctx) return 500 end + if conf.passthrough then + -- do we need a buffer to cache entire LLM response? + -- i think so, we can do something like the following, just read, no return + local res_body, err = res:read_body() + if not res_body then + core.log.error("failed to read response body: ", err) + return 500 + end + ngx_req.set_body_data(res_body) + return + end if core.table.try_read_attr(conf, "model", "options", "stream") then local chunk, err @@ -105,4 +88,40 @@ function _M.delayed_access(conf, ctx) end 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() + 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 + ctx.request_table = request_table + + if conf.model.options and conf.model.options.stream then + request_table.stream = true + ngx.ctx.disable_proxy_buffering = true + return + end + + return send_request(conf, ctx) +end + + +function _M.delayed_access(conf, ctx) + return send_request(conf, ctx) +end + return _M diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 8ff29fcfa9b1..0ae3d78751a3 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -49,6 +49,26 @@ add_block_preprocessor(sub { default_type 'application/json'; + location /anything { + content_by_lua_block { + local json = require("cjson.safe") + + if ngx.req.get_method() ~= "POST" then + ngx.status = 400 + ngx.say("Unsupported request method: ", ngx.req.get_method()) + end + ngx.req.read_body() + local body = ngx.req.get_body_data() + + if body ~= "SELECT * FROM STUDENTS" then + ngx.status = 503 + ngx.say("passthrough doesn't work") + return + end + ngx.say('{"foo", "bar"}') + } + } + location /v1/chat/completions { content_by_lua_block { local json = require("cjson.safe") @@ -88,18 +108,25 @@ add_block_preprocessor(sub { local esc = body:gsub('"\\\""', '\"') body, err = json.decode(esc) - if body.messages and #body.messages > 1 then - ngx.status = 200 - ngx.say([[$resp]]) - return - else + if not body.messages or #body.messages < 1 then ngx.status = 400 ngx.say([[{ "error": "bad request"}]]) return end - else - ngx.status = 401 + + if body.messages[1].content == "write an SQL query to get all rows from student table" then + ngx.print("SELECT * FROM STUDENTS") + return + end + + ngx.status = 200 + ngx.say([[$resp]]) + return end + + + ngx.status = 503 + ngx.say("reached the end of the test suite") } } @@ -304,7 +331,6 @@ qr/\{ "content": "1 \+ 1 = 2\.", "role": "assistant" \}/ - === TEST 7: send request with empty body --- request POST /anything @@ -505,3 +531,68 @@ options_works } --- response_body_chomp path override works + + + +=== TEST 14: set route with right auth header +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "type": "header", + "name": "Authorization", + "value": "Bearer token" + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0 + } + }, + "override": { + "host": "localhost", + "port": 6724 + }, + "ssl_verify": false, + "passthrough": true + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "127.0.0.1:6724": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 15: send request with wrong method should work +--- request +POST /anything +{ "messages": [ { "role": "user", "content": "write an SQL query to get all rows from student table" } ] } +--- more_headers +Authorization: Bearer token +--- error_code: 200 +--- response_body +{"foo", "bar"} From bd8309b48d5160567a01b188a01d8955f395efed Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 3 Sep 2024 23:25:37 +0545 Subject: [PATCH 50/85] clean up --- apisix/plugins/ai-proxy.lua | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index e0bb8a77a6b9..d37d75e60fe6 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -52,8 +52,6 @@ local function send_request(conf, ctx) end if conf.passthrough then - -- do we need a buffer to cache entire LLM response? - -- i think so, we can do something like the following, just read, no return local res_body, err = res:read_body() if not res_body then core.log.error("failed to read response body: ", err) @@ -64,19 +62,20 @@ local function send_request(conf, ctx) end if core.table.try_read_attr(conf, "model", "options", "stream") then - local chunk, err local content_length = 0 while true do - chunk, err = res.body_reader() -- will read chunk by chunk + local chunk, err = res.body_reader() -- will read chunk by chunk if err then core.log.error("failed to read response chunk: ", err) - return 500 + return end - content_length = content_length + (#chunk or 0) + if not chunk then + return + end + content_length = content_length + #chunk ngx.print(chunk) ngx.flush(true) end - core.response.set_header("Content-Length", content_length) httpc:set_keepalive(10000, 100) else local res_body, err = res:read_body() From 6f5d15847544e458c124b7ae4e7a57516af011e2 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 11:39:01 +0545 Subject: [PATCH 51/85] don't handle upstream when proxy_buffering --- apisix/init.lua | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index 1606f9bbdc1b..736767a4e9b6 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -722,12 +722,12 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - _M.handle_upstream(api_ctx, route, enable_websocket) - if ngx.ctx.disable_proxy_buffering then stash_ngx_ctx() return ngx.exec("@disable_proxy_buffering") end + + _M.handle_upstream(api_ctx, route, enable_websocket) end From d4f3a5a8486b919e09ac32226b58eacda6e4f938 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 11:41:04 +0545 Subject: [PATCH 52/85] test scheme fix --- apisix/plugins/ai-proxy/drivers/openai.lua | 4 ++-- apisix/plugins/ai-proxy/schema.lua | 3 +++ t/plugin/ai-proxy.t | 18 ++++++++++-------- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index e1382cf5a196..15453d614b55 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -18,7 +18,6 @@ local _M = {} local core = require("apisix.core") local http = require("resty.http") -local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME") local upstream = require("apisix.upstream") local ngx = ngx local pairs = pairs @@ -40,9 +39,10 @@ function _M.request(conf, request_table, ctx) local custom_host = core.table.try_read_attr(conf, "override", "host") local custom_port = core.table.try_read_attr(conf, "override", "port") + local custom_scheme = core.table.try_read_attr(conf, "override", "scheme") local ok, err = httpc:connect({ - scheme = test_scheme or "https", + scheme = custom_scheme or "https", host = custom_host or DEFAULT_HOST, port = custom_port or DEFAULT_PORT, ssl_verify = conf.ssl_verify, diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index fc87f47d3f3d..9e7c1c62f7f6 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -112,6 +112,9 @@ local model_schema = { type = "string", description = "Overrieds the request path to the AI provider endpoints", }, + scheme = { + type = "string" + } } } }, diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 0ae3d78751a3..78e2e5ba4a90 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -14,9 +14,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # -BEGIN { - $ENV{AI_PROXY_TEST_SCHEME} = "http"; -} use t::APISIX 'no_plan'; @@ -232,7 +229,8 @@ qr/.*provider: some-unique is not supported.*/ }, "override": { "host": "localhost", - "port": 6724 + "port": 6724, + "scheme": "http" }, "ssl_verify": false } @@ -294,7 +292,8 @@ Unauthorized }, "override": { "host": "localhost", - "port": 6724 + "port": 6724, + "scheme": "http" }, "ssl_verify": false } @@ -415,7 +414,8 @@ request format doesn't match schema: property "messages" is required }, "override": { "host": "localhost", - "port": 6724 + "port": 6724, + "scheme": "http" }, "ssl_verify": false } @@ -489,7 +489,8 @@ options_works "override": { "host": "localhost", "port": 6724, - "path": "/random" + "path": "/random", + "scheme": "http" }, "ssl_verify": false } @@ -561,7 +562,8 @@ path override works }, "override": { "host": "localhost", - "port": 6724 + "port": 6724, + "scheme": "http" }, "ssl_verify": false, "passthrough": true From 1bfeac9b95dcc7344e995369dd28249a1484e90e Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 11:41:20 +0545 Subject: [PATCH 53/85] test path fix --- apisix/plugins/ai-proxy/drivers/openai.lua | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 15453d614b55..a558bfcdc2a1 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -39,6 +39,7 @@ function _M.request(conf, request_table, ctx) local custom_host = core.table.try_read_attr(conf, "override", "host") local custom_port = core.table.try_read_attr(conf, "override", "port") + local custom_path = core.table.try_read_attr(conf, "override", "path") local custom_scheme = core.table.try_read_attr(conf, "override", "scheme") local ok, err = httpc:connect({ @@ -61,7 +62,7 @@ function _M.request(conf, request_table, ctx) }, keepalive = conf.keepalive, ssl_verify = conf.ssl_verify, - path = "/v1/chat/completions", + path = custom_path or "/v1/chat/completions", } if conf.auth.type == "header" then From 25975c0e0068c79e20d868940cec872898da7867 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 11:43:30 +0545 Subject: [PATCH 54/85] support SSE in tests --- ci/common.sh | 21 +++++ ci/linux_openresty_common_runner.sh | 2 + ci/redhat-ci.sh | 2 + t/APISIX.pm | 35 +++++++++ t/plugin/ai-proxy.t | 116 ++++++++++++++++++++++++++++ t/sse_server_example/go.mod | 3 + t/sse_server_example/main.go | 42 ++++++++++ 7 files changed, 221 insertions(+) create mode 100644 t/sse_server_example/go.mod create mode 100644 t/sse_server_example/main.go diff --git a/ci/common.sh b/ci/common.sh index 146b7aa5080a..ae5d12b2b7c6 100644 --- a/ci/common.sh +++ b/ci/common.sh @@ -203,3 +203,24 @@ function start_grpc_server_example() { ss -lntp | grep 10051 | grep grpc_server && break done } + + +function start_sse_server_example() { + # build sse_server_example + pushd t/sse_server_example + go build + ./sse_server_example 7737 2>&1 & + + for (( i = 0; i <= 10; i++ )); do + sleep 0.5 + SSE_PROC=`ps -ef | grep sse_server_example | grep -v grep || echo "none"` + if [[ $SSE_PROC == "none" || "$i" -eq 10 ]]; then + echo "failed to start sse_server_example" + ss -antp | grep 7737 || echo "no proc listen port 7737" + exit 1 + else + break + fi + done + popd +} diff --git a/ci/linux_openresty_common_runner.sh b/ci/linux_openresty_common_runner.sh index ea2e8b41c8bb..1b73ceec92c6 100755 --- a/ci/linux_openresty_common_runner.sh +++ b/ci/linux_openresty_common_runner.sh @@ -77,6 +77,8 @@ script() { start_grpc_server_example + start_sse_server_example + # APISIX_ENABLE_LUACOV=1 PERL5LIB=.:$PERL5LIB prove -Itest-nginx/lib -r t FLUSH_ETCD=1 TEST_EVENTS_MODULE=$TEST_EVENTS_MODULE prove --timer -Itest-nginx/lib -I./ -r $TEST_FILE_SUB_DIR | tee /tmp/test.result rerun_flaky_tests /tmp/test.result diff --git a/ci/redhat-ci.sh b/ci/redhat-ci.sh index 3cad10b5992b..da9839d4e699 100755 --- a/ci/redhat-ci.sh +++ b/ci/redhat-ci.sh @@ -77,6 +77,8 @@ install_dependencies() { yum install -y iproute procps start_grpc_server_example + start_sse_server_example + # installing grpcurl install_grpcurl diff --git a/t/APISIX.pm b/t/APISIX.pm index f11d6fd60dbf..2635b64aa4a5 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -845,6 +845,41 @@ _EOC_ } } + location \@disable_proxy_buffering { + + proxy_http_version 1.1; + proxy_set_header Host \$upstream_host; + proxy_set_header Upgrade \$upstream_upgrade; + proxy_set_header Connection \$upstream_connection; + proxy_set_header X-Real-IP \$remote_addr; + proxy_pass_header Date; + + ### the following x-forwarded-* headers is to send to upstream server + proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto \$var_x_forwarded_proto; + proxy_set_header X-Forwarded-Host \$var_x_forwarded_host; + proxy_set_header X-Forwarded-Port \$var_x_forwarded_port; + + proxy_pass \$upstream_scheme://apisix_backend\$upstream_uri; + + header_filter_by_lua_block { + apisix.http_header_filter_phase() + } + + body_filter_by_lua_block { + apisix.http_body_filter_phase() + } + + log_by_lua_block { + apisix.http_log_phase() + } + + proxy_buffering off; + access_by_lua_block { + apisix.disable_proxy_buffering_access_phase() + } + } + $grpc_location $dubbo_location diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 78e2e5ba4a90..82c89290ed07 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -598,3 +598,119 @@ Authorization: Bearer token --- error_code: 200 --- response_body {"foo", "bar"} + + + +=== TEST 16: set route with stream = true (SSE) +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "route_type": "llm/chat", + "auth": { + "type": "header", + "name": "Authorization", + "value": "Bearer token" + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0, + "stream": true + } + }, + "override": { + "host": "localhost", + "port": 7737, + "scheme": "http" + }, + "ssl_verify": false + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 17: test is SSE works as expected +--- LAST +--- config + location /t { + content_by_lua_block { + local http = require("resty.http") + local httpc = http.new() + local core = require("apisix.core") + + local ok, err = httpc:connect({ + scheme = "http", + host = "localhost", + port = ngx.var.server_port, + }) + + if not ok then + ngx.status = 500 + ngx.say(err) + return + end + + local params = { + method = "POST", + headers = { + ["Content-Type"] = "application/json", + }, + path = "/anything", + body = [[{ + "messages": [ + { "role": "system", "content": "some content" } + ] + }]], + } + + local res, err = httpc:request(params) + if not res then + ngx.status = 500 + ngx.say(err) + return + end + + local final_res = {} + while true do + local chunk, err = res.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 + core.table.insert_tail(final_res, chunk) + end + + ngx.print(#final_res .. final_res[6]) + } + } +--- response_body_like eval +qr/6data: \[DONE\]\n\n/ diff --git a/t/sse_server_example/go.mod b/t/sse_server_example/go.mod new file mode 100644 index 000000000000..9cc909d0338e --- /dev/null +++ b/t/sse_server_example/go.mod @@ -0,0 +1,3 @@ +module foo.bar/apache/sse_server_example + +go 1.17 diff --git a/t/sse_server_example/main.go b/t/sse_server_example/main.go new file mode 100644 index 000000000000..890cf1a6d700 --- /dev/null +++ b/t/sse_server_example/main.go @@ -0,0 +1,42 @@ +package main + +import ( + "fmt" + "log" + "net/http" + "os" + "time" +) + +func sseHandler(w http.ResponseWriter, r *http.Request) { + // Set the headers for SSE + w.Header().Set("Content-Type", "text/event-stream") + w.Header().Set("Cache-Control", "no-cache") + w.Header().Set("Connection", "keep-alive") + + // A simple loop that sends a message every 2 seconds + for i := 0; i < 5; i++ { + // Create a message to send to the client + fmt.Fprintf(w, "data: %s\n\n", time.Now().Format(time.RFC3339)) + + // Flush the data immediately to the client + if f, ok := w.(http.Flusher); ok { + f.Flush() + } else { + log.Println("Unable to flush data to client.") + break + } + + time.Sleep(500 * time.Millisecond) + } + fmt.Fprintf(w, "data: %s\n\n", "[DONE]") +} + +func main() { + // Create a simple route + http.HandleFunc("/v1/chat/completions", sseHandler) + port := os.Args[1] + // Start the server + log.Println("Starting server on :", port) + log.Fatal(http.ListenAndServe(":" + port, nil)) +} From 7c77ed61b209ef62d420eecd8876ce4738f0235e Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 11:51:26 +0545 Subject: [PATCH 55/85] cleanup --- apisix/constants.lua | 4 +--- apisix/plugins/ai-proxy/drivers/openai.lua | 7 +------ 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/apisix/constants.lua b/apisix/constants.lua index 83566686c37d..05223a29389f 100644 --- a/apisix/constants.lua +++ b/apisix/constants.lua @@ -42,7 +42,5 @@ return { ["/ssls"] = true, ["/stream_routes"] = true, ["/plugin_metadata"] = true, - }, - CHAT = "llm/chat", - COMPLETION = "llm/completions", + } } diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index a558bfcdc2a1..b03bd915ef36 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -18,17 +18,13 @@ local _M = {} local core = require("apisix.core") local http = require("resty.http") -local upstream = require("apisix.upstream") -local ngx = ngx + local pairs = pairs -- globals local DEFAULT_HOST = "api.openai.com" local DEFAULT_PORT = 443 -local path_mapper = { - ["llm/chat"] = "/v1/chat/completions", -} function _M.request(conf, request_table, ctx) local httpc, err = http.new() @@ -81,7 +77,6 @@ function _M.request(conf, request_table, ctx) return 500, "failed to send request to LLM server: " .. err end - -- TOOD: keepalive maintainance return res, nil, httpc end From fa46abe762cfaf195738710f0ccb4f9d9c7d7401 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 12:03:32 +0545 Subject: [PATCH 56/85] CLEANUP --- apisix/core/resolver.lua | 20 ----------- apisix/plugins/ai-proxy.lua | 5 +-- apisix/plugins/ai-proxy/schema.lua | 6 +--- apisix/upstream.lua | 56 +----------------------------- t/APISIX.pm | 1 - t/plugin/ai-proxy.t | 8 ----- t/sse_server_example/main.go | 17 +++++++++ 7 files changed, 22 insertions(+), 91 deletions(-) diff --git a/apisix/core/resolver.lua b/apisix/core/resolver.lua index 790d8945360a..3568a9762063 100644 --- a/apisix/core/resolver.lua +++ b/apisix/core/resolver.lua @@ -24,7 +24,6 @@ local log = require("apisix.core.log") local utils = require("apisix.core.utils") local dns_utils = require("resty.dns.utils") local config_local = require("apisix.core.config_local") -local ipmatcher = require("resty.ipmatcher") local HOSTS_IP_MATCH_CACHE = {} @@ -94,23 +93,4 @@ function _M.parse_domain(host) end -function _M.parse_domain_for_node(node) - local host = node.domain or node.host - if not ipmatcher.parse_ipv4(host) - and not ipmatcher.parse_ipv6(host) - then - node.domain = host - - local ip, err = _M.parse_domain(host) - if ip then - node.host = ip - end - - if err then - log.error("dns resolver domain: ", host, " error: ", err) - end - end -end - - return _M diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index d37d75e60fe6..0bdf800f2125 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -67,16 +67,17 @@ local function send_request(conf, ctx) local chunk, err = res.body_reader() -- will read chunk by chunk if err then core.log.error("failed to read response chunk: ", err) - return + break end if not chunk then - return + break end content_length = content_length + #chunk ngx.print(chunk) ngx.flush(true) end httpc:set_keepalive(10000, 100) + return else local res_body, err = res:read_body() if not res_body then diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 9e7c1c62f7f6..98ca26207ee2 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -124,10 +124,6 @@ local model_schema = { _M.plugin_schema = { type = "object", properties = { - route_type = { - type = "string", - enum = { "llm/chat", "passthrough" } - }, auth = auth_schema, model = model_schema, passthrough = { type = "boolean", default = false }, @@ -143,7 +139,7 @@ _M.plugin_schema = { keepalive_pool = {type = "integer", minimum = 1, default = 30}, ssl_verify = {type = "boolean", default = true }, }, - required = {"route_type", "model", "auth"} + required = {"model", "auth"} } _M.chat_request_schema = { diff --git a/apisix/upstream.lua b/apisix/upstream.lua index d3a7c34cd129..ce3738b0deb1 100644 --- a/apisix/upstream.lua +++ b/apisix/upstream.lua @@ -20,7 +20,7 @@ local discovery = require("apisix.discovery.init").discovery local upstream_util = require("apisix.utils.upstream") local apisix_ssl = require("apisix.ssl") local events = require("apisix.events") -local resolver = require("apisix.core.resolver") + local error = error local tostring = tostring local ipairs = ipairs @@ -254,60 +254,6 @@ local function fill_node_info(up_conf, scheme, is_stream) end -function _M.set_upstream(upstream_info, ctx) - local nodes = upstream_info.nodes - local new_nodes = {} - if core.table.isarray(nodes) then - for _, node in ipairs(nodes) do - resolver.parse_domain_for_node(node) - table_insert(new_nodes, node) - end - else - for addr, weight in pairs(nodes) do - local node = {} - local port, host - host, port = core.utils.parse_addr(addr) - node.host = host - resolver.parse_domain_for_node(node) - node.port = port - node.weight = weight - table_insert(new_nodes, node) - end - end - - local up_conf = { - name = upstream_info.name, - type = upstream_info.type, - hash_on = upstream_info.hash_on, - pass_host = upstream_info.pass_host, - upstream_host = upstream_info.upstream_host, - key = upstream_info.key, - nodes = new_nodes, - timeout = upstream_info.timeout, - scheme = upstream_info.scheme - } - - local ok, err = _M.check_schema(up_conf) - if not ok then - core.log.error("failed to validate generated upstream: ", err) - return 500, err - end - - local matched_route = ctx.matched_route - up_conf.parent = matched_route - local upstream_key = up_conf.type .. "#route_" .. - matched_route.value.id .. "_" .. upstream_info.vid - if upstream_info.node_tid then - upstream_key = upstream_key .. "_" .. upstream_info.node_tid - end - core.log.info("upstream_key: ", upstream_key) - _M.set(ctx, upstream_key, ctx.conf_version, up_conf) - if upstream_info.scheme == "https" then - _M.set_scheme(ctx, up_conf) - end -end - - function _M.set_by_route(route, api_ctx) if api_ctx.upstream_conf then -- upstream_conf has been set by traffic-split plugin diff --git a/t/APISIX.pm b/t/APISIX.pm index 2635b64aa4a5..f9e61ff5bd68 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -266,7 +266,6 @@ env APISIX_PROFILE; env PATH; # for searching external plugin runner's binary env TEST_NGINX_HTML_DIR; env OPENSSL_BIN; -env AI_PROXY_TEST_SCHEME; _EOC_ diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 82c89290ed07..2fe73245cca8 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -148,7 +148,6 @@ __DATA__ content_by_lua_block { local plugin = require("apisix.plugins.ai-proxy") local ok, err = plugin.check_schema({ - route_type = "llm/chat", model = { provider = "openai", name = "gpt-4", @@ -178,7 +177,6 @@ passed content_by_lua_block { local plugin = require("apisix.plugins.ai-proxy") local ok, err = plugin.check_schema({ - route_type = "llm/chat", model = { provider = "some-unique", name = "gpt-4", @@ -213,7 +211,6 @@ qr/.*provider: some-unique is not supported.*/ "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { "type": "header", "name": "Authorization", @@ -276,7 +273,6 @@ Unauthorized "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { "type": "header", "name": "Authorization", @@ -398,7 +394,6 @@ request format doesn't match schema: property "messages" is required "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { "type": "header", "name": "Authorization", @@ -472,7 +467,6 @@ options_works "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { "type": "header", "name": "Authorization", @@ -546,7 +540,6 @@ path override works "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { "type": "header", "name": "Authorization", @@ -612,7 +605,6 @@ Authorization: Bearer token "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { "type": "header", "name": "Authorization", diff --git a/t/sse_server_example/main.go b/t/sse_server_example/main.go index 890cf1a6d700..9a9c1688376e 100644 --- a/t/sse_server_example/main.go +++ b/t/sse_server_example/main.go @@ -1,3 +1,20 @@ +/* + * 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. + */ + package main import ( From 45e4f98006f3dd10457fc0e6d831adbf4795380e Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 16:40:54 +0545 Subject: [PATCH 57/85] cleanup --- apisix/core/request.lua | 5 ----- apisix/upstream.lua | 1 - 2 files changed, 6 deletions(-) diff --git a/apisix/core/request.lua b/apisix/core/request.lua index 801ce583ffd3..7d4842e6ccfc 100644 --- a/apisix/core/request.lua +++ b/apisix/core/request.lua @@ -341,11 +341,6 @@ function _M.get_request_body_table() return nil, { message = "could not get body: " .. (err or "request body is empty") } end - body, err = body:gsub("\\\"", "\"") -- remove escaping in JSON - if not body then - return nil, { message = "failed to remove escaping from body. err: " .. err} - end - local body_tab, err = json.decode(body) if not body_tab then return nil, { message = "could not get parse JSON request body: " .. err } diff --git a/apisix/upstream.lua b/apisix/upstream.lua index ce3738b0deb1..a0be2471bacd 100644 --- a/apisix/upstream.lua +++ b/apisix/upstream.lua @@ -26,7 +26,6 @@ local tostring = tostring local ipairs = ipairs local pairs = pairs local pcall = pcall -local table_insert = table.insert local ngx_var = ngx.var local is_http = ngx.config.subsystem == "http" local upstreams From 99af8677d53b0c296645f5393e3cf7eb286a4edf Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 16:41:34 +0545 Subject: [PATCH 58/85] use `endpoint` instead of several config fields --- apisix/plugins/ai-proxy/drivers/openai.lua | 20 +++++++++-------- apisix/plugins/ai-proxy/schema.lua | 13 +---------- t/plugin/ai-proxy.t | 25 ++++++---------------- 3 files changed, 18 insertions(+), 40 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index b03bd915ef36..8d7b1fd815cb 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -18,6 +18,7 @@ local _M = {} local core = require("apisix.core") local http = require("resty.http") +local url = require("socket.url") local pairs = pairs @@ -33,17 +34,18 @@ function _M.request(conf, request_table, ctx) end httpc:set_timeout(conf.timeout) - local custom_host = core.table.try_read_attr(conf, "override", "host") - local custom_port = core.table.try_read_attr(conf, "override", "port") - local custom_path = core.table.try_read_attr(conf, "override", "path") - local custom_scheme = core.table.try_read_attr(conf, "override", "scheme") + 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 = custom_scheme or "https", - host = custom_host or DEFAULT_HOST, - port = custom_port or DEFAULT_PORT, + 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 = custom_host or DEFAULT_HOST, + ssl_server_name = parsed_url.host or DEFAULT_HOST, pool_size = conf.keepalive and conf.keepalive_pool, }) @@ -58,7 +60,7 @@ function _M.request(conf, request_table, ctx) }, keepalive = conf.keepalive, ssl_verify = conf.ssl_verify, - path = custom_path or "/v1/chat/completions", + path = parsed_url.path or "/v1/chat/completions", } if conf.auth.type == "header" then diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 98ca26207ee2..290fb6495505 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -100,21 +100,10 @@ local model_schema = { override = { type = "object", properties = { - host = { + endpoint = { type = "string", description = "To be specified to override the host of the AI provider", }, - port = { - type = "integer", - description = "To be specified to override the AI provider port", - }, - path = { - type = "string", - description = "Overrieds the request path to the AI provider endpoints", - }, - scheme = { - type = "string" - } } } }, diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 2fe73245cca8..e6afd0e4ed79 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -225,9 +225,7 @@ qr/.*provider: some-unique is not supported.*/ } }, "override": { - "host": "localhost", - "port": 6724, - "scheme": "http" + "endpoint": "http://localhost:6724" }, "ssl_verify": false } @@ -287,9 +285,7 @@ Unauthorized } }, "override": { - "host": "localhost", - "port": 6724, - "scheme": "http" + "endpoint": "http://localhost:6724" }, "ssl_verify": false } @@ -408,9 +404,7 @@ request format doesn't match schema: property "messages" is required } }, "override": { - "host": "localhost", - "port": 6724, - "scheme": "http" + "endpoint": "http://localhost:6724" }, "ssl_verify": false } @@ -481,10 +475,7 @@ options_works } }, "override": { - "host": "localhost", - "port": 6724, - "path": "/random", - "scheme": "http" + "endpoint": "http://localhost:6724/random" }, "ssl_verify": false } @@ -554,9 +545,7 @@ path override works } }, "override": { - "host": "localhost", - "port": 6724, - "scheme": "http" + "endpoint": "http://localhost:6724" }, "ssl_verify": false, "passthrough": true @@ -620,9 +609,7 @@ Authorization: Bearer token } }, "override": { - "host": "localhost", - "port": 7737, - "scheme": "http" + "endpoint": "http://localhost:7737" }, "ssl_verify": false } From 11aef59517ebc0234937aae7a8640af48e9e7028 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 16:43:02 +0545 Subject: [PATCH 59/85] cleanup --- apisix/constants.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/constants.lua b/apisix/constants.lua index 05223a29389f..0b3ec160b53d 100644 --- a/apisix/constants.lua +++ b/apisix/constants.lua @@ -42,5 +42,5 @@ return { ["/ssls"] = true, ["/stream_routes"] = true, ["/plugin_metadata"] = true, - } + }, } From 4eb0ff428b136c9c7b703f05a1315b0f56ef139b Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 16:44:22 +0545 Subject: [PATCH 60/85] `delayed_access` -> `disable_proxy_buffering_access_phase` --- apisix/init.lua | 2 +- apisix/plugins/ai-proxy.lua | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index 736767a4e9b6..7143bb37af6a 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -737,7 +737,7 @@ function _M.disable_proxy_buffering_access_phase() local plugins = plugin.filter(api_ctx, api_ctx.matched_route) -- plugins to be run after proxy_buffering is disabled - plugin.run_plugin("delayed_access", plugins, api_ctx) + plugin.run_plugin("disable_proxy_buffering_access_phase", plugins, api_ctx) end diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 0bdf800f2125..6b089e2471ed 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -120,7 +120,7 @@ function _M.access(conf, ctx) end -function _M.delayed_access(conf, ctx) +function _M.disable_proxy_buffering_access_phase(conf, ctx) return send_request(conf, ctx) end From 92ebd9d05c8138abb97301131b5b5d29ee4a2e96 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 16:51:27 +0545 Subject: [PATCH 61/85] cleanup --- apisix/plugins/ai-proxy.lua | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 6b089e2471ed..0dbbdf46191d 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -110,7 +110,7 @@ function _M.access(conf, ctx) end ctx.request_table = request_table - if conf.model.options and conf.model.options.stream then + if core.table.try_read_attr(conf, "model", "options", "stream") then request_table.stream = true ngx.ctx.disable_proxy_buffering = true return From c0dc59ed94ed97bbf8214d46d2fefa9d576fdc7c Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 4 Sep 2024 23:22:26 +0545 Subject: [PATCH 62/85] scalable auth schema --- apisix/core/utils.lua | 15 +++++++ apisix/plugins/ai-proxy/drivers/openai.lua | 15 +++++-- apisix/plugins/ai-proxy/schema.lua | 28 +++++-------- t/core/utils.t | 22 ++++++++++ t/plugin/ai-proxy.t | 48 +++++++++++----------- 5 files changed, 83 insertions(+), 45 deletions(-) diff --git a/apisix/core/utils.lua b/apisix/core/utils.lua index cfea756542bd..ff4d715de609 100644 --- a/apisix/core/utils.lua +++ b/apisix/core/utils.lua @@ -215,6 +215,21 @@ function _M.uri_safe_encode(uri) end +function _M.table_to_query_params(tbl) + if not tbl or type(tbl) ~= "table" then + return "" + end + + local query_params = {} + for key, value in pairs(tbl) do + local encoded_key = _M.uri_safe_encode(tostring(key)) + local encoded_value = _M.uri_safe_encode(tostring(value)) + table.insert(query_params, encoded_key .. "=" .. encoded_value) + end + return table.concat(query_params, "&") +end + + function _M.validate_header_field(field) for i = 1, #field do local b = str_byte(field, i, i) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 8d7b1fd815cb..487160d213e6 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -53,14 +53,21 @@ function _M.request(conf, request_table, ctx) return nil, "failed to connect to LLM server: " .. err end + local query_params = core.utils.table_to_query_params(conf.auth.params) + local path = (parsed_url.path or "/v1/chat/completions") .. query_params + + for key, value in pairs(conf.auth.body or {}) do + request_table[key] = value + end + + local headers = (conf.auth.header or {}) + headers["Content-Type"] = "application/json" local params = { method = "POST", - headers = { - ["Content-Type"] = "application/json", - }, + headers = headers, keepalive = conf.keepalive, ssl_verify = conf.ssl_verify, - path = parsed_url.path or "/v1/chat/completions", + path = path, } if conf.auth.type == "header" then diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 290fb6495505..71ca1686091e 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -16,26 +16,20 @@ -- local _M = {} +local auth_item_schema = { + type = "object", + patternProperties = { + ["^[a-zA-Z0-9._-]+$"] = { + type = "string" + } + } +} + local auth_schema = { type = "object", - properties = { - type = { - type = "string", - enum = {"header", "param"} - }, - name = { - type = "string", - description = "Name of the param/header carrying Authorization or API key.", - minLength = 1, - }, - value = { - type = "string", - description = "Full auth-header/param value.", - minLength = 1, - -- TODO encrypted = true, - }, + patternProperties = { + ["^(query|body|header)$"] = auth_item_schema }, - required = { "type", "name", "value" }, additionalProperties = false, } diff --git a/t/core/utils.t b/t/core/utils.t index 9faa545e17e1..2e56151fe054 100644 --- a/t/core/utils.t +++ b/t/core/utils.t @@ -393,3 +393,25 @@ res:nil res:5 res:12 res:7 + + + +=== TEST 13: table_to_query_params +--- config + location /t { + content_by_lua_block { + local table_to_query_params = require("apisix.core.utils").table_to_query_params + local cases = { + {input = {}, out = ""}, + {input = {"a"}, out = "1=a", }, + {input = {"a", "b"}, out = "1=a&2=b", }, + {input = {a = "b", c = "d"}, out = "c=d&a=b", }, + } + for _, case in ipairs(cases) do + local got = table_to_query_params(case.input) + assert(got == case.out, string.format("got: %s but expected: %s", got, case.out)) + end + } + } +--- request +GET /t diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index e6afd0e4ed79..083fbd989d69 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -153,9 +153,9 @@ __DATA__ name = "gpt-4", }, auth = { - type = "header", - value = "some value", - name = "some name", + header = { + some_header = "some_value" + } } }) @@ -182,9 +182,9 @@ passed name = "gpt-4", }, auth = { - type = "header", - value = "some value", - name = "some name", + header = { + some_header = "some_value" + } } }) @@ -212,9 +212,9 @@ qr/.*provider: some-unique is not supported.*/ "plugins": { "ai-proxy": { "auth": { - "type": "header", - "name": "Authorization", - "value": "Bearer wrongtoken" + "header": { + "Authorization": "Bearer wrongtoken" + } }, "model": { "provider": "openai", @@ -272,9 +272,9 @@ Unauthorized "plugins": { "ai-proxy": { "auth": { - "type": "header", - "name": "Authorization", - "value": "Bearer token" + "header": { + "Authorization": "Bearer token" + } }, "model": { "provider": "openai", @@ -391,9 +391,9 @@ request format doesn't match schema: property "messages" is required "plugins": { "ai-proxy": { "auth": { - "type": "header", - "name": "Authorization", - "value": "Bearer token" + "header": { + "Authorization": "Bearer token" + } }, "model": { "provider": "openai", @@ -462,9 +462,9 @@ options_works "plugins": { "ai-proxy": { "auth": { - "type": "header", - "name": "Authorization", - "value": "Bearer token" + "header": { + "Authorization": "Bearer token" + } }, "model": { "provider": "openai", @@ -532,9 +532,9 @@ path override works "plugins": { "ai-proxy": { "auth": { - "type": "header", - "name": "Authorization", - "value": "Bearer token" + "header": { + "Authorization": "Bearer token" + } }, "model": { "provider": "openai", @@ -595,9 +595,9 @@ Authorization: Bearer token "plugins": { "ai-proxy": { "auth": { - "type": "header", - "name": "Authorization", - "value": "Bearer token" + "header": { + "Authorization": "Bearer token" + } }, "model": { "provider": "openai", From 99cf3a489ab717c17cd98a79f6f44b8ca8023474 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 09:37:25 +0545 Subject: [PATCH 63/85] fix lint --- apisix/plugins/ai-proxy.lua | 4 +++- t/sse_server_example/main.go | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 0dbbdf46191d..50718aa218a5 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -18,8 +18,10 @@ 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 ngx_ctx = ngx.ctx local plugin_name = "ai-proxy" local _M = { diff --git a/t/sse_server_example/main.go b/t/sse_server_example/main.go index 9a9c1688376e..47f79953014a 100644 --- a/t/sse_server_example/main.go +++ b/t/sse_server_example/main.go @@ -14,7 +14,7 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - + package main import ( From 243b5f5ae64010764f1846a9e9341947402beca7 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 09:50:10 +0545 Subject: [PATCH 64/85] remove body as auth param --- apisix/plugins/ai-proxy/drivers/openai.lua | 4 ---- apisix/plugins/ai-proxy/schema.lua | 2 +- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 487160d213e6..ec969fefebb1 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -56,10 +56,6 @@ function _M.request(conf, request_table, ctx) local query_params = core.utils.table_to_query_params(conf.auth.params) local path = (parsed_url.path or "/v1/chat/completions") .. query_params - for key, value in pairs(conf.auth.body or {}) do - request_table[key] = value - end - local headers = (conf.auth.header or {}) headers["Content-Type"] = "application/json" local params = { diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 71ca1686091e..720c5f2bf292 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -28,7 +28,7 @@ local auth_item_schema = { local auth_schema = { type = "object", patternProperties = { - ["^(query|body|header)$"] = auth_item_schema + ["^(query|header)$"] = auth_item_schema }, additionalProperties = false, } From 217b5afdaee16a16bf869a3dd129cc877c5190c4 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 09:50:30 +0545 Subject: [PATCH 65/85] query param auth test --- apisix/plugins/ai-proxy/drivers/openai.lua | 6 +- t/plugin/ai-proxy2.t | 200 +++++++++++++++++++++ 2 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 t/plugin/ai-proxy2.t diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index ec969fefebb1..2aa43eb19ce9 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -53,7 +53,11 @@ function _M.request(conf, request_table, ctx) return nil, "failed to connect to LLM server: " .. err end - local query_params = core.utils.table_to_query_params(conf.auth.params) + local query_params = core.utils.table_to_query_params(conf.auth.query) + if query_params and query_params ~= "" then + query_params = "?" .. query_params + end + local path = (parsed_url.path or "/v1/chat/completions") .. query_params local headers = (conf.auth.header or {}) diff --git a/t/plugin/ai-proxy2.t b/t/plugin/ai-proxy2.t new file mode 100644 index 000000000000..6e398e5665a4 --- /dev/null +++ b/t/plugin/ai-proxy2.t @@ -0,0 +1,200 @@ +# +# 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. +# + +use t::APISIX 'no_plan'; + +log_level("info"); +repeat_each(1); +no_long_string(); +no_root_location(); + + +my $resp_file = 't/assets/ai-proxy-response.json'; +open(my $fh, '<', $resp_file) or die "Could not open file '$resp_file' $!"; +my $resp = do { local $/; <$fh> }; +close($fh); + +print "Hello, World!\n"; +print $resp; + + +add_block_preprocessor(sub { + my ($block) = @_; + + if (!defined $block->request) { + $block->set_value("request", "GET /t"); + } + + my $http_config = $block->http_config // <<_EOC_; + server { + server_name openai; + listen 6724; + + default_type 'application/json'; + + location /v1/chat/completions { + content_by_lua_block { + local json = require("cjson.safe") + + if ngx.req.get_method() ~= "POST" then + ngx.status = 400 + ngx.say("Unsupported request method: ", ngx.req.get_method()) + end + ngx.req.read_body() + local body, err = ngx.req.get_body_data() + body, err = json.decode(body) + + local query_auth = ngx.req.get_uri_args()["api_key"] + + if query_auth ~= "apikey" then + ngx.status = 401 + ngx.say("Unauthorized") + return + end + + + ngx.status = 200 + ngx.say("passed") + } + } + } +_EOC_ + + $block->set_value("http_config", $http_config); +}); + +run_tests(); + +__DATA__ + +=== TEST 1: set route with wrong query param +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "auth": { + "query": { + "api_key": "wrong_key" + } + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0 + } + }, + "override": { + "endpoint": "http://localhost:6724" + }, + "ssl_verify": false + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 2: send request +--- request +POST /anything +{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } +--- error_code: 401 +--- response_body +Unauthorized + + + +=== TEST 3: set route with right query param +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/anything", + "plugins": { + "ai-proxy": { + "auth": { + "query": { + "api_key": "apikey" + } + }, + "model": { + "provider": "openai", + "name": "gpt-35-turbo-instruct", + "options": { + "max_tokens": 512, + "temperature": 1.0 + } + }, + "override": { + "endpoint": "http://localhost:6724" + }, + "ssl_verify": false + } + }, + "upstream": { + "type": "roundrobin", + "nodes": { + "canbeanything.com": 1 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 4: send request +--- request +POST /anything +{ "messages": [ { "role": "system", "content": "You are a mathematician" }, { "role": "user", "content": "What is 1+1?"} ] } +--- error_code: 200 +--- response_body +passed From 6661a0e06298e12c621dfbc489e71a8edc277945 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 11:17:10 +0545 Subject: [PATCH 66/85] fix test --- t/core/utils.t | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/t/core/utils.t b/t/core/utils.t index 2e56151fe054..ab6ced9601a5 100644 --- a/t/core/utils.t +++ b/t/core/utils.t @@ -405,7 +405,7 @@ res:7 {input = {}, out = ""}, {input = {"a"}, out = "1=a", }, {input = {"a", "b"}, out = "1=a&2=b", }, - {input = {a = "b", c = "d"}, out = "c=d&a=b", }, + {input = {a = "b", c = "d"}, out = "a=b&c=d", }, } for _, case in ipairs(cases) do local got = table_to_query_params(case.input) From 1c00e2c9d09e19d60be0a8b74a7d5ca006da2d87 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 11:18:05 +0545 Subject: [PATCH 67/85] fix lint --- apisix/plugins/ai-proxy.lua | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 50718aa218a5..340cc738564e 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -75,8 +75,8 @@ local function send_request(conf, ctx) break end content_length = content_length + #chunk - ngx.print(chunk) - ngx.flush(true) + ngx_print(chunk) + ngx_flush(true) end httpc:set_keepalive(10000, 100) return @@ -114,7 +114,7 @@ function _M.access(conf, ctx) if core.table.try_read_attr(conf, "model", "options", "stream") then request_table.stream = true - ngx.ctx.disable_proxy_buffering = true + ngx_ctx.disable_proxy_buffering = true return end From 37cdd7b0b1537842fbdf94babfba059e7041730d Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 13:59:18 +0545 Subject: [PATCH 68/85] remove forgotten header appender part --- apisix/plugins/ai-proxy/drivers/openai.lua | 4 ---- 1 file changed, 4 deletions(-) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 2aa43eb19ce9..1441c5611602 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -70,10 +70,6 @@ function _M.request(conf, request_table, ctx) path = path, } - if conf.auth.type == "header" then - params.headers[conf.auth.name] = conf.auth.value - end - if conf.model.options then for opt, val in pairs(conf.model.options) do request_table[opt] = val From 1062cc2d0700282926d81b4f49e0df708bb0b909 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 14:43:31 +0545 Subject: [PATCH 69/85] =?UTF-8?q?remove=20subrequest=20=F0=9F=98=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apisix/cli/ngx_tpl.lua | 58 ---------------------------------- apisix/init.lua | 14 --------- apisix/plugins/ai-proxy.lua | 62 ++++++++++++++----------------------- t/APISIX.pm | 35 --------------------- 4 files changed, 24 insertions(+), 145 deletions(-) diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index ffc8ac55a915..4b7ff4102bc1 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -809,64 +809,6 @@ http { } } - location @disable_proxy_buffering { - # http server location configuration snippet starts - {% if http_server_location_configuration_snippet then %} - {* http_server_location_configuration_snippet *} - {% end %} - # http server location configuration snippet ends - - proxy_http_version 1.1; - proxy_set_header Host $upstream_host; - proxy_set_header Upgrade $upstream_upgrade; - proxy_set_header Connection $upstream_connection; - proxy_set_header X-Real-IP $remote_addr; - proxy_pass_header Date; - - ### the following x-forwarded-* headers is to send to upstream server - proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto $var_x_forwarded_proto; - proxy_set_header X-Forwarded-Host $var_x_forwarded_host; - proxy_set_header X-Forwarded-Port $var_x_forwarded_port; - - {% if enabled_plugins["proxy-cache"] then %} - ### the following configuration is to cache response content from upstream server - proxy_cache $upstream_cache_zone; - proxy_cache_valid any {% if proxy_cache.cache_ttl then %} {* proxy_cache.cache_ttl *} {% else %} 10s {% end %}; - proxy_cache_min_uses 1; - proxy_cache_methods GET HEAD POST; - proxy_cache_lock_timeout 5s; - proxy_cache_use_stale off; - proxy_cache_key $upstream_cache_key; - proxy_no_cache $upstream_no_cache; - proxy_cache_bypass $upstream_cache_bypass; - - {% end %} - - proxy_pass $upstream_scheme://apisix_backend$upstream_uri; - - {% if enabled_plugins["proxy-mirror"] then %} - mirror /proxy_mirror; - {% end %} - - header_filter_by_lua_block { - apisix.http_header_filter_phase() - } - - body_filter_by_lua_block { - apisix.http_body_filter_phase() - } - - log_by_lua_block { - apisix.http_log_phase() - } - - proxy_buffering off; - access_by_lua_block { - apisix.disable_proxy_buffering_access_phase() - } - } - location @grpc_pass { access_by_lua_block { diff --git a/apisix/init.lua b/apisix/init.lua index 7143bb37af6a..404b75358b18 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -722,24 +722,10 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - if ngx.ctx.disable_proxy_buffering then - stash_ngx_ctx() - return ngx.exec("@disable_proxy_buffering") - end - _M.handle_upstream(api_ctx, route, enable_websocket) end -function _M.disable_proxy_buffering_access_phase() - ngx.ctx = fetch_ctx() - local api_ctx = ngx.ctx.api_ctx - local plugins = plugin.filter(api_ctx, api_ctx.matched_route) - - -- plugins to be run after proxy_buffering is disabled - plugin.run_plugin("disable_proxy_buffering_access_phase", plugins, api_ctx) -end - function _M.dubbo_access_phase() ngx.ctx = fetch_ctx() diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 340cc738564e..f88ee139afe9 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -44,8 +44,30 @@ end local CONTENT_TYPE_JSON = "application/json" -local function send_request(conf, ctx) - local request_table = ctx.request_table +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() + 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) local res, err, httpc = ai_driver.request(conf, request_table, ctx) if not res then @@ -90,40 +112,4 @@ local function send_request(conf, ctx) end 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() - 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 - ctx.request_table = request_table - - if core.table.try_read_attr(conf, "model", "options", "stream") then - request_table.stream = true - ngx_ctx.disable_proxy_buffering = true - return - end - - return send_request(conf, ctx) -end - - -function _M.disable_proxy_buffering_access_phase(conf, ctx) - return send_request(conf, ctx) -end - return _M diff --git a/t/APISIX.pm b/t/APISIX.pm index f9e61ff5bd68..50f7cfaecab6 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -844,41 +844,6 @@ _EOC_ } } - location \@disable_proxy_buffering { - - proxy_http_version 1.1; - proxy_set_header Host \$upstream_host; - proxy_set_header Upgrade \$upstream_upgrade; - proxy_set_header Connection \$upstream_connection; - proxy_set_header X-Real-IP \$remote_addr; - proxy_pass_header Date; - - ### the following x-forwarded-* headers is to send to upstream server - proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto \$var_x_forwarded_proto; - proxy_set_header X-Forwarded-Host \$var_x_forwarded_host; - proxy_set_header X-Forwarded-Port \$var_x_forwarded_port; - - proxy_pass \$upstream_scheme://apisix_backend\$upstream_uri; - - header_filter_by_lua_block { - apisix.http_header_filter_phase() - } - - body_filter_by_lua_block { - apisix.http_body_filter_phase() - } - - log_by_lua_block { - apisix.http_log_phase() - } - - proxy_buffering off; - access_by_lua_block { - apisix.disable_proxy_buffering_access_phase() - } - } - $grpc_location $dubbo_location From 7b236238ceeac0a302f709584302cc0393ff2742 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 14:52:04 +0545 Subject: [PATCH 70/85] =?UTF-8?q?Revert=20"remove=20subrequest=20?= =?UTF-8?q?=F0=9F=98=A9"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 1062cc2d0700282926d81b4f49e0df708bb0b909. --- apisix/cli/ngx_tpl.lua | 58 ++++++++++++++++++++++++++++++++++ apisix/init.lua | 14 +++++++++ apisix/plugins/ai-proxy.lua | 62 +++++++++++++++++++++++-------------- t/APISIX.pm | 35 +++++++++++++++++++++ 4 files changed, 145 insertions(+), 24 deletions(-) diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index 4b7ff4102bc1..ffc8ac55a915 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -809,6 +809,64 @@ http { } } + location @disable_proxy_buffering { + # http server location configuration snippet starts + {% if http_server_location_configuration_snippet then %} + {* http_server_location_configuration_snippet *} + {% end %} + # http server location configuration snippet ends + + proxy_http_version 1.1; + proxy_set_header Host $upstream_host; + proxy_set_header Upgrade $upstream_upgrade; + proxy_set_header Connection $upstream_connection; + proxy_set_header X-Real-IP $remote_addr; + proxy_pass_header Date; + + ### the following x-forwarded-* headers is to send to upstream server + proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto $var_x_forwarded_proto; + proxy_set_header X-Forwarded-Host $var_x_forwarded_host; + proxy_set_header X-Forwarded-Port $var_x_forwarded_port; + + {% if enabled_plugins["proxy-cache"] then %} + ### the following configuration is to cache response content from upstream server + proxy_cache $upstream_cache_zone; + proxy_cache_valid any {% if proxy_cache.cache_ttl then %} {* proxy_cache.cache_ttl *} {% else %} 10s {% end %}; + proxy_cache_min_uses 1; + proxy_cache_methods GET HEAD POST; + proxy_cache_lock_timeout 5s; + proxy_cache_use_stale off; + proxy_cache_key $upstream_cache_key; + proxy_no_cache $upstream_no_cache; + proxy_cache_bypass $upstream_cache_bypass; + + {% end %} + + proxy_pass $upstream_scheme://apisix_backend$upstream_uri; + + {% if enabled_plugins["proxy-mirror"] then %} + mirror /proxy_mirror; + {% end %} + + header_filter_by_lua_block { + apisix.http_header_filter_phase() + } + + body_filter_by_lua_block { + apisix.http_body_filter_phase() + } + + log_by_lua_block { + apisix.http_log_phase() + } + + proxy_buffering off; + access_by_lua_block { + apisix.disable_proxy_buffering_access_phase() + } + } + location @grpc_pass { access_by_lua_block { diff --git a/apisix/init.lua b/apisix/init.lua index 404b75358b18..7143bb37af6a 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -722,10 +722,24 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end + if ngx.ctx.disable_proxy_buffering then + stash_ngx_ctx() + return ngx.exec("@disable_proxy_buffering") + end + _M.handle_upstream(api_ctx, route, enable_websocket) end +function _M.disable_proxy_buffering_access_phase() + ngx.ctx = fetch_ctx() + local api_ctx = ngx.ctx.api_ctx + local plugins = plugin.filter(api_ctx, api_ctx.matched_route) + + -- plugins to be run after proxy_buffering is disabled + plugin.run_plugin("disable_proxy_buffering_access_phase", plugins, api_ctx) +end + function _M.dubbo_access_phase() ngx.ctx = fetch_ctx() diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index f88ee139afe9..340cc738564e 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -44,30 +44,8 @@ end local CONTENT_TYPE_JSON = "application/json" -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() - 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 function send_request(conf, ctx) + local request_table = ctx.request_table local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) local res, err, httpc = ai_driver.request(conf, request_table, ctx) if not res then @@ -112,4 +90,40 @@ function _M.access(conf, ctx) end 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() + 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 + ctx.request_table = request_table + + if core.table.try_read_attr(conf, "model", "options", "stream") then + request_table.stream = true + ngx_ctx.disable_proxy_buffering = true + return + end + + return send_request(conf, ctx) +end + + +function _M.disable_proxy_buffering_access_phase(conf, ctx) + return send_request(conf, ctx) +end + return _M diff --git a/t/APISIX.pm b/t/APISIX.pm index 50f7cfaecab6..f9e61ff5bd68 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -844,6 +844,41 @@ _EOC_ } } + location \@disable_proxy_buffering { + + proxy_http_version 1.1; + proxy_set_header Host \$upstream_host; + proxy_set_header Upgrade \$upstream_upgrade; + proxy_set_header Connection \$upstream_connection; + proxy_set_header X-Real-IP \$remote_addr; + proxy_pass_header Date; + + ### the following x-forwarded-* headers is to send to upstream server + proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; + proxy_set_header X-Forwarded-Proto \$var_x_forwarded_proto; + proxy_set_header X-Forwarded-Host \$var_x_forwarded_host; + proxy_set_header X-Forwarded-Port \$var_x_forwarded_port; + + proxy_pass \$upstream_scheme://apisix_backend\$upstream_uri; + + header_filter_by_lua_block { + apisix.http_header_filter_phase() + } + + body_filter_by_lua_block { + apisix.http_body_filter_phase() + } + + log_by_lua_block { + apisix.http_log_phase() + } + + proxy_buffering off; + access_by_lua_block { + apisix.disable_proxy_buffering_access_phase() + } + } + $grpc_location $dubbo_location From 4278fd53bff4fd802860e5f2ea86d0246397678c Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 14:55:22 +0545 Subject: [PATCH 71/85] use encode_args instead of custom func --- apisix/core/utils.lua | 15 --------------- apisix/plugins/ai-proxy/drivers/openai.lua | 2 +- t/core/utils.t | 22 ---------------------- 3 files changed, 1 insertion(+), 38 deletions(-) diff --git a/apisix/core/utils.lua b/apisix/core/utils.lua index ff4d715de609..cfea756542bd 100644 --- a/apisix/core/utils.lua +++ b/apisix/core/utils.lua @@ -215,21 +215,6 @@ function _M.uri_safe_encode(uri) end -function _M.table_to_query_params(tbl) - if not tbl or type(tbl) ~= "table" then - return "" - end - - local query_params = {} - for key, value in pairs(tbl) do - local encoded_key = _M.uri_safe_encode(tostring(key)) - local encoded_value = _M.uri_safe_encode(tostring(value)) - table.insert(query_params, encoded_key .. "=" .. encoded_value) - end - return table.concat(query_params, "&") -end - - function _M.validate_header_field(field) for i = 1, #field do local b = str_byte(field, i, i) diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 1441c5611602..5509b46e6cea 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -53,7 +53,7 @@ function _M.request(conf, request_table, ctx) return nil, "failed to connect to LLM server: " .. err end - local query_params = core.utils.table_to_query_params(conf.auth.query) + local query_params = core.string.encode_args(conf.auth.query) if query_params and query_params ~= "" then query_params = "?" .. query_params end diff --git a/t/core/utils.t b/t/core/utils.t index ab6ced9601a5..9faa545e17e1 100644 --- a/t/core/utils.t +++ b/t/core/utils.t @@ -393,25 +393,3 @@ res:nil res:5 res:12 res:7 - - - -=== TEST 13: table_to_query_params ---- config - location /t { - content_by_lua_block { - local table_to_query_params = require("apisix.core.utils").table_to_query_params - local cases = { - {input = {}, out = ""}, - {input = {"a"}, out = "1=a", }, - {input = {"a", "b"}, out = "1=a&2=b", }, - {input = {a = "b", c = "d"}, out = "a=b&c=d", }, - } - for _, case in ipairs(cases) do - local got = table_to_query_params(case.input) - assert(got == case.out, string.format("got: %s but expected: %s", got, case.out)) - end - } - } ---- request -GET /t From d915292b99f5b8326f48117cdc3863929bd028f6 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 15:13:17 +0545 Subject: [PATCH 72/85] cleanup --- apisix/init.lua | 2 +- apisix/plugins/ai-proxy.lua | 3 +-- apisix/plugins/ai-proxy/drivers/openai.lua | 9 ++++++--- apisix/plugins/ai-proxy/schema.lua | 3 ++- 4 files changed, 10 insertions(+), 7 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index 7143bb37af6a..c0539cf8f16b 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -722,7 +722,7 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - if ngx.ctx.disable_proxy_buffering then + if api_ctx.disable_proxy_buffering then stash_ngx_ctx() return ngx.exec("@disable_proxy_buffering") end diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 340cc738564e..d423f7609761 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -21,7 +21,6 @@ local pcall = pcall local ngx_req = ngx.req local ngx_print = ngx.print local ngx_flush = ngx.flush -local ngx_ctx = ngx.ctx local plugin_name = "ai-proxy" local _M = { @@ -114,7 +113,7 @@ function _M.access(conf, ctx) if core.table.try_read_attr(conf, "model", "options", "stream") then request_table.stream = true - ngx_ctx.disable_proxy_buffering = true + ctx.disable_proxy_buffering = true return end diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 5509b46e6cea..46f9cb0ca5eb 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -53,9 +53,12 @@ function _M.request(conf, request_table, ctx) return nil, "failed to connect to LLM server: " .. err end - local query_params = core.string.encode_args(conf.auth.query) - if query_params and query_params ~= "" then - query_params = "?" .. query_params + local query_params = "" + if conf.auth.query and type(conf.auth.query) == "table" then + query_params = core.string.encode_args(conf.auth.query) + if query_params and query_params ~= "" then + query_params = "?" .. query_params + end end local path = (parsed_url.path or "/v1/chat/completions") .. query_params diff --git a/apisix/plugins/ai-proxy/schema.lua b/apisix/plugins/ai-proxy/schema.lua index 720c5f2bf292..382644dc2147 100644 --- a/apisix/plugins/ai-proxy/schema.lua +++ b/apisix/plugins/ai-proxy/schema.lua @@ -28,7 +28,8 @@ local auth_item_schema = { local auth_schema = { type = "object", patternProperties = { - ["^(query|header)$"] = auth_item_schema + header = auth_item_schema, + query = auth_item_schema, }, additionalProperties = false, } From 8c3bcb37bc3d69e7fe8f055108a758acac8621c6 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 15:20:53 +0545 Subject: [PATCH 73/85] clean upstream.lua --- apisix/upstream.lua | 1 - 1 file changed, 1 deletion(-) diff --git a/apisix/upstream.lua b/apisix/upstream.lua index a0be2471bacd..eb5e467daaca 100644 --- a/apisix/upstream.lua +++ b/apisix/upstream.lua @@ -20,7 +20,6 @@ local discovery = require("apisix.discovery.init").discovery local upstream_util = require("apisix.utils.upstream") local apisix_ssl = require("apisix.ssl") local events = require("apisix.events") - local error = error local tostring = tostring local ipairs = ipairs From bcca3f28d875fa26ef041c51a46e30a381f7c70d Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 5 Sep 2024 20:13:52 +0545 Subject: [PATCH 74/85] optimize --- apisix/plugins/ai-proxy.lua | 22 ++++++++++++++++++---- apisix/plugins/ai-proxy/drivers/openai.lua | 1 + 2 files changed, 19 insertions(+), 4 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index d423f7609761..40abedfd000a 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -43,6 +43,15 @@ end local CONTENT_TYPE_JSON = "application/json" +local function keepalive_or_close(conf, httpc) + if conf.set_keepalive then + httpc:set_keepalive(10000, 100) + return + end + httpc:close() +end + + local function send_request(conf, ctx) local request_table = ctx.request_table local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) @@ -62,10 +71,15 @@ local function send_request(conf, ctx) return end + local body_reader = res.body_reader + if not body_reader then + core.log.error("LLM sent no response body") + return 500 + end + if core.table.try_read_attr(conf, "model", "options", "stream") then - local content_length = 0 while true do - local chunk, err = res.body_reader() -- will read chunk by chunk + local chunk, err = body_reader() -- will read chunk by chunk if err then core.log.error("failed to read response chunk: ", err) break @@ -73,11 +87,10 @@ local function send_request(conf, ctx) if not chunk then break end - content_length = content_length + #chunk ngx_print(chunk) ngx_flush(true) end - httpc:set_keepalive(10000, 100) + keepalive_or_close(conf, httpc) return else local res_body, err = res:read_body() @@ -85,6 +98,7 @@ local function send_request(conf, ctx) core.log.error("failed to read response body: ", err) return 500 end + keepalive_or_close(conf, httpc) return res.status, res_body end end diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 46f9cb0ca5eb..6912d54972cc 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -21,6 +21,7 @@ local http = require("resty.http") local url = require("socket.url") local pairs = pairs +local type = type -- globals local DEFAULT_HOST = "api.openai.com" From 929cbb1cea02de290789e762fb2e9d067e7ce7b8 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 6 Sep 2024 10:19:29 +0545 Subject: [PATCH 75/85] fix lint --- t/plugin/ai-proxy.t | 1 - 1 file changed, 1 deletion(-) diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 083fbd989d69..85f3ee64ec1d 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -635,7 +635,6 @@ passed === TEST 17: test is SSE works as expected ---- LAST --- config location /t { content_by_lua_block { From e149170a97adcc34cf3b08dcc3335371a23454c6 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 6 Sep 2024 10:57:24 +0545 Subject: [PATCH 76/85] pass data to upstream in streaming way --- apisix/plugins/ai-proxy.lua | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 40abedfd000a..6b71de1dae7f 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -61,22 +61,30 @@ local function send_request(conf, ctx) return 500 end - if conf.passthrough then - local res_body, err = res:read_body() - if not res_body then - core.log.error("failed to read response body: ", err) - return 500 - end - ngx_req.set_body_data(res_body) - return - end - local body_reader = res.body_reader if not body_reader then core.log.error("LLM sent no response body") return 500 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 while true do local chunk, err = body_reader() -- will read chunk by chunk From d346d3834536c15c68e6229ed200c64756ba74a8 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Mon, 9 Sep 2024 11:24:36 +0545 Subject: [PATCH 77/85] =?UTF-8?q?Reapply=20"remove=20subrequest=20?= =?UTF-8?q?=F0=9F=98=A9"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 7b236238ceeac0a302f709584302cc0393ff2742. --- apisix/cli/ngx_tpl.lua | 58 ---------------------------------- apisix/init.lua | 14 --------- apisix/plugins/ai-proxy.lua | 62 ++++++++++++++----------------------- t/APISIX.pm | 35 --------------------- 4 files changed, 24 insertions(+), 145 deletions(-) diff --git a/apisix/cli/ngx_tpl.lua b/apisix/cli/ngx_tpl.lua index ffc8ac55a915..4b7ff4102bc1 100644 --- a/apisix/cli/ngx_tpl.lua +++ b/apisix/cli/ngx_tpl.lua @@ -809,64 +809,6 @@ http { } } - location @disable_proxy_buffering { - # http server location configuration snippet starts - {% if http_server_location_configuration_snippet then %} - {* http_server_location_configuration_snippet *} - {% end %} - # http server location configuration snippet ends - - proxy_http_version 1.1; - proxy_set_header Host $upstream_host; - proxy_set_header Upgrade $upstream_upgrade; - proxy_set_header Connection $upstream_connection; - proxy_set_header X-Real-IP $remote_addr; - proxy_pass_header Date; - - ### the following x-forwarded-* headers is to send to upstream server - proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto $var_x_forwarded_proto; - proxy_set_header X-Forwarded-Host $var_x_forwarded_host; - proxy_set_header X-Forwarded-Port $var_x_forwarded_port; - - {% if enabled_plugins["proxy-cache"] then %} - ### the following configuration is to cache response content from upstream server - proxy_cache $upstream_cache_zone; - proxy_cache_valid any {% if proxy_cache.cache_ttl then %} {* proxy_cache.cache_ttl *} {% else %} 10s {% end %}; - proxy_cache_min_uses 1; - proxy_cache_methods GET HEAD POST; - proxy_cache_lock_timeout 5s; - proxy_cache_use_stale off; - proxy_cache_key $upstream_cache_key; - proxy_no_cache $upstream_no_cache; - proxy_cache_bypass $upstream_cache_bypass; - - {% end %} - - proxy_pass $upstream_scheme://apisix_backend$upstream_uri; - - {% if enabled_plugins["proxy-mirror"] then %} - mirror /proxy_mirror; - {% end %} - - header_filter_by_lua_block { - apisix.http_header_filter_phase() - } - - body_filter_by_lua_block { - apisix.http_body_filter_phase() - } - - log_by_lua_block { - apisix.http_log_phase() - } - - proxy_buffering off; - access_by_lua_block { - apisix.disable_proxy_buffering_access_phase() - } - } - location @grpc_pass { access_by_lua_block { diff --git a/apisix/init.lua b/apisix/init.lua index c0539cf8f16b..404b75358b18 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -722,24 +722,10 @@ function _M.http_access_phase() plugin.run_plugin("access", plugins, api_ctx) end - if api_ctx.disable_proxy_buffering then - stash_ngx_ctx() - return ngx.exec("@disable_proxy_buffering") - end - _M.handle_upstream(api_ctx, route, enable_websocket) end -function _M.disable_proxy_buffering_access_phase() - ngx.ctx = fetch_ctx() - local api_ctx = ngx.ctx.api_ctx - local plugins = plugin.filter(api_ctx, api_ctx.matched_route) - - -- plugins to be run after proxy_buffering is disabled - plugin.run_plugin("disable_proxy_buffering_access_phase", plugins, api_ctx) -end - function _M.dubbo_access_phase() ngx.ctx = fetch_ctx() diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 6b71de1dae7f..f8514194798e 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -52,8 +52,30 @@ local function keepalive_or_close(conf, httpc) end -local function send_request(conf, ctx) - local request_table = ctx.request_table +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() + 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) local res, err, httpc = ai_driver.request(conf, request_table, ctx) if not res then @@ -111,40 +133,4 @@ local function send_request(conf, ctx) end 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() - 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 - ctx.request_table = request_table - - if core.table.try_read_attr(conf, "model", "options", "stream") then - request_table.stream = true - ctx.disable_proxy_buffering = true - return - end - - return send_request(conf, ctx) -end - - -function _M.disable_proxy_buffering_access_phase(conf, ctx) - return send_request(conf, ctx) -end - return _M diff --git a/t/APISIX.pm b/t/APISIX.pm index f9e61ff5bd68..50f7cfaecab6 100644 --- a/t/APISIX.pm +++ b/t/APISIX.pm @@ -844,41 +844,6 @@ _EOC_ } } - location \@disable_proxy_buffering { - - proxy_http_version 1.1; - proxy_set_header Host \$upstream_host; - proxy_set_header Upgrade \$upstream_upgrade; - proxy_set_header Connection \$upstream_connection; - proxy_set_header X-Real-IP \$remote_addr; - proxy_pass_header Date; - - ### the following x-forwarded-* headers is to send to upstream server - proxy_set_header X-Forwarded-For \$proxy_add_x_forwarded_for; - proxy_set_header X-Forwarded-Proto \$var_x_forwarded_proto; - proxy_set_header X-Forwarded-Host \$var_x_forwarded_host; - proxy_set_header X-Forwarded-Port \$var_x_forwarded_port; - - proxy_pass \$upstream_scheme://apisix_backend\$upstream_uri; - - header_filter_by_lua_block { - apisix.http_header_filter_phase() - } - - body_filter_by_lua_block { - apisix.http_body_filter_phase() - } - - log_by_lua_block { - apisix.http_log_phase() - } - - proxy_buffering off; - access_by_lua_block { - apisix.disable_proxy_buffering_access_phase() - } - } - $grpc_location $dubbo_location From 3429e03e4aa166aef9377902730f80be8b8e4f66 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Mon, 9 Sep 2024 16:26:21 +0545 Subject: [PATCH 78/85] update doc --- docs/en/latest/plugins/ai-proxy.md | 57 +++++++++++++++--------------- 1 file changed, 28 insertions(+), 29 deletions(-) diff --git a/docs/en/latest/plugins/ai-proxy.md b/docs/en/latest/plugins/ai-proxy.md index 8813882e03be..f27e6af1dc98 100644 --- a/docs/en/latest/plugins/ai-proxy.md +++ b/docs/en/latest/plugins/ai-proxy.md @@ -46,32 +46,29 @@ Proxying requests to OpenAI is supported for now, other AI models will be suppor | `messages.role` | String | Yes | Role of the message (`system`, `user`, `assistant`) | | `messages.content` | String | Yes | Content of the message | -- Completion API - -| Name | Type | Required | Description | -| -------- | ------ | -------- | --------------------------------- | -| `prompt` | String | Yes | Prompt to be sent to the upstream | - ## Plugin Attributes -| Field | Type | Description | Required | -| ---------------------------------- | ------- | --------------------------------------------------------------------------------------------- | -------- | -| `route_type` | String | Specifies the type of route (`llm/chat`, `llm/completions`, `passthrough`) | Yes | -| `auth` | Object | Authentication configuration | Yes | -| `auth.source` | String | Source of the authentication (`header`, `param`) | Yes | -| `auth.name` | String | Name of the param/header carrying Authorization or API key. Minimum length: 1 | Yes | -| `auth.value` | String | Full auth-header/param value. Minimum length: 1. Encrypted. | Yes | -| `model` | Object | Model configuration | Yes | -| `model.provider` | String | AI provider request format. Translates requests to/from specified backend compatible formats. | Yes | -| `model.name` | String | Model name to execute. | Yes | -| `model.options` | Object | Key/value settings for the model | No | -| `model.options.max_tokens` | Integer | Defines the max_tokens, if using chat or completion models. Default: 256 | No | -| `model.options.temperature` | Number | Defines the matching temperature, if using chat or completion models. Range: 0.0 - 5.0 | No | -| `model.options.top_p` | Number | Defines the top-p probability mass, if supported. Range: 0.0 - 1.0 | No | -| `model.options.upstream_host` | String | To be specified to override the host of the AI provider | No | -| `model.options.upstream_port` | Integer | To be specified to override the AI provider port | No | -| `model.options.upstream_path` | String | To be specified to override the URL to the AI provider endpoints | No | -| `model.options.response_streaming` | Boolean | Stream response by SSE. Default: false | No | +| **Field** | **Required** | **Type** | **Description** | +| ------------------------- | ------------ | -------- | ------------------------------------------------------------------------------------ | +| auth | Yes | Object | Authentication configuration | +| auth.header | No | Object | Authentication headers. Key must match pattern `^[a-zA-Z0-9._-]+$`. | +| auth.query | No | Object | Authentication query parameters. Key must match pattern `^[a-zA-Z0-9._-]+$`. | +| model.provider | Yes | String | Name of the AI service provider (`openai`). | +| model.name | Yes | String | Model name to execute. | +| model.options | No | Object | Key/value settings for the model | +| model.options.max_tokens | No | Integer | Defines the max tokens if using chat or completion models. Default: 256 | +| model.options.input_cost | No | Number | Cost per 1M tokens in your prompt. Minimum: 0 | +| model.options.output_cost | No | Number | Cost per 1M tokens in the output of the AI. Minimum: 0 | +| model.options.temperature | No | Number | Matching temperature for models. Range: 0.0 - 5.0 | +| model.options.top_p | No | Number | Top-p probability mass. Range: 0 - 1 | +| model.options.stream | No | Boolean | Stream response by SSE. Default: false | +| model.override.endpoint | No | String | Override the endpoint of the AI provider | +| passthrough | No | Boolean | If enabled, the response from LLM will be sent to the upstream. Default: false | +| timeout | No | Integer | Timeout in milliseconds for requests to LLM. Range: 1 - 60000. Default: 3000 | +| keepalive | No | Boolean | Enable keepalive for requests to LLM. Default: true | +| keepalive_timeout | No | Integer | Keepalive timeout in milliseconds for requests to LLM. Minimum: 1000. Default: 60000 | +| keepalive_pool | No | Integer | Keepalive pool size for requests to LLM. Minimum: 1. Default: 30 | +| ssl_verify | No | Boolean | SSL verification for requests to LLM. Default: true | ## Example usage @@ -84,10 +81,10 @@ curl "http://127.0.0.1:9180/apisix/admin/routes/1" -X PUT \ "uri": "/anything", "plugins": { "ai-proxy": { - "route_type": "llm/chat", "auth": { - "header_name": "Authorization", - "header_value": "Bearer " + "header": { + "Authorization": "Bearer " + } }, "model": { "provider": "openai", @@ -103,12 +100,14 @@ curl "http://127.0.0.1:9180/apisix/admin/routes/1" -X PUT \ "type": "roundrobin", "nodes": { "somerandom.com:443": 1 - } + }, + "scheme": "https", + "pass_host": "node" } }' ``` -The upstream node can be any arbitrary value because it won't be contacted. +Since `passthrough` is not enabled upstream node can be any arbitrary value because it won't be contacted. Now send a request: From 7c082904fb603c901b82fe1a9235d3219cab44a9 Mon Sep 17 00:00:00 2001 From: Shreemaan Abhishek Date: Tue, 10 Sep 2024 11:26:02 +0545 Subject: [PATCH 79/85] Apply suggestions from code review Co-authored-by: Traky Deng --- docs/en/latest/plugins/ai-proxy.md | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/en/latest/plugins/ai-proxy.md b/docs/en/latest/plugins/ai-proxy.md index f27e6af1dc98..a6a4e35426eb 100644 --- a/docs/en/latest/plugins/ai-proxy.md +++ b/docs/en/latest/plugins/ai-proxy.md @@ -29,10 +29,10 @@ description: This document contains information about the Apache APISIX ai-proxy ## Description -The `ai-proxy` plugin simplifies access to AI providers and models by defining a standard request format -that allows configuring key fields in plugin configuration to embed into the request. +The `ai-proxy` plugin simplifies access to LLM providers and models by defining a standard request format +that allows key fields in plugin configuration to be embedded into the request. -Proxying requests to OpenAI is supported for now, other AI models will be supported soon. +Proxying requests to OpenAI is supported now. Other LLM services will be supported soon. ## Request Format From 977cf685672cb15f8a9b0dd0db71587bc0f113c0 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Tue, 10 Sep 2024 11:27:13 +0545 Subject: [PATCH 80/85] code review --- apisix/init.lua | 1 - t/plugin/ai-proxy.t | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/apisix/init.lua b/apisix/init.lua index 404b75358b18..103a8c1d7584 100644 --- a/apisix/init.lua +++ b/apisix/init.lua @@ -726,7 +726,6 @@ function _M.http_access_phase() end - function _M.dubbo_access_phase() ngx.ctx = fetch_ctx() end diff --git a/t/plugin/ai-proxy.t b/t/plugin/ai-proxy.t index 85f3ee64ec1d..445e406f60ab 100644 --- a/t/plugin/ai-proxy.t +++ b/t/plugin/ai-proxy.t @@ -102,8 +102,7 @@ add_block_preprocessor(sub { if header_auth == "Bearer token" or query_auth == "apikey" then ngx.req.read_body() local body, err = ngx.req.get_body_data() - local esc = body:gsub('"\\\""', '\"') - body, err = json.decode(esc) + body, err = json.decode(body) if not body.messages or #body.messages < 1 then ngx.status = 400 From 073b3f8b45a9c9c20fb3f3fb3a19dfc80ae3eaeb Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 11 Sep 2024 09:43:16 +0545 Subject: [PATCH 81/85] code review --- apisix/core/request.lua | 2 +- apisix/plugins/ai-proxy.lua | 18 ++++++++++-------- apisix/plugins/ai-proxy/drivers/openai.lua | 12 +++--------- t/sse_server_example/main.go | 13 ++++++------- 4 files changed, 20 insertions(+), 25 deletions(-) diff --git a/apisix/core/request.lua b/apisix/core/request.lua index 7d4842e6ccfc..fef4bf17e3f7 100644 --- a/apisix/core/request.lua +++ b/apisix/core/request.lua @@ -335,7 +335,7 @@ function _M.get_body(max_size, ctx) end -function _M.get_request_body_table() +function _M.get_json_request_body_table() local body, err = _M.get_body() if not body then return nil, { message = "could not get body: " .. (err or "request body is empty") } diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index f8514194798e..5834d648efa6 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -18,6 +18,8 @@ local core = require("apisix.core") local schema = require("apisix.plugins.ai-proxy.schema") local require = require local pcall = pcall +local internal_server_error = ngx.HTTP_INTERNAL_SERVER_ERROR +local bad_request = ngx.HTTP_BAD_REQUEST local ngx_req = ngx.req local ngx_print = ngx.print local ngx_flush = ngx.flush @@ -55,17 +57,17 @@ 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 + return bad_request, "unsupported content-type: " .. ct end - local request_table, err = core.request.get_request_body_table() + local request_table, err = core.request.get_json_request_body_table() if not request_table then - return 400, err + return bad_request, 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 + return bad_request, "request format doesn't match schema: " .. err end if conf.model.name then @@ -80,13 +82,13 @@ function _M.access(conf, ctx) 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 + return internal_server_error end local body_reader = res.body_reader if not body_reader then core.log.error("LLM sent no response body") - return 500 + return internal_server_error end if conf.passthrough then @@ -107,7 +109,7 @@ function _M.access(conf, ctx) return end - if core.table.try_read_attr(conf, "model", "options", "stream") then + if request_table.stream then while true do local chunk, err = body_reader() -- will read chunk by chunk if err then @@ -126,7 +128,7 @@ function _M.access(conf, ctx) local res_body, err = res:read_body() if not res_body then core.log.error("failed to read response body: ", err) - return 500 + return internal_server_error end keepalive_or_close(conf, httpc) return res.status, res_body diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index 6912d54972cc..ff3d44387908 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -26,6 +26,7 @@ local type = type -- globals local DEFAULT_HOST = "api.openai.com" local DEFAULT_PORT = 443 +local DEFAULT_PATH = "/v1/chat/completions" function _M.request(conf, request_table, ctx) @@ -54,15 +55,7 @@ function _M.request(conf, request_table, ctx) return nil, "failed to connect to LLM server: " .. err end - local query_params = "" - if conf.auth.query and type(conf.auth.query) == "table" then - query_params = core.string.encode_args(conf.auth.query) - if query_params and query_params ~= "" then - query_params = "?" .. query_params - end - end - - local path = (parsed_url.path or "/v1/chat/completions") .. query_params + local path = (parsed_url.path or DEFAULT_PATH) local headers = (conf.auth.header or {}) headers["Content-Type"] = "application/json" @@ -72,6 +65,7 @@ function _M.request(conf, request_table, ctx) keepalive = conf.keepalive, ssl_verify = conf.ssl_verify, path = path, + query = conf.auth.query } if conf.model.options then diff --git a/t/sse_server_example/main.go b/t/sse_server_example/main.go index 47f79953014a..e509c5c64dc1 100644 --- a/t/sse_server_example/main.go +++ b/t/sse_server_example/main.go @@ -31,19 +31,18 @@ func sseHandler(w http.ResponseWriter, r *http.Request) { w.Header().Set("Cache-Control", "no-cache") w.Header().Set("Connection", "keep-alive") + f, ok := w.(http.Flusher); + if !ok { + fmt.Fprintf(w, "[ERROR]") + return + } // A simple loop that sends a message every 2 seconds for i := 0; i < 5; i++ { // Create a message to send to the client fmt.Fprintf(w, "data: %s\n\n", time.Now().Format(time.RFC3339)) // Flush the data immediately to the client - if f, ok := w.(http.Flusher); ok { - f.Flush() - } else { - log.Println("Unable to flush data to client.") - break - } - + f.Flush() time.Sleep(500 * time.Millisecond) } fmt.Fprintf(w, "data: %s\n\n", "[DONE]") From 5831108aa55e16362788dbb0d5d49d3649ead61f Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Wed, 11 Sep 2024 10:57:15 +0545 Subject: [PATCH 82/85] code review --- apisix/plugins/ai-proxy.lua | 2 +- apisix/plugins/ai-proxy/drivers/openai.lua | 3 +-- t/sse_server_example/main.go | 2 +- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 5834d648efa6..c83f33408603 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -81,7 +81,7 @@ function _M.access(conf, ctx) local ai_driver = require("apisix.plugins.ai-proxy.drivers." .. conf.model.provider) 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) + core.log.error("failed to send request to LLM service: ", err) return internal_server_error end diff --git a/apisix/plugins/ai-proxy/drivers/openai.lua b/apisix/plugins/ai-proxy/drivers/openai.lua index ff3d44387908..c8f7f4b6223f 100644 --- a/apisix/plugins/ai-proxy/drivers/openai.lua +++ b/apisix/plugins/ai-proxy/drivers/openai.lua @@ -21,7 +21,6 @@ local http = require("resty.http") local url = require("socket.url") local pairs = pairs -local type = type -- globals local DEFAULT_HOST = "api.openai.com" @@ -77,7 +76,7 @@ function _M.request(conf, request_table, ctx) local res, err = httpc:request(params) if not res then - return 500, "failed to send request to LLM server: " .. err + return nil, err end return res, nil, httpc diff --git a/t/sse_server_example/main.go b/t/sse_server_example/main.go index e509c5c64dc1..ab976c86094a 100644 --- a/t/sse_server_example/main.go +++ b/t/sse_server_example/main.go @@ -36,7 +36,7 @@ func sseHandler(w http.ResponseWriter, r *http.Request) { fmt.Fprintf(w, "[ERROR]") return } - // A simple loop that sends a message every 2 seconds + // A simple loop that sends a message every 500ms for i := 0; i < 5; i++ { // Create a message to send to the client fmt.Fprintf(w, "data: %s\n\n", time.Now().Format(time.RFC3339)) From ab4c37e8611d650e0758e1c02516a14a2442d9ec Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 12 Sep 2024 08:27:58 +0545 Subject: [PATCH 83/85] update priority --- apisix/plugins/ai-proxy.lua | 2 +- conf/config.yaml.example | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index c83f33408603..75548b17dbf6 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -27,7 +27,7 @@ local ngx_flush = ngx.flush local plugin_name = "ai-proxy" local _M = { version = 0.5, - priority = 1004, + priority = 1055, name = plugin_name, schema = schema, } diff --git a/conf/config.yaml.example b/conf/config.yaml.example index 17b385216a88..d01be37272cc 100644 --- a/conf/config.yaml.example +++ b/conf/config.yaml.example @@ -478,6 +478,7 @@ plugins: # plugin list (sorted by priority) - body-transformer # priority: 1080 - ai-prompt-template # priority: 1071 - ai-prompt-decorator # priority: 1070 + - ai-proxy # priority: 1055 - proxy-mirror # priority: 1010 - proxy-rewrite # priority: 1008 - workflow # priority: 1006 From cdd37db20a8e3c1dad7c2f7be8408350e731761b Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 12 Sep 2024 08:43:50 +0545 Subject: [PATCH 84/85] update priority --- apisix/plugins/ai-proxy.lua | 2 +- conf/config.yaml.example | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/apisix/plugins/ai-proxy.lua b/apisix/plugins/ai-proxy.lua index 75548b17dbf6..8a0d8fa970d4 100644 --- a/apisix/plugins/ai-proxy.lua +++ b/apisix/plugins/ai-proxy.lua @@ -27,7 +27,7 @@ local ngx_flush = ngx.flush local plugin_name = "ai-proxy" local _M = { version = 0.5, - priority = 1055, + priority = 999, name = plugin_name, schema = schema, } diff --git a/conf/config.yaml.example b/conf/config.yaml.example index d01be37272cc..afbe22f9e633 100644 --- a/conf/config.yaml.example +++ b/conf/config.yaml.example @@ -478,7 +478,6 @@ plugins: # plugin list (sorted by priority) - body-transformer # priority: 1080 - ai-prompt-template # priority: 1071 - ai-prompt-decorator # priority: 1070 - - ai-proxy # priority: 1055 - proxy-mirror # priority: 1010 - proxy-rewrite # priority: 1008 - workflow # priority: 1006 @@ -487,6 +486,7 @@ plugins: # plugin list (sorted by priority) - limit-count # priority: 1002 - limit-req # priority: 1001 #- node-status # priority: 1000 + - ai-proxy # priority: 999 #- brotli # priority: 996 - gzip # priority: 995 - server-info # priority: 990 From e0e3e1575c0acba035490db9d8bff0cb8106a3d6 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 12 Sep 2024 09:21:48 +0545 Subject: [PATCH 85/85] fix test --- t/admin/plugins.t | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/t/admin/plugins.t b/t/admin/plugins.t index 101efad7caa6..bf3d485e8b31 100644 --- a/t/admin/plugins.t +++ b/t/admin/plugins.t @@ -99,10 +99,10 @@ proxy-mirror proxy-rewrite workflow api-breaker -ai-proxy limit-conn limit-count limit-req +ai-proxy gzip server-info traffic-split