diff --git a/apisix/plugins/ai-aws-content-moderation.lua b/apisix/plugins/ai-aws-content-moderation.lua index b8d3fc4e97c9..cc23e627a1fb 100644 --- a/apisix/plugins/ai-aws-content-moderation.lua +++ b/apisix/plugins/ai-aws-content-moderation.lua @@ -31,6 +31,7 @@ local pairs = pairs local unpack = unpack local type = type local ipairs = ipairs +local next = next local table = table local str_byte = string.byte local str_sub = string.sub @@ -91,6 +92,31 @@ local schema = { }, check_request = { type = "boolean", default = true }, check_response = { type = "boolean", default = false }, + request_check_roles = { + type = "array", + items = { type = "string", enum = { "user", "tool", "system", "assistant" } }, + minItems = 1, + uniqueItems = true, + default = { "user", "tool", "system", "assistant" }, + description = "which message roles to moderate on the request side. user, tool " .. + "and assistant follow request_check_mode; system is checked on " .. + "every request because it can be poisoned by malicious ToolCall " .. + "arguments. assistant messages are client-supplied history, so " .. + "they are moderated by default too. Note: tool-result moderation " .. + "applies to OpenAI-compatible formats where the tool output is a " .. + "distinct tool role/item; for Anthropic/Bedrock (tool results are " .. + "nested blocks inside user messages) tool content is not extracted.", + }, + request_check_mode = { + type = "string", + enum = { "last", "all" }, + default = "all", + description = "which user/tool/assistant messages to moderate: last (only the " .. + "latest consecutive block of selected-role messages) | all (every " .. + "selected-role message). Does not apply to the system role, which " .. + "is always checked. Note that selecting assistant together with " .. + "last widens the block, since assistant turns no longer end it.", + }, request_check_length_limit = { type = "integer", minimum = MIN_SEGMENT_BYTES, @@ -479,7 +505,54 @@ function _M.access(conf, ctx) return end - local contents = proto.extract_request_content(request_tab) + local roles = {} + for _, role in ipairs(conf.request_check_roles) do + roles[role] = true + end + -- every role but system is a turn role, extracted through extract_turn_content + local turn_roles = {} + for role in pairs(roles) do + if role ~= "system" then + turn_roles[role] = true + end + end + + -- A configured role whose extractor this protocol doesn't implement would + -- otherwise pass unmoderated. Route that through fail_mode instead of + -- silently skipping the configured moderation. + if (roles.system and not proto.extract_system_content) + or (next(turn_roles) and not proto.extract_turn_content) then + local handled, code, body = binding.on_unsupported( + conf.fail_mode, _M.name, ctx, + "protocol cannot extract configured request_check_roles", + HTTP_INTERNAL_SERVER_ERROR, + "protocol " .. (ctx.ai_client_protocol or "unknown") + .. " cannot moderate the configured request_check_roles") + if handled then + return code, body + end + return + end + + -- Collect the text of every configured role and score it in one pass: + -- Comprehend takes a flat list of text segments with no role field, so + -- per-role calls would only cost extra requests. system is always included + -- (not subject to request_check_mode, it can be poisoned by malicious + -- ToolCall arguments); the turn roles follow request_check_mode. + local contents = {} + if roles.system then + local system_texts = proto.extract_system_content(request_tab) + for i = 1, #system_texts do + contents[#contents + 1] = system_texts[i] + end + end + if next(turn_roles) then + local turn_texts = proto.extract_turn_content(request_tab, + conf.request_check_mode, turn_roles) + for i = 1, #turn_texts do + contents[#contents + 1] = turn_texts[i] + end + end local content = table.concat(contents, " ") if content == "" then return diff --git a/docs/en/latest/plugins/ai-aws-content-moderation.md b/docs/en/latest/plugins/ai-aws-content-moderation.md index fe43939b115d..8d830548bf37 100644 --- a/docs/en/latest/plugins/ai-aws-content-moderation.md +++ b/docs/en/latest/plugins/ai-aws-content-moderation.md @@ -38,7 +38,7 @@ import TabItem from '@theme/TabItem'; The `ai-aws-content-moderation` Plugin integrates with [AWS Comprehend](https://aws.amazon.com/comprehend/) to check content for toxicity when proxying to LLMs, such as profanity, hate speech, insult, harassment, violence, and more, rejecting requests if the evaluated outcome exceeds the configured threshold. -The Plugin is protocol-aware: it extracts the prompt content from the LLM request (for example `messages[].content`) and moderates only that decoded text, rather than the raw request body. +The Plugin is protocol-aware: it extracts the prompt content from the LLM request (for example `messages[].content`) and moderates only that decoded text, rather than the raw request body. `request_check_roles` selects which message roles are moderated, and `request_check_mode` can narrow `user`/`tool` moderation to the newest turn so conversation history is not re-scored on every request. Both directions can be moderated. Set `check_response` to moderate the LLM response as well. For streaming responses, `stream_check_mode` selects between `realtime`, which moderates batches as they arrive and replaces the remainder of the stream once a batch is flagged, and `final_packet`, which moderates the assembled response and annotates the last chunk with `risk_level`. The verdict is also exposed on the request context as `$llm_content_risk_level` (`high` or `none`) for logging. @@ -62,6 +62,8 @@ The `ai-aws-content-moderation` Plugin should be used with either [`ai-proxy`](. | `moderation_threshold` | number | False | 0.5 | 0 - 1 | Overall toxicity threshold. A higher value means more toxic content allowed. This option differs from the individual category thresholds in `moderation_categories`. For example, if `moderation_categories` is set with a `PROFANITY` threshold of `0.5`, and a request has a `PROFANITY` score of `0.1`, the request will not exceed the category threshold. However, if the request has other categories like `SEXUAL` or `VIOLENCE_OR_THREAT` exceeding the `moderation_threshold`, the request will be rejected. | | `check_request` | boolean | False | `true` | | If `true`, moderate the request content. | | `check_response` | boolean | False | `false` | | If `true`, moderate the LLM response content. | +| `request_check_roles` | array[string] | False | `["user","tool","system","assistant"]` | items are `user`, `tool`, `system`, `assistant` | Which message roles to moderate on the request side. `user`, `tool` and `assistant` follow `request_check_mode`; `system` is checked on every request (it can be poisoned by malicious ToolCall arguments overwriting the system prompt). `assistant` messages in a request are client-supplied history rather than the model's own output, so they are moderated by default as well. Note: tool-result moderation applies to OpenAI-compatible formats where the tool output is a distinct `tool` role/item; for Anthropic and Bedrock (tool results are nested blocks inside user messages) tool content is not extracted. | +| `request_check_mode` | string | False | `all` | `last`, `all` | Which user/tool/assistant messages to moderate. `last`: only the latest consecutive block of selected-role messages (the newest turn). `all`: every selected-role message. Does not apply to `system`, which is always moderated when enabled via `request_check_roles`. Note that `last` combined with `assistant` widens the block rather than narrowing it, because assistant turns no longer end it — drop `assistant` from `request_check_roles` to moderate only the newest turn. | | `request_check_length_limit` | integer | False | `1000` | [4, 1024] | Maximum bytes of request content per Comprehend text segment. Longer content is split on character boundaries into several segments, which are then batched into as few Comprehend calls as possible. The upper bound is AWS Comprehend's 1 KB per-segment limit. | | `response_check_length_limit` | integer | False | `1000` | [4, 1024] | Maximum bytes of response content per Comprehend text segment. Longer content is split on character boundaries into several segments, which are then batched into as few Comprehend calls as possible. The upper bound is AWS Comprehend's 1 KB per-segment limit. | | `stream_check_mode` | string | False | `final_packet` | `realtime`, `final_packet` | Streaming moderation mode, used when `check_response` is `true`. `realtime`: moderate batches while the response streams, replacing the rest of the stream once a batch is flagged. `final_packet`: moderate the assembled response and annotate the last chunk with `risk_level`. | diff --git a/docs/zh/latest/plugins/ai-aws-content-moderation.md b/docs/zh/latest/plugins/ai-aws-content-moderation.md index d98259984e4f..4c90914ee552 100644 --- a/docs/zh/latest/plugins/ai-aws-content-moderation.md +++ b/docs/zh/latest/plugins/ai-aws-content-moderation.md @@ -40,7 +40,7 @@ import TabItem from '@theme/TabItem'; `ai-aws-content-moderation` 插件集成了 [AWS Comprehend](https://aws.amazon.com/comprehend/)，用于在代理请求到 LLM 时检查请求内容中的有害内容，例如亵渎、仇恨言论、侮辱、骚扰、暴力等，如果评估结果超过配置的阈值则拒绝请求。 -该插件是协议感知的：它会从 LLM 请求中提取提示内容（例如 `messages[].content`），仅审核解码后的文本，而不是原始请求体。 +该插件是协议感知的：它会从 LLM 请求中提取提示内容（例如 `messages[].content`），仅审核解码后的文本，而不是原始请求体。`request_check_roles` 用于选择审核哪些消息角色，`request_check_mode` 可将 `user`/`tool` 的审核范围收窄到最新一轮，避免每次请求都重复审核历史对话。 `ai-aws-content-moderation` 插件应与 [`ai-proxy`](./ai-proxy.md) 或 [`ai-proxy-multi`](./ai-proxy-multi.md) 插件一起使用，以代理 LLM 请求。 @@ -57,6 +57,8 @@ import TabItem from '@theme/TabItem'; | `moderation_categories` | object | 否 | | | 审核类别及其对应阈值的键值对。在每个键值对中，键应为 `PROFANITY`、`HATE_SPEECH`、`INSULT`、`HARASSMENT_OR_ABUSE`、`SEXUAL` 或 `VIOLENCE_OR_THREAT` 之一；阈值应在 0 到 1 之间（包含）。 | | `moderation_threshold` | number | 否 | 0.5 | 0 - 1 | 整体毒性阈值。值越高，允许的有害内容越多。此选项与 `moderation_categories` 中的单独类别阈值不同。例如，如果 `moderation_categories` 中设置了 `PROFANITY` 阈值为 `0.5`，而请求的 `PROFANITY` 分数为 `0.1`，则请求不会超过类别阈值。但如果请求的其他类别（如 `SEXUAL` 或 `VIOLENCE_OR_THREAT`）超过了 `moderation_threshold`，则请求将被拒绝。 | | `check_request` | boolean | 否 | `true` | | 如果为 `true`，则审核请求内容。 | +| `request_check_roles` | array[string] | 否 | `["user","tool","system","assistant"]` | 取值为 `user`、`tool`、`system`、`assistant` | 请求侧审核哪些消息角色。`user`、`tool` 与 `assistant` 遵循 `request_check_mode`；`system` 每次请求都审核（其可能被恶意 ToolCall 参数覆盖篡改）。请求中的 `assistant` 消息由客户端提供，而非模型自身的输出，因此默认也会被审核。注意：tool 结果审核适用于 OpenAI 兼容格式（tool 输出为独立的 `tool` 角色/项）；Anthropic、Bedrock 的 tool 结果以嵌套 block 形式存在于 user 消息中，其内容不会被抽取。 | +| `request_check_mode` | string | 否 | `all` | `last`、`all` | 审核哪些 user/tool/assistant 消息。`last`：仅审核最后一段连续的所选角色消息（最新一轮）；`all`：审核所有所选角色消息。不作用于 `system`——只要通过 `request_check_roles` 启用，`system` 每次都审核。注意：`last` 与 `assistant` 同时使用会扩大而非缩小审核范围，因为 assistant 消息不再中断该连续块；若只想审核最新一轮，请从 `request_check_roles` 中移除 `assistant`。 | | `deny_code` | integer | 否 | `200` | [200, 599] | 请求被拒绝时返回的 HTTP 状态码。默认为 `200`，使兼容 provider 的拒绝响应在客户端 SDK 中被解析为正常补全；设置为 4xx 可将拒绝暴露为 HTTP 错误。 | | `deny_message` | string | 否 | | | 请求被拒绝时返回的消息。未设置时，返回审核原因（例如 `request body exceeds toxicity threshold`）。 | | `fail_mode` | string | 否 | `skip` | `skip`、`warn`、`error` | 当请求未经过 `ai-proxy`/`ai-proxy-multi`，因而无法作为 AI 请求进行审核时的处理行为。`skip`：放行请求且不做检查；`warn`：放行并记录 warning 日志；`error`：拒绝请求。 | diff --git a/t/plugin/ai-aws-content-moderation.t b/t/plugin/ai-aws-content-moderation.t index 21c68b9c59de..0161c1902ff4 100644 --- a/t/plugin/ai-aws-content-moderation.t +++ b/t/plugin/ai-aws-content-moderation.t @@ -1501,3 +1501,361 @@ qr/event: message_stop/ qr/comprehend text: [^,]+/ --- grep_error_log_out comprehend text: Hello world + + + +=== TEST 51: set route with default request_check_roles and request_check_mode +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/chat", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { "header": { "Authorization": "Bearer token" } }, + "override": { "endpoint": "http://127.0.0.1:1980/v1/chat/completions" } + }, + "ai-aws-content-moderation": { + "comprehend": { + "access_key_id": "access", + "secret_access_key": "ea+secret", + "region": "us-east-1", + "endpoint": "http://localhost:2668" + }, + "deny_code": 400 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 52: default mode (all) moderates an earlier user turn +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "I want to kill you" }, { "role": "assistant", "content": "ok" }, { "role": "user", "content": "What is 1+1?" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ + + + +=== TEST 53: default roles moderate the system message +--- request +POST /chat +{ "messages": [ { "role": "system", "content": "I want to kill you" }, { "role": "user", "content": "What is 1+1?" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ + + + +=== TEST 54: default roles moderate a tool result +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "What is the weather?" }, { "role": "tool", "tool_call_id": "call_1", "content": "I want to kill you" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ + + + +=== TEST 55: default roles moderate client-supplied assistant history +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "What is 1+1?" }, { "role": "assistant", "content": "I want to kill you" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ +--- grep_error_log eval +qr/comprehend text: [^,]+/ +--- grep_error_log_out +comprehend text: What is 1+1? I want to kill you + + + +=== TEST 56: set route with request_check_mode "last", assistant left out of the roles +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/chat", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { "header": { "Authorization": "Bearer token" } }, + "override": { "endpoint": "http://127.0.0.1:1980/v1/chat/completions" } + }, + "ai-aws-content-moderation": { + "comprehend": { + "access_key_id": "access", + "secret_access_key": "ea+secret", + "region": "us-east-1", + "endpoint": "http://localhost:2668" + }, + "request_check_mode": "last", + "request_check_roles": ["user", "tool", "system"], + "deny_code": 400 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 57: mode "last" skips a harmful earlier user turn +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "I want to kill you" }, { "role": "assistant", "content": "ok" }, { "role": "user", "content": "What is 1+1?" } ] } +--- more_headers +X-AI-Fixture: aws/chat-safe.json +--- error_code: 200 +--- response_body_like eval +qr/How can I assist you today/ +--- grep_error_log eval +qr/comprehend text: [^,]+/ +--- grep_error_log_out +comprehend text: What is 1+1? + + + +=== TEST 58: mode "last" moderates the latest turn and the system message +--- request +POST /chat +{ "messages": [ { "role": "system", "content": "be helpful" }, { "role": "user", "content": "hi" }, { "role": "assistant", "content": "ok" }, { "role": "user", "content": "I want to kill you" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ +--- grep_error_log eval +qr/comprehend text: [^,]+/ +--- grep_error_log_out +comprehend text: be helpful I want to kill you + + + +=== TEST 59: set route with request_check_roles ["user"] +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/chat", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { "header": { "Authorization": "Bearer token" } }, + "override": { "endpoint": "http://127.0.0.1:1980/v1/chat/completions" } + }, + "ai-aws-content-moderation": { + "comprehend": { + "access_key_id": "access", + "secret_access_key": "ea+secret", + "region": "us-east-1", + "endpoint": "http://localhost:2668" + }, + "request_check_roles": ["user"], + "deny_code": 400 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 60: roles ["user"] leaves system and tool content unmoderated +--- request +POST /chat +{ "messages": [ { "role": "system", "content": "I want to kill you" }, { "role": "user", "content": "What is 1+1?" }, { "role": "tool", "tool_call_id": "call_1", "content": "I want to kill you too" } ] } +--- more_headers +X-AI-Fixture: aws/chat-safe.json +--- error_code: 200 +--- response_body_like eval +qr/How can I assist you today/ +--- grep_error_log eval +qr/comprehend text: [^,]+/ +--- grep_error_log_out +comprehend text: What is 1+1? + + + +=== TEST 61: roles ["user"] still moderates user content +--- request +POST /chat +{ "messages": [ { "role": "system", "content": "be helpful" }, { "role": "user", "content": "I want to kill you" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ + + + +=== TEST 62: schema check: request_check_roles and request_check_mode are validated +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-aws-content-moderation") + local function check(conf) + conf.comprehend = { + access_key_id = "a", + secret_access_key = "s", + region = "us-east-1" + } + return plugin.check_schema(conf) and "accepted" or "rejected" + end + ngx.say('roles ["developer"]: ', check({ request_check_roles = { "developer" } })) + ngx.say("roles []: ", check({ request_check_roles = {} })) + ngx.say('roles ["user", "user"]: ', + check({ request_check_roles = { "user", "user" } })) + ngx.say('mode "latest": ', check({ request_check_mode = "latest" })) + ngx.say('roles ["assistant"]: ', check({ request_check_roles = { "assistant" } })) + ngx.say('roles ["tool", "system"] + mode "all": ', + check({ request_check_roles = { "tool", "system" }, + request_check_mode = "all" })) + } + } +--- response_body +roles ["developer"]: rejected +roles []: rejected +roles ["user", "user"]: rejected +mode "latest": rejected +roles ["assistant"]: accepted +roles ["tool", "system"] + mode "all": accepted + + + +=== TEST 63: set route with fail_mode error on a protocol without role extractors +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/chat", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { "header": { "Authorization": "Bearer token" } }, + "override": { "endpoint": "http://127.0.0.1:1980/v1/chat/completions" } + }, + "ai-aws-content-moderation": { + "comprehend": { + "access_key_id": "access", + "secret_access_key": "ea+secret", + "region": "us-east-1", + "endpoint": "http://localhost:2668" + }, + "fail_mode": "error" + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 64: passthrough protocol cannot extract roles, fail_mode decides +--- request +POST /chat +{ "prompt": "I want to kill you" } +--- error_code: 500 +--- response_body_chomp +protocol passthrough cannot moderate the configured request_check_roles + + + +=== TEST 65: set route with request_check_mode "last" and assistant among the roles +--- config + location /t { + content_by_lua_block { + local t = require("lib.test_admin").test + local code, body = t('/apisix/admin/routes/1', + ngx.HTTP_PUT, + [[{ + "uri": "/chat", + "plugins": { + "ai-proxy": { + "provider": "openai", + "auth": { "header": { "Authorization": "Bearer token" } }, + "override": { "endpoint": "http://127.0.0.1:1980/v1/chat/completions" } + }, + "ai-aws-content-moderation": { + "comprehend": { + "access_key_id": "access", + "secret_access_key": "ea+secret", + "region": "us-east-1", + "endpoint": "http://localhost:2668" + }, + "request_check_mode": "last", + "request_check_roles": ["user", "assistant"], + "deny_code": 400 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 66: assistant turns no longer end the block, so "last" reaches back into history +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "I want to kill you" }, { "role": "assistant", "content": "ok" }, { "role": "user", "content": "What is 1+1?" } ] } +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ +--- grep_error_log eval +qr/comprehend text: [^,]+/ +--- grep_error_log_out +comprehend text: I want to kill you ok What is 1+1?