feat: ai-proxy plugin

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
 
 

apisix/cli/config.lua

@@ -217,6 +217,7 @@ local _M = {
     "proxy-rewrite",
     "workflow",
     "api-breaker",
+    "ai-proxy",
     "limit-conn",
     "limit-count",
     "limit-req",

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 {

apisix/constants.lua

@@ -43,4 +43,6 @@ return {
         ["/stream_routes"] = true,
         ["/plugin_metadata"] = true,
     },
+    CHAT = "llm/chat",
+    COMPLETION = "llm/completions",
 }

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
@@ -722,7 +723,19 @@ function _M.http_access_phase()
         plugin.run_plugin("access", plugins, api_ctx)
     end
 
-    _M.handle_upstream(api_ctx, route, enable_websocket)
+    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 +906,16 @@ 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
 
 

apisix/plugins/ai-proxy.lua

@@ -0,0 +1,103 @@
+--
+-- 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")
+local require = require
+
+local ngx_req = ngx.req
+local ngx = ngx
+
+local plugin_name = "ai-proxy"
+local _M = {
+    version = 0.5,
+    priority = 1004,
+    name = plugin_name,
+    schema = schema,
+}
+
+
+function _M.check_schema(conf)
+    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()
+    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 = {}
+
+    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
+
+    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)
+        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

apisix/plugins/ai-proxy/drivers/openai.lua

@@ -0,0 +1,67 @@
+--
+-- Licensed to the Apache Software Foundation (ASF) under one or more
+-- contributor license agreements.  See the NOTICE file distributed with
+-- this work for additional information regarding copyright ownership.
+-- The ASF licenses this file to You under the Apache License, Version 2.0
+-- (the "License"); you may not use this file except in compliance with
+-- the License.  You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+local _M = {}
+
+local core = require("apisix.core")
+local test_scheme = os.getenv("AI_PROXY_TEST_SCHEME")
+local ngx = ngx
+local pairs = pairs
+
+-- 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, 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)
+    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]
+    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
+        local args = core.request.get_uri_args(ctx)
+        args[conf.auth.name] = conf.auth.value
+        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
+
+return _M