ai proxy

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

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

apisix/cli/ngx_tpl.lua

@@ -809,6 +809,85 @@ 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 %}
+            {* 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/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 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 }
+    end
+
+    return body_tab
+end
+
+
 function _M.get_scheme(ctx)
     if not ctx then
         ctx = ngx.ctx.api_ctx

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

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,6 +762,25 @@ 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
+        stash_ngx_ctx()
+        return ngx.exec("@disable_proxy_buffering")
+    end
+end
+
+
+function _M.disable_proxy_buffering_access_phase()
+    ngx.ctx = fetch_ctx()
 end
 
 

apisix/plugins/ai-proxy.lua

@@ -0,0 +1,112 @@
+--
+-- Licensed to the Apache Software Foundation (ASF) under one or more
+-- contributor license agreements.  See the NOTICE file distributed with
+-- this work for additional information regarding copyright ownership.
+-- The ASF licenses this file to You under the Apache License, Version 2.0
+-- (the "License"); you may not use this file except in compliance with
+-- the License.  You may obtain a copy of the License at
+--
+--     http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing, software
+-- distributed under the License is distributed on an "AS IS" BASIS,
+-- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+-- See the License for the specific language governing permissions and
+-- limitations under the License.
+--
+local core = require("apisix.core")
+local schema = require("apisix.plugins.ai-proxy.schema")
+local require = require
+local pcall = pcall
+
+local ngx_req = ngx.req
+
+local plugin_name = "ai-proxy"
+local _M = {
+    version = 0.5,
+    priority = 1004,
+    name = plugin_name,
+    schema = schema,
+}
+
+
+function _M.check_schema(conf)
+    local ai_driver = pcall(require, "apisix.plugins.ai-proxy.drivers." .. conf.model.provider)
+    if not ai_driver then
+        return false, "provider: " .. conf.model.provider .. " is not supported."
+    end
+    return core.schema.check(schema.plugin_schema, conf)
+end
+
+
+local CONTENT_TYPE_JSON = "application/json"
+
+
+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
+    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
+        ctx.disable_proxy_buffering = true
+    else
+        ctx.subrequest = 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)
+    -------------------------------- MODIFIED --------------------------------
+    local res, err = ai_driver.request(conf, request_table, ctx)
+    if not res then
+        return 500, "failed to proxy LLM request: " .. err
+    end
+
+    local data
+
+    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
+        ngx_req.set_body_data(res.data)
+        return
+    end
+
+    if core.table.try_read_attr(conf, "model", "options", "stream") then
+        local reader = res.body_reader
+        while true do
+            local buffer, err = reader()
+            if err then
+                ngx.log(ngx.ERR, err)
+                break
+            end
+
+            ngx.print(buffer)
+            ngx.flush(true) -- just a example, need more verification
+        end
+    else
+        return 200, res.data
+    end
+    -- we may have to simulate an SSE (chunked) response through the scheme mentioned in
+    -- https://github.com/openresty/lua-nginx-module/issues/1736#issuecomment-650143112
+    -- return 200, res.data
+    --------------------------------------------------------------------------
+end
+
+return _M