feat: add ai-request-rewrite plugin (#12036)

Author: LiteSun

apisix/cli/config.lua

@@ -271,6 +271,7 @@ local _M = {
     "serverless-post-function",
     "ext-plugin-post-req",
     "ext-plugin-post-resp",
+    "ai-request-rewrite",
   },
   stream_plugins = { "ip-restriction", "limit-conn", "mqtt-proxy", "syslog" },
   plugin_attr = {

apisix/plugins/ai-request-rewrite.lua

@@ -0,0 +1,226 @@
+--
+-- 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 require = require
+local pcall = pcall
+local ngx = ngx
+local req_set_body_data = ngx.req.set_body_data
+local HTTP_BAD_REQUEST = ngx.HTTP_BAD_REQUEST
+local HTTP_INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR
+
+local plugin_name = "ai-request-rewrite"
+
+local auth_item_schema = {
+    type = "object",
+    patternProperties = {
+        ["^[a-zA-Z0-9._-]+$"] = {
+            type = "string"
+        }
+    }
+}
+
+local auth_schema = {
+    type = "object",
+    properties = {
+        header = auth_item_schema,
+        query = auth_item_schema
+    },
+    additionalProperties = false
+}
+
+local model_options_schema = {
+    description = "Key/value settings for the model",
+    type = "object",
+    properties = {
+        model = {
+            type = "string",
+            description = "Model to execute. Examples: \"gpt-3.5-turbo\" for openai, " ..
+            "\"deepseek-chat\" for deekseek, or \"qwen-turbo\" for openai-compatible services"
+        }
+    },
+    additionalProperties = true
+}
+
+local schema = {
+    type = "object",
+    properties = {
+        prompt = {
+            type = "string",
+            description = "The prompt to rewrite client request."
+        },
+        provider = {
+            type = "string",
+            description = "Name of the AI service provider.",
+            enum = {"openai", "openai-compatible", "deepseek"} -- add more providers later
+        },
+        auth = auth_schema,
+        options = model_options_schema,
+        timeout = {
+            type = "integer",
+            minimum = 1,
+            maximum = 60000,
+            default = 30000,
+            description = "Total timeout in milliseconds for requests to LLM service, " ..
+            "including connect, send, and read timeouts."
+        },
+        keepalive = {
+            type = "boolean",
+            default = true
+        },
+        keepalive_pool = {
+            type = "integer",
+            minimum = 1,
+            default = 30
+        },
+        ssl_verify = {
+            type = "boolean",
+            default = true
+        },
+        override = {
+            type = "object",
+            properties = {
+                endpoint = {
+                    type = "string",
+                    description = "To be specified to override " ..
+                    "the endpoint of the AI service provider."
+                }
+            }
+        }
+    },
+    required = {"prompt", "provider", "auth"}
+}
+
+local _M = {
+    version = 0.1,
+    priority = 1073,
+    name = plugin_name,
+    schema = schema
+}
+
+local function request_to_llm(conf, request_table, ctx)
+    local ok, ai_driver = pcall(require, "apisix.plugins.ai-drivers." .. conf.provider)
+    if not ok then
+        return nil, nil, "failed to load ai-driver: " .. conf.provider
+    end
+
+    local extra_opts = {
+        endpoint = core.table.try_read_attr(conf, "override", "endpoint"),
+        query_params = conf.auth.query or {},
+        headers = (conf.auth.header or {}),
+        model_options = conf.options
+    }
+
+    local res, err, httpc = ai_driver:request(conf, request_table, extra_opts)
+    if err then
+        return nil, nil, err
+    end
+
+    local resp_body, err = res:read_body()
+    httpc:close()
+    if err then
+        return nil, nil, err
+    end
+
+    return res, resp_body
+end
+
+
+local function parse_llm_response(res_body)
+    local response_table, err = core.json.decode(res_body)
+
+    if err then
+        return nil, "failed to decode llm response " .. ", err: " .. err
+    end
+
+    if not response_table.choices or not response_table.choices[1] then
+        return nil, "'choices' not in llm response"
+    end
+
+    local message = response_table.choices[1].message
+    if not message then
+        return nil, "'message' not in llm response choices"
+    end
+
+    return message.content
+end
+
+
+function _M.check_schema(conf)
+    -- openai-compatible should be used with override.endpoint
+    if conf.provider == "openai-compatible" then
+        local override = conf.override
+
+        if not override or not override.endpoint then
+            return false, "override.endpoint is required for openai-compatible provider"
+        end
+    end
+
+    return core.schema.check(schema, conf)
+end
+
+
+function _M.access(conf, ctx)
+    local client_request_body, err = core.request.get_body()
+    if err then
+        core.log.warn("failed to get request body: ", err)
+        return HTTP_BAD_REQUEST
+    end
+
+    if not client_request_body then
+        core.log.warn("missing request body")
+        return
+    end
+
+    -- Prepare request for LLM service
+    local ai_request_table = {
+        messages = {
+            {
+                role = "system",
+                content = conf.prompt
+            },
+            {
+                role = "user",
+                content = client_request_body
+            }
+        },
+        stream = false
+    }
+
+    -- Send request to LLM service
+    local res, resp_body, err = request_to_llm(conf, ai_request_table, ctx)
+    if err then
+        core.log.error("failed to request to LLM service: ", err)
+        return HTTP_INTERNAL_SERVER_ERROR
+    end
+
+    -- Handle LLM response
+    if res.status > 299 then
+        core.log.error("LLM service returned error status: ", res.status)
+        return HTTP_INTERNAL_SERVER_ERROR
+    end
+
+    -- Parse LLM response
+    local llm_response, err = parse_llm_response(resp_body)
+    if err then
+        core.log.error("failed to parse LLM response: ", err)
+        return HTTP_INTERNAL_SERVER_ERROR
+    end
+
+    req_set_body_data(llm_response)
+end
+
+return _M

docs/en/latest/config.json

@@ -103,7 +103,8 @@
             "plugins/ai-proxy",
             "plugins/ai-proxy-multi",
             "plugins/attach-consumer-label",
-            "plugins/ai-rag"
+            "plugins/ai-rag",
+            "plugins/ai-request-rewrite"
           ]
         },
         {