feat: add ai-prompt-guard plugin

apisix/cli/config.lua

@@ -216,6 +216,7 @@ local _M = {
     "body-transformer",
     "ai-prompt-template",
     "ai-prompt-decorator",
+    "ai-prompt-guard",
     "ai-rag",
     "ai-content-moderation",
     "proxy-mirror",

apisix/plugins/ai-prompt-guard.lua

@@ -0,0 +1,142 @@
+--
+-- 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 ngx = ngx
+local ipairs = ipairs
+local table = table
+local re_compile  = require("resty.core.regex").re_match_compile
+local re_find = ngx.re.find
+
+local plugin_name = "ai-prompt-guard"
+
+local schema = {
+    type = "object",
+    properties = {
+        match_all_roles = {
+            type = "boolean",
+            default = false,
+        },
+        match_all_conversation_history = {
+            type = "boolean",
+            default = false,
+        },
+        allow_patterns = {
+            type = "array",
+            items = {type = "string"},
+            default = {},
+        },
+        deny_patterns = {
+            type = "array",
+            items = {type = "string"},
+            default = {},
+        },
+    },
+}
+
+local _M = {
+    version = 0.1,
+    priority = 1072,
+    name = plugin_name,
+    schema = schema,
+}
+
+function _M.check_schema(conf)
+    local ok, err = core.schema.check(schema, conf)
+    if not ok then
+        return false, err
+    end
+
+    -- Validate allow_patterns
+    for _, pattern in ipairs(conf.allow_patterns) do
+        local compiled = re_compile(pattern, "jou")
+        if not compiled then
+            return false, "invalid allow_pattern: " .. pattern
+        end
+    end
+
+    -- Validate deny_patterns
+    for _, pattern in ipairs(conf.deny_patterns) do
+        local compiled = re_compile(pattern, "jou")
+        if not compiled then
+            return false, "invalid deny_pattern: " .. pattern
+        end
+    end
+
+    return true
+end
+
+local function get_content_to_check(conf, messages)
+    local contents = {}
+    if conf.match_all_conversation_history then
+        for _, msg in ipairs(messages) do
+            if msg.content then
+                core.table.insert(contents, msg.content)
+            end
+        end
+    else
+        if #messages > 0 then
+            local last_msg = messages[#messages]
+            if last_msg.content then
+                core.table.insert(contents, last_msg.content)
+            end
+        end
+    end
+    return table.concat(contents, " ")
+end
+
+function _M.access(conf, ctx)
+    local body = core.request.get_body()
+    if not body then
+        core.log.error("Empty request body")
+        return 400, {message = "Empty request body"}
+    end
+
+    local json_body, err = core.json.decode(body)
+    if err then
+        return 400, {message = err}
+    end
+
+    local messages = json_body.messages or {}
+    if not conf.match_all_roles and #messages > 0 and messages[#messages].role ~= "user" then
+        return
+    end
+
+    local content_to_check = get_content_to_check(conf, messages)
+
+     -- Allow patterns check
+     if #conf.allow_patterns > 0 then
+        local any_allowed = false
+        for _, pattern in ipairs(conf.allow_patterns) do
+            if re_find(content_to_check, pattern, "jou") then
+                any_allowed = true
+                break
+            end
+        end
+        if not any_allowed then
+            return 400, {message = "Request doesn't match allow patterns"}
+        end
+    end
+
+    -- Deny patterns check
+    for _, pattern in ipairs(conf.deny_patterns) do
+        if re_find(content_to_check, pattern, "jou") then
+            return 400, {message = "Request contains prohibited content"}
+        end
+    end
+end
+
+return _M

conf/config.yaml.example

@@ -480,6 +480,7 @@ plugins:                           # plugin list (sorted by priority)
   - body-transformer               # priority: 1080
   - ai-prompt-template             # priority: 1071
   - ai-prompt-decorator            # priority: 1070
+  - ai-prompt-guard                # priority: 1072
   - ai-rag                         # priority: 1060
   - ai-content-moderation          # priority: 1040 TODO: compare priority with other ai plugins
   - proxy-mirror                   # priority: 1010

docs/en/latest/config.json

@@ -82,7 +82,8 @@
             "plugins/ext-plugin-post-resp",
             "plugins/inspect",
             "plugins/ocsp-stapling",
-            "plugins/ai-content-moderation"
+            "plugins/ai-content-moderation",
+            "plugins/ai-prompt-guard"
           ]
         },
         {

docs/en/latest/plugins/ai-prompt-guard.md

@@ -0,0 +1,89 @@
+---
+title: ai-prompt-guard
+keywords:
+  - Apache APISIX
+  - API Gateway
+  - Plugin
+  - ai-prompt-guard
+description: This document contains information about the Apache APISIX ai-prompt-guard Plugin.
+---
+
+<!--
+#
+# 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.
+#
+-->
+
+## Description
+
+The `ai-prompt-guard` plugin safeguards your AI endpoints by inspecting and validating incoming prompt messages. It checks the content of requests against user-defined allowed and denied patterns to ensure that only approved inputs are processed. Based on its configuration, the plugin can either examine just the latest message or the entire conversation history, and it can be set to check prompts from all roles or only from end users.
+
+When both **allow** and **deny** patterns are configured, the plugin first ensures that at least one allowed pattern is matched. If none match, the request is rejected with a _"Request doesn't match allow patterns"_ error. If an allowed pattern is found, it then checks for any occurrences of denied patterns—rejecting the request with a _"Request contains prohibited content"_ error if any are detected.
+
+## Plugin Attributes
+
+| **Field**                      | **Required** | **Type**  | **Description**                                                                                                                                                      |
+| ------------------------------ | ------------ | --------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| match_all_roles                | No           | boolean   | If set to `true`, the plugin will check prompt messages from all roles. Otherwise, it only validates when its role is `"user"`. Default is `false`. |
+| match_all_conversation_history | No           | boolean   | When enabled, all messages in the conversation history are concatenated and checked. If `false`, only the content of the last message is examined. Default is `false`. |
+| allow_patterns                 | No           | array     | A list of regex patterns. When provided, the prompt must match **at least one** pattern to be considered valid.                                                      |
+| deny_patterns                  | No           | array     | A list of regex patterns. If any of these patterns match the prompt content, the request is rejected.                                                                  |
+
+## Example usage
+
+Create a route with the `ai-prompt-guard` 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": "/v1/chat/completions",
+    "plugins": {
+        "ai-prompt-guard": {
+           "match_all_roles": true,
+              "allow_patterns": [
+                 "goodword"
+              ],
+          "deny_patterns": [
+            "badword"
+          ]
+     }
+  },
+    "upstream": {
+      "type": "roundrobin",
+      "nodes": {
+        "api.openai.com:443": 1
+      },
+      "pass_host": "node",
+      "scheme": "https"
+    }
+  }'
+```
+
+Now send a request:
+
+```shell
+curl http://127.0.0.1:9080/v1/chat/completions -i -XPOST  -H 'Content-Type: application/json' -d '{
+  "model": "gpt-4",
+  "messages": [{ "role": "user", "content": "badword request" }]
+}' -H "Authorization: Bearer <your token here>"
+```
+
+The request will fail with 400 error and following response.
+
+```bash
+{"message":"Request doesn't match allow patterns"}
+```

t/admin/plugins.t

@@ -94,6 +94,7 @@ opa
 authz-keycloak
 proxy-cache
 body-transformer
+ai-prompt-guard
 ai-prompt-template
 ai-prompt-decorator
 ai-rag