From 178c0be3791c84b0133d771b10c3c7cfb6976ac8 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Thu, 30 Jul 2026 18:03:53 +0545 Subject: [PATCH 1/2] fix(ai-aws-content-moderation): respect Comprehend's segment limits and reuse the client DetectToxicContent takes at most 10 text segments per call, each at most 1 KB. The plugin sent the whole prompt or the whole completion as a single segment, so any content over 1 KB came back as a 400: the request side failed closed with a 500, and the response side logged the error and let the content through, which silently disabled check_response on normal-length responses. On that path the verdict was never computed either, yet final_packet still annotated the stream with risk_level, reporting the request-side (or stale) verdict as if the response had been scored. Content is now split on UTF-8 character boundaries into segments bounded by the new request_check_length_limit / response_check_length_limit (default 1000 bytes, capped at Comprehend's 1 KB), and those segments are batched up to the per-call limits so a long body costs as few round trips as possible. When Comprehend fails, the stream is left without a risk_level rather than carrying a verdict that was never computed. The Comprehend client is also built once per request instead of once per batch, and timeout / keepalive / keepalive_timeout are exposed so the connection is reused across calls instead of a fresh TLS handshake each time. --- apisix/plugins/ai-aws-content-moderation.lua | 184 +++++++- .../plugins/ai-aws-content-moderation.md | 9 +- t/fixtures/aws/chat-long-harmful.json | 21 + .../aws/chat-streaming-long-harmful.sse | 14 + t/plugin/ai-aws-content-moderation.t | 437 +++++++++++++++++- 5 files changed, 629 insertions(+), 36 deletions(-) create mode 100644 t/fixtures/aws/chat-long-harmful.json create mode 100644 t/fixtures/aws/chat-streaming-long-harmful.sse diff --git a/apisix/plugins/ai-aws-content-moderation.lua b/apisix/plugins/ai-aws-content-moderation.lua index c40fe84b40b8..89393fef7812 100644 --- a/apisix/plugins/ai-aws-content-moderation.lua +++ b/apisix/plugins/ai-aws-content-moderation.lua @@ -25,16 +25,26 @@ local aws_instance local http = require("resty.http") -local ngx = ngx -local ngx_ok = ngx.OK -local pairs = pairs -local unpack = unpack -local type = type -local ipairs = ipairs -local table = table +local ngx = ngx +local ngx_ok = ngx.OK +local pairs = pairs +local unpack = unpack +local type = type +local ipairs = ipairs +local table = table +local str_byte = string.byte +local str_sub = string.sub local HTTP_INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR local HTTP_BAD_REQUEST = ngx.HTTP_BAD_REQUEST +-- DetectToxicContent takes at most 10 text segments per call, each at most +-- 1 KB, with the whole list capped at 10 KB. Anything over that is rejected +-- with TextSizeLimitExceededException, so content is split and batched to fit. +-- https://docs.aws.amazon.com/comprehend/latest/APIReference/API_DetectToxicContent.html +local MAX_SEGMENT_BYTES = 1024 +local MAX_SEGMENTS_PER_CALL = 10 +local MAX_CALL_BYTES = 10000 + local moderation_categories_pattern = "^(PROFANITY|HATE_SPEECH|INSULT|".. "HARASSMENT_OR_ABUSE|SEXUAL|VIOLENCE_OR_THREAT)$" local schema = { @@ -76,6 +86,22 @@ local schema = { }, check_request = { type = "boolean", default = true }, check_response = { type = "boolean", default = false }, + request_check_length_limit = { + type = "integer", + minimum = 1, + maximum = MAX_SEGMENT_BYTES, + default = 1000, + description = "max bytes of request content per Comprehend text segment; " .. + "longer content is split into several segments", + }, + response_check_length_limit = { + type = "integer", + minimum = 1, + maximum = MAX_SEGMENT_BYTES, + default = 1000, + description = "max bytes of response content per Comprehend text segment; " .. + "longer content is split into several segments", + }, stream_check_mode = { type = "string", enum = { "realtime", "final_packet" }, @@ -106,6 +132,24 @@ local schema = { "client SDKs; set a 4xx to surface denies as HTTP errors instead.", }, deny_message = { type = "string" }, + timeout = { + type = "integer", + minimum = 1, + default = 10000, + description = "Comprehend request timeout in milliseconds", + }, + keepalive = { + type = "boolean", + default = true, + description = "if true, keep the Comprehend connection alive for reuse", + }, + keepalive_timeout = { + type = "integer", + minimum = 1000, + default = 60000, + description = "idle time in milliseconds before a pooled Comprehend " .. + "connection is closed", + }, fail_mode = binding.schema_property("skip"), }, encrypt_fields = { "comprehend.secret_access_key" }, @@ -138,11 +182,15 @@ local function set_risk_level(ctx, flagged) end --- Score content with AWS Comprehend detectToxicContent. --- `subject` names the moderated text in the deny reason ("request"/"response" body). --- Returns (reason, nil) when a category/toxicity threshold is exceeded, --- (nil, err) on a service error, and (nil, nil) when the content is clean. -local function detect_toxic(conf, ctx, content, subject) +-- One Comprehend client per request: realtime moderation fires a call per +-- batch, and rebuilding the credentials and the service object every time is +-- pure overhead. The connection itself is reused across requests through the +-- keepalive pool. +local function get_comprehend_client(conf, ctx) + if ctx.aws_cm_client then + return ctx.aws_cm_client, ctx.aws_cm_endpoint + end + local comprehend = conf.comprehend if not aws_instance then @@ -160,18 +208,76 @@ local function detect_toxic(conf, ctx, content, subject) aws_instance.config.endpoint = endpoint aws_instance.config.ssl_verify = comprehend.ssl_verify - local comprehend_client = aws_instance:Comprehend({ + ctx.aws_cm_client = aws_instance:Comprehend({ credentials = credentials, endpoint = endpoint, region = comprehend.region, port = port, + timeout = conf.timeout, + keepalive_idle_timeout = conf.keepalive and conf.keepalive_timeout or nil, }) + ctx.aws_cm_endpoint = endpoint + return ctx.aws_cm_client, endpoint +end + - local res, err = comprehend_client:detectToxicContent({ +-- Last byte of the segment starting at `from`: at most `limit` bytes, cut on a +-- UTF-8 character boundary so a multibyte character is never torn in half. +local function segment_end(content, from, limit) + local last = from + limit - 1 + if last >= #content then + return #content + end + + -- 0x80..0xBF are continuation bytes: walk back to the start of the + -- character straddling the cut and end the segment just before it + local cut = last + for _ = 1, 3 do + local byte = str_byte(content, cut + 1) + if byte < 0x80 or byte >= 0xC0 then + break + end + cut = cut - 1 + end + + local byte = str_byte(content, cut + 1) + if cut < from or (byte >= 0x80 and byte < 0xC0) then + -- one character wider than the limit, or invalid UTF-8: cut on the byte + return last + end + return cut +end + + +-- Collect the next batch of text segments starting at `cursor`, staying within +-- Comprehend's per-segment, per-list and per-call size limits. +-- Returns the segments and the cursor to resume from. +local function next_batch(content, cursor, limit) + local segments = {} + local bytes = 0 + + while cursor <= #content and #segments < MAX_SEGMENTS_PER_CALL do + local stop = segment_end(content, cursor, limit) + local piece = str_sub(content, cursor, stop) + if #segments > 0 and bytes + #piece > MAX_CALL_BYTES then + break + end + segments[#segments + 1] = { Text = piece } + bytes = bytes + #piece + cursor = stop + 1 + end + + return segments, cursor +end + + +-- Score one batch of segments. Returns the deny reason when a +-- category/toxicity threshold is exceeded, (nil, err) on a service error. +local function detect_batch(conf, ctx, segments, subject) + local client, endpoint = get_comprehend_client(conf, ctx) + local res, err = client:detectToxicContent({ LanguageCode = "en", - TextSegments = {{ - Text = content - }}, + TextSegments = segments, }) if not res then return nil, "failed to send request to " .. endpoint .. ": " .. err @@ -187,17 +293,39 @@ local function detect_toxic(conf, ctx, content, subject) for _, item in pairs(result.Labels) do local threshold = conf.moderation_categories[item.Name] if threshold and item.Score > threshold then - set_risk_level(ctx, true) return subject .. " exceeds " .. item.Name .. " threshold" end end end if result.Toxicity > conf.moderation_threshold then - set_risk_level(ctx, true) return subject .. " exceeds toxicity threshold" end end +end + + +-- Score content with AWS Comprehend detectToxicContent, splitting it into +-- segments the service accepts and batching those into as few calls as +-- possible. `subject` names the moderated text in the deny reason +-- ("request"/"response" body). Returns (reason, nil) when a category/toxicity +-- threshold is exceeded, (nil, err) on a service error, and (nil, nil) when +-- the content is clean. +local function detect_toxic(conf, ctx, content, subject, length_limit) + local cursor = 1 + while cursor <= #content do + local segments + segments, cursor = next_batch(content, cursor, length_limit) + + local reason, err = detect_batch(conf, ctx, segments, subject) + if err then + return nil, err + end + if reason then + set_risk_level(ctx, true) + return reason + end + end set_risk_level(ctx, false) end @@ -233,14 +361,17 @@ local function moderate_response(ctx, conf, content) -- nothing to score, but keep the risk_level contract satisfied so the -- streamed annotation still reports a verdict set_risk_level(ctx, false) + ctx.aws_cm_response_scored = true return end - local reason, err = detect_toxic(conf, ctx, content, "response body") + local reason, err = detect_toxic(conf, ctx, content, "response body", + conf.response_check_length_limit) if err then core.log.error(err) return nil, nil, err end + ctx.aws_cm_response_scored = true if reason then return conf.deny_code, build_deny_message(ctx, conf, reason) end @@ -255,11 +386,19 @@ local function annotate_stream(ctx, body) return end + -- Comprehend failed, so there is no verdict for this response. Leaving + -- risk_level off is the honest outcome: annotating would report the + -- request-side verdict, or a stale one, as if the response was scored. + local scored = ctx.aws_cm_response_scored + if not scored then + core.log.warn("response was not scored, streaming without a risk_level annotation") + end + local events = sse.decode(body) local raw_events = {} local contains_done_event = false for _, event in ipairs(events) do - if proto.is_data_event(event) then + if scored and proto.is_data_event(event) then local data, err = core.json.decode(event.data) -- a scalar or cjson.null decode can't be indexed; only annotate -- well-formed JSON object frames and leave anything else untouched @@ -341,7 +480,8 @@ function _M.access(conf, ctx) return end - local reason, err = detect_toxic(conf, ctx, content, "request body") + local reason, err = detect_toxic(conf, ctx, content, "request body", + conf.request_check_length_limit) if err then core.log.error(err) return HTTP_INTERNAL_SERVER_ERROR, err diff --git a/docs/en/latest/plugins/ai-aws-content-moderation.md b/docs/en/latest/plugins/ai-aws-content-moderation.md index dcddf60ff0d8..9bb49d6d95af 100644 --- a/docs/en/latest/plugins/ai-aws-content-moderation.md +++ b/docs/en/latest/plugins/ai-aws-content-moderation.md @@ -42,7 +42,9 @@ The Plugin is protocol-aware: it extracts the prompt content from the LLM reques 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. -If AWS Comprehend cannot be reached, the request and the buffered (non-streaming) response both fail closed with a `500`, so unmoderated content is never proxied. Streaming response moderation is best-effort: once the first bytes have been sent to the client the response cannot be blocked, so a Comprehend failure there lets the remaining stream through. +AWS Comprehend accepts at most 10 text segments per call, each at most 1 KB. Content longer than `request_check_length_limit` / `response_check_length_limit` is therefore split on character boundaries and the segments are batched into as few calls as possible. + +If AWS Comprehend cannot be reached, the request and the buffered (non-streaming) response both fail closed with a `500`, so unmoderated content is never proxied. Streaming response moderation is best-effort: once the first bytes have been sent to the client the response cannot be blocked, so a Comprehend failure there lets the remaining stream through and the stream is left without a `risk_level` annotation. The `ai-aws-content-moderation` Plugin should be used with either [`ai-proxy`](./ai-proxy.md) or [`ai-proxy-multi`](./ai-proxy-multi.md) Plugin for proxying LLM requests. @@ -60,11 +62,16 @@ 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_length_limit` | integer | False | `1000` | [1, 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` | [1, 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`. | | `stream_check_cache_size` | integer | False | `128` | >= 1 | Maximum characters per moderation batch in `realtime` mode. | | `stream_check_interval` | number | False | `3` | >= 0.1 | Seconds between batch checks in `realtime` mode. | | `deny_code` | integer | False | `200` | [200, 599] | HTTP status code returned when a request is rejected. Defaults to `200` so the provider-compatible refusal parses as a normal completion in client SDKs; set a 4xx to surface denies as HTTP errors instead. Streaming responses denied mid-stream keep the status already sent to the client. | | `deny_message` | string | False | | | Message returned when a request or response is rejected. If unset, the moderation reason (for example `request body exceeds toxicity threshold`) is returned. | +| `timeout` | integer | False | `10000` | >= 1 | Comprehend request timeout in milliseconds. | +| `keepalive` | boolean | False | `true` | | If `true`, keep the Comprehend connection alive for reuse. | +| `keepalive_timeout` | integer | False | `60000` | >= 1000 | Idle time in milliseconds before a pooled Comprehend connection is closed. | | `fail_mode` | string | False | `skip` | `skip`, `warn`, `error` | Behavior when the request did not pass through `ai-proxy`/`ai-proxy-multi` and therefore cannot be moderated as an AI request. `skip`: let the request pass through unchecked; `warn`: pass through and log a warning; `error`: reject the request. | ## Examples diff --git a/t/fixtures/aws/chat-long-harmful.json b/t/fixtures/aws/chat-long-harmful.json new file mode 100644 index 000000000000..13803422c61d --- /dev/null +++ b/t/fixtures/aws/chat-long-harmful.json @@ -0,0 +1,21 @@ +{ + "choices": [ + { + "finish_reason": "stop", + "index": 0, + "message": { + "content": "The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is cI will kill you.", + "role": "assistant" + } + } + ], + "created": 1723780938, + "id": "chatcmpl-9wiSIg5LYrrpxwsr2PubSQnbtod1P", + "model": "gpt-3.5-turbo", + "object": "chat.completion", + "usage": { + "completion_tokens": 300, + "prompt_tokens": 8, + "total_tokens": 308 + } +} diff --git a/t/fixtures/aws/chat-streaming-long-harmful.sse b/t/fixtures/aws/chat-streaming-long-harmful.sse new file mode 100644 index 000000000000..e0f1488764e2 --- /dev/null +++ b/t/fixtures/aws/chat-streaming-long-harmful.sse @@ -0,0 +1,14 @@ +data: {"id":"chatcmpl-aws2","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]} + +data: {"id":"chatcmpl-aws2","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads ar"},"finish_reason":null}]} + +data: {"id":"chatcmpl-aws2","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads ar"},"finish_reason":null}]} + +data: {"id":"chatcmpl-aws2","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads are clear for travel. The weather today is calm and the roads ar"},"finish_reason":null}]} + +data: {"id":"chatcmpl-aws2","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":"and then I will kill you right now!"},"finish_reason":null}]} + +data: {"id":"chatcmpl-aws2","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":400,"total_tokens":410}} + +data: [DONE] + diff --git a/t/plugin/ai-aws-content-moderation.t b/t/plugin/ai-aws-content-moderation.t index 01872c5101ad..37b01e52ef24 100644 --- a/t/plugin/ai-aws-content-moderation.t +++ b/t/plugin/ai-aws-content-moderation.t @@ -36,8 +36,9 @@ _EOC_ $block->set_value("main_config", $main_config); - # Mock AWS Comprehend detectToxicContent: looks up the extracted text - # (TextSegments[1].Text) in the canned responses fixture. + # Mock AWS Comprehend detectToxicContent: looks up each extracted text + # segment in the canned responses fixture and answers with one result per + # segment, the way the real service does. my $http_config = $block->http_config // <<_EOC_; server { listen 2668; @@ -80,20 +81,53 @@ _EOC_ ngx.status(503) ngx.say("[INTERNAL FAILURE]: failed to decoded request body: ", err) end - local result = body.TextSegments[1].Text - local final_response = responses[result] - - -- Response-side text is free-form LLM output, not a fixture - -- key, so fall back to flagging anything violent. - if not final_response then - if result:find("kill", 1, true) then - final_response = responses["toxic"] - else - final_response = responses["good_request"] + local segments = body.TextSegments + ngx.log(ngx.WARN, "comprehend call: segments=", #segments, + " connection_requests=", ngx.var.connection_requests) + + -- Comprehend takes at most 10 segments of 1 KB each, 10 KB + -- for the list; over that it answers 400. Enforce it here so + -- the tests catch content that is sent unsplit. + local total = 0 + for _, segment in ipairs(segments) do + total = total + #segment.Text + end + if #segments > 10 or total > 10240 then + ngx.status = 400 + ngx.say(json.encode({ + __type = "TextSizeLimitExceededException", + Message = "Input text size exceeds limit." + })) + return + end + + local results = {} + for i, segment in ipairs(segments) do + local text = segment.Text + if #text > 1024 then + ngx.status = 400 + ngx.say(json.encode({ + __type = "TextSizeLimitExceededException", + Message = "Input text size exceeds limit." + })) + return + end + + local canned = responses[text] + + -- Response-side text is free-form LLM output, not a + -- fixture key, so fall back to flagging anything violent. + if not canned then + if text:find("kill", 1, true) then + canned = responses["toxic"] + else + canned = responses["good_request"] + end end + results[i] = canned.ResultList[1] end - ngx.say(json.encode(final_response)) + ngx.say(json.encode({ ResultList = results })) } } } @@ -871,3 +905,380 @@ qr/code:nil\nrisk_level:none\n.*"risk_level":"none"/s qr/code:nil\n.*data: 12345.*"risk_level":"none"/s --- no_error_log [error] + + + +=== TEST 33: set route with the default check length limits +--- 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 34: request over 1 KB is split into segments Comprehend accepts +--- request eval +"POST /chat +{ \"messages\": [ { \"role\": \"user\", \"content\": \"" . ("a" x 2395) . " kill\" } ] }" +--- error_code: 400 +--- response_body_like eval +qr/request body exceeds toxicity threshold/ +--- grep_error_log eval +qr/comprehend call: segments=\d+/ +--- grep_error_log_out +comprehend call: segments=3 + + + +=== TEST 35: clean request over 1 KB passes +--- request eval +"POST /chat +{ \"messages\": [ { \"role\": \"user\", \"content\": \"" . ("a" x 2400) . "\" } ] }" +--- error_code: 200 +--- grep_error_log eval +qr/comprehend call: segments=\d+/ +--- grep_error_log_out +comprehend call: segments=3 + + + +=== TEST 36: set route with a check length limit that needs several batches +--- 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_length_limit": 100 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 37: segments past the tenth go in a second call on the same connection +--- request eval +"POST /chat +{ \"messages\": [ { \"role\": \"user\", \"content\": \"" . ("a" x 1100) . "\" } ] }" +--- error_code: 200 +--- grep_error_log eval +qr/comprehend call: segments=\d+ connection_requests=\d+/ +--- grep_error_log_out +comprehend call: segments=10 connection_requests=1 +comprehend call: segments=1 connection_requests=2 + + + +=== TEST 38: same route with keepalive turned off +--- 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_length_limit": 100, + "keepalive": false + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 39: without keepalive every batch opens a new connection +--- request eval +"POST /chat +{ \"messages\": [ { \"role\": \"user\", \"content\": \"" . ("a" x 1100) . "\" } ] }" +--- error_code: 200 +--- grep_error_log eval +qr/comprehend call: segments=\d+ connection_requests=\d+/ +--- grep_error_log_out +comprehend call: segments=10 connection_requests=1 +comprehend call: segments=1 connection_requests=1 + + + +=== TEST 40: set route with check_response enabled (non-streaming) +--- 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" + }, + "check_request": false, + "check_response": true, + "deny_code": 400 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 41: toxic LLM response over 1 KB is denied +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "good_request" } ] } +--- more_headers +X-AI-Fixture: aws/chat-long-harmful.json +--- error_code: 400 +--- response_body_like eval +qr/response body exceeds toxicity threshold/ +--- grep_error_log eval +qr/comprehend call: segments=\d+/ +--- grep_error_log_out +comprehend call: segments=2 + + + +=== TEST 42: set route with stream_check_mode = realtime and a batch over 1 KB +--- 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" + }, + "check_request": false, + "check_response": true, + "deny_message": "the response was withheld", + "stream_check_mode": "realtime", + "stream_check_cache_size": 1500 + } + } + }]] + ) + + if code >= 300 then + ngx.status = code + end + ngx.say(body) + } + } +--- response_body +passed + + + +=== TEST 43: a realtime batch over 1 KB is still moderated +--- request +POST /chat +{ "messages": [ { "role": "user", "content": "good_request" } ], "stream": true } +--- more_headers +X-AI-Fixture: aws/chat-streaming-long-harmful.sse +X-AI-Fixture-Flush-Events: true +--- error_code: 200 +--- response_body_like eval +qr/the response was withheld/ +--- error_log +comprehend call: segments=2 +--- no_error_log +failed to get moderation results from response + + + +=== TEST 44: final_packet does not annotate a verdict it could not compute +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-aws-content-moderation") + local ctx = { + picked_ai_instance = { provider = "openai" }, + ai_client_protocol = "openai-chat", + var = { + request_type = "ai_stream", + llm_response_text = "hello there", + llm_request_done = true, + -- the request check already ran and found nothing + llm_content_risk_level = "none", + }, + } + local conf = { + comprehend = { + access_key_id = "access", + secret_access_key = "secret", + region = "us-east-1", + -- nothing listening, so the response is never scored + endpoint = "http://localhost:2669" + }, + check_response = true, + stream_check_mode = "final_packet", + } + plugin.check_schema(conf) + local body = 'data: {"choices":[{"delta":{"content":"hi"}}]}\n\n' + local code, new_body = plugin.lua_body_filter(conf, ctx, {}, body) + ngx.say("code:", code or "nil") + ngx.say("annotated:", new_body:find("risk_level", 1, true) and "yes" or "no") + } + } +--- response_body +code:nil +annotated:no +--- error_log +response was not scored, streaming without a risk_level annotation + + + +=== TEST 45: schema check: length limits are bounded by Comprehend's 1 KB cap +--- config + location /t { + content_by_lua_block { + local plugin = require("apisix.plugins.ai-aws-content-moderation") + local function conf(extra) + local c = { + comprehend = { + access_key_id = "a", + secret_access_key = "s", + region = "us-east-1" + } + } + for k, v in pairs(extra or {}) do + c[k] = v + end + return c + end + + ngx.say("request limit over 1 KB: ", + plugin.check_schema(conf({request_check_length_limit = 1025})) + and "accepted" or "rejected") + ngx.say("response limit over 1 KB: ", + plugin.check_schema(conf({response_check_length_limit = 1025})) + and "accepted" or "rejected") + ngx.say("zero request limit: ", + plugin.check_schema(conf({request_check_length_limit = 0})) + and "accepted" or "rejected") + ngx.say("1 KB limit: ", + plugin.check_schema(conf({response_check_length_limit = 1024})) + and "accepted" or "rejected") + + local c = conf() + plugin.check_schema(c) + ngx.say("defaults: ", c.request_check_length_limit, " ", + c.response_check_length_limit, " ", c.timeout, " ", + c.keepalive and "on" or "off", " ", c.keepalive_timeout) + } + } +--- response_body +request limit over 1 KB: rejected +response limit over 1 KB: rejected +zero request limit: rejected +1 KB limit: accepted +defaults: 1000 1000 10000 on 60000 From a0211bb54570170f5522a294219a5af522c00301 Mon Sep 17 00:00:00 2001 From: Abhishek Choudhary Date: Fri, 31 Jul 2026 09:27:18 +0545 Subject: [PATCH 2/2] fix(ai-aws-content-moderation): address review - pass ssl_verify and the endpoint through the Comprehend service config instead of mutating the shared aws instance - require at least 4 bytes per segment, the width of the widest UTF-8 character, so a segment can always hold one whole character - have the Comprehend mock reject segments that are not valid UTF-8, and cover a multibyte body that spans several segments --- apisix/plugins/ai-aws-content-moderation.lua | 15 +++-- .../plugins/ai-aws-content-moderation.md | 4 +- t/plugin/ai-aws-content-moderation.t | 66 +++++++++++++++++++ 3 files changed, 78 insertions(+), 7 deletions(-) diff --git a/apisix/plugins/ai-aws-content-moderation.lua b/apisix/plugins/ai-aws-content-moderation.lua index 89393fef7812..d06c9e8a0f69 100644 --- a/apisix/plugins/ai-aws-content-moderation.lua +++ b/apisix/plugins/ai-aws-content-moderation.lua @@ -41,9 +41,14 @@ local HTTP_BAD_REQUEST = ngx.HTTP_BAD_REQUEST -- 1 KB, with the whole list capped at 10 KB. Anything over that is rejected -- with TextSizeLimitExceededException, so content is split and batched to fit. -- https://docs.aws.amazon.com/comprehend/latest/APIReference/API_DetectToxicContent.html +-- MAX_CALL_BYTES keeps a margin under the 10 KB list cap: AWS does not say +-- whether the cap counts the encoded list or just the text, and an extra call +-- is much cheaper than a rejected one that leaves content unmoderated. local MAX_SEGMENT_BYTES = 1024 local MAX_SEGMENTS_PER_CALL = 10 local MAX_CALL_BYTES = 10000 +-- a UTF-8 character is up to 4 bytes, so a smaller segment could not hold one +local MIN_SEGMENT_BYTES = 4 local moderation_categories_pattern = "^(PROFANITY|HATE_SPEECH|INSULT|".. "HARASSMENT_OR_ABUSE|SEXUAL|VIOLENCE_OR_THREAT)$" @@ -88,7 +93,7 @@ local schema = { check_response = { type = "boolean", default = false }, request_check_length_limit = { type = "integer", - minimum = 1, + minimum = MIN_SEGMENT_BYTES, maximum = MAX_SEGMENT_BYTES, default = 1000, description = "max bytes of request content per Comprehend text segment; " .. @@ -96,7 +101,7 @@ local schema = { }, response_check_length_limit = { type = "integer", - minimum = 1, + minimum = MIN_SEGMENT_BYTES, maximum = MAX_SEGMENT_BYTES, default = 1000, description = "max bytes of response content per Comprehend text segment; " .. @@ -205,14 +210,14 @@ local function get_comprehend_client(conf, ctx) local default_endpoint = "https://comprehend." .. comprehend.region .. ".amazonaws.com" local scheme, host, port = unpack(http:parse_uri(comprehend.endpoint or default_endpoint)) local endpoint = scheme .. "://" .. host - aws_instance.config.endpoint = endpoint - aws_instance.config.ssl_verify = comprehend.ssl_verify + -- aws_instance is shared, so keep per-route settings on the service config ctx.aws_cm_client = aws_instance:Comprehend({ credentials = credentials, endpoint = endpoint, region = comprehend.region, port = port, + ssl_verify = comprehend.ssl_verify, timeout = conf.timeout, keepalive_idle_timeout = conf.keepalive and conf.keepalive_timeout or nil, }) @@ -242,7 +247,7 @@ local function segment_end(content, from, limit) local byte = str_byte(content, cut + 1) if cut < from or (byte >= 0x80 and byte < 0xC0) then - -- one character wider than the limit, or invalid UTF-8: cut on the byte + -- the content is not valid UTF-8: cut on the byte return last end return cut diff --git a/docs/en/latest/plugins/ai-aws-content-moderation.md b/docs/en/latest/plugins/ai-aws-content-moderation.md index 9bb49d6d95af..fe43939b115d 100644 --- a/docs/en/latest/plugins/ai-aws-content-moderation.md +++ b/docs/en/latest/plugins/ai-aws-content-moderation.md @@ -62,8 +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_length_limit` | integer | False | `1000` | [1, 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` | [1, 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. | +| `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`. | | `stream_check_cache_size` | integer | False | `128` | >= 1 | Maximum characters per moderation batch in `realtime` mode. | | `stream_check_interval` | number | False | `3` | >= 0.1 | Seconds between batch checks in `realtime` mode. | diff --git a/t/plugin/ai-aws-content-moderation.t b/t/plugin/ai-aws-content-moderation.t index 37b01e52ef24..abee14dd1222 100644 --- a/t/plugin/ai-aws-content-moderation.t +++ b/t/plugin/ai-aws-content-moderation.t @@ -81,6 +81,7 @@ _EOC_ ngx.status(503) ngx.say("[INTERNAL FAILURE]: failed to decoded request body: ", err) end + local utf8 = require("lua-utf8") local segments = body.TextSegments ngx.log(ngx.WARN, "comprehend call: segments=", #segments, " connection_requests=", ngx.var.connection_requests) @@ -104,6 +105,9 @@ _EOC_ local results = {} for i, segment in ipairs(segments) do local text = segment.Text + if not utf8.len(text) then + ngx.log(ngx.ERR, "comprehend got a segment that is not valid utf-8") + end if #text > 1024 then ngx.status = 400 ngx.say(json.encode({ @@ -1265,6 +1269,13 @@ response was not scored, streaming without a risk_level annotation ngx.say("zero request limit: ", plugin.check_schema(conf({request_check_length_limit = 0})) and "accepted" or "rejected") + -- a segment narrower than 4 bytes could not hold one UTF-8 character + ngx.say("3 byte request limit: ", + plugin.check_schema(conf({request_check_length_limit = 3})) + and "accepted" or "rejected") + ngx.say("4 byte response limit: ", + plugin.check_schema(conf({response_check_length_limit = 4})) + and "accepted" or "rejected") ngx.say("1 KB limit: ", plugin.check_schema(conf({response_check_length_limit = 1024})) and "accepted" or "rejected") @@ -1280,5 +1291,60 @@ response was not scored, streaming without a risk_level annotation request limit over 1 KB: rejected response limit over 1 KB: rejected zero request limit: rejected +3 byte request limit: rejected +4 byte response limit: accepted 1 KB limit: accepted defaults: 1000 1000 10000 on 60000 + + + +=== TEST 46: set route with the default check length limits +--- 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 47: multibyte content is split on character boundaries +--- request eval +"POST /chat +{ \"messages\": [ { \"role\": \"user\", \"content\": \"" . ("\xe6\x97\xa5" x 400) . "\" } ] }" +--- error_code: 200 +--- grep_error_log eval +qr/comprehend call: segments=\d+/ +--- grep_error_log_out +comprehend call: segments=2 +--- no_error_log +not valid utf-8