feat: ai-content-moderation plugin

apisix-master-0.rockspec

@@ -82,7 +82,7 @@ dependencies = {
     "lua-resty-t1k = 1.1.5",
     "brotli-ffi = 0.3-1",
     "lua-ffi-zlib = 0.6-0",
-    "api7-lua-resty-aws == 2.0.1-1",
+    "api7-lua-resty-aws == 2.0.2-1",
 }
 
 build = {

apisix/cli/config.lua

@@ -215,6 +215,7 @@ local _M = {
     "body-transformer",
     "ai-prompt-template",
     "ai-prompt-decorator",
+    "ai-content-moderation",
     "proxy-mirror",
     "proxy-rewrite",
     "workflow",

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_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/plugins/ai-content-moderation.lua

@@ -0,0 +1,170 @@
+--
+-- 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 aws = require("resty.aws")
+local http = require("resty.http")
+local fetch_secrets = require("apisix.secret").fetch_secrets
+
+local aws_instance = aws()
+local next = next
+local pairs = pairs
+local unpack = unpack
+local type = type
+local ipairs = ipairs
+
+local aws_comprehend_schema = {
+    type = "object",
+    properties = {
+        access_key_id = { type = "string" },
+        secret_access_key = { type = "string" },
+        region = { type = "string" },
+        endpoint = {
+            type = "string",
+            pattern = [[^https?://]]
+        },
+    },
+    required = { "access_key_id", "secret_access_key", "region", }
+}
+
+local schema = {
+    type = "object",
+    properties = {
+        provider = {
+            type = "object",
+            properties = {
+                aws_comprehend = aws_comprehend_schema
+            },
+            -- change to oneOf/enum while implementing support for other services
+            required = { "aws_comprehend" }
+        },
+        moderation_categories = {
+            type = "object",
+            patternProperties = {
+                -- luacheck: push max code line length 300
+                ["^(PROFANITY|HATE_SPEECH|INSULT|HARASSMENT_OR_ABUSE|SEXUAL|VIOLENCE_OR_THREAT)$"] = {
+                -- luacheck: pop
+                type = "number",
+                    minimum = 0,
+                    maximum = 1
+                }
+            },
+            additionalProperties = false
+        },
+        toxicity_level = {
+            type = "number",
+            minimum = 0,
+            maximum = 1,
+            default = 0.5
+        },
+        type = {
+            type = "string",
+            enum = { "openai" },
+        }
+    },
+    required = { "provider", "type" },
+}
+
+
+local _M = {
+    version  = 0.1,
+    priority = 1040, -- TODO: might change
+    name     = "ai-content-moderation",
+    schema   = schema,
+}
+
+
+function _M.check_schema(conf)
+    return core.schema.check(schema, conf)
+end
+
+function _M.rewrite(conf, ctx)
+    conf = fetch_secrets(conf, true, conf, "")
+    if not conf then
+        return 500, "failed to retrieve secrets from conf"
+    end
+
+    local body, err = core.request.get_body_table()
+    if not body then
+        return 400, err
+    end
+
+    local msgs = body.messages
+    if not msgs or type(msgs) ~= "table" or #msgs < 1 then
+        return 400, "messages not found in request body"
+    end
+
+    local provider = conf.provider[next(conf.provider)]
+
+    -- TODO support secret
+    local credentials = aws_instance:Credentials({
+        accessKeyId = provider.access_key_id,
+        secretAccessKey = provider.secret_access_key,
+        sessionToken = provider.session_token,
+    })
+
+    local default_endpoint = "https://comprehend." .. provider.region .. ".amazonaws.com"
+    local scheme, host, port = unpack(http:parse_uri(provider.endpoint or default_endpoint))
+    local endpoint = scheme .. "://" .. host
+    aws_instance.config.endpoint = endpoint
+    aws_instance.config.ssl_verify = false
+
+    local comprehend = aws_instance:Comprehend({
+        credentials = credentials,
+        endpoint = endpoint,
+        region = provider.region,
+        port = port,
+    })
+
+    local ai_module = require("apisix.plugins.ai." .. conf.type)
+    local create_request_text_segments = ai_module.create_request_text_segments
+
+    local text_segments = create_request_text_segments(msgs)
+    local res, err = comprehend:detectToxicContent({
+        LanguageCode = "en",
+        TextSegments = text_segments,
+    })
+
+    if not res then
+        core.log.error("failed to send request to ", provider, ": ", err)
+        return 500, err
+    end
+
+    local results = res.body and res.body.ResultList
+    if not results or type(results) ~= "table" or #results < 1 then
+        return 500, "failed to get moderation results from response"
+    end
+
+    for _, result in ipairs(results) do
+        if conf.moderation_categories then
+            for _, item in pairs(result.Labels) do
+                if not conf.moderation_categories[item.Name] then
+                    goto continue
+                end
+                if item.Score > conf.moderation_categories[item.Name] then
+                    return 400, "request body exceeds " .. item.Name .. " threshold"
+                end
+                ::continue::
+            end
+        end
+
+        if result.Toxicity > conf.toxicity_level then
+            return 400, "request body exceeds toxicity threshold"
+        end
+    end
+end
+
+return _M

apisix/plugins/ai/openai.lua

@@ -0,0 +1,33 @@
+--
+-- 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 ipairs = ipairs
+
+local _M = {}
+
+
+function _M.create_request_text_segments(msgs)
+    local text_segments = {}
+    for _, msg in ipairs(msgs) do
+        core.table.insert_tail(text_segments, {
+            Text = msg.content
+        })
+    end
+    return text_segments
+end
+
+return  _M

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-content-moderation          # priority: 1040 TODO: compare priority with other ai plugins
   - proxy-mirror                   # priority: 1010
   - proxy-rewrite                  # priority: 1008
   - workflow                       # priority: 1006

docs/en/latest/config.json

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