fix(ai-aws-content-moderation): respect Comprehend's segment limits and reuse the client

apisix/plugins/ai-aws-content-moderation.lua

@@ -25,16 +25,31 @@ 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
+-- 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)$"
 local schema = {
@@ -76,6 +91,22 @@ local schema = {
         },
         check_request = { type = "boolean", default = true },
         check_response = { type = "boolean", default = false },
+        request_check_length_limit = {
+            type = "integer",
+            minimum = MIN_SEGMENT_BYTES,
+            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 = MIN_SEGMENT_BYTES,
+            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 +137,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 +187,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
@@ -157,21 +210,79 @@ local function detect_toxic(conf, ctx, content, subject)
     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
 
-    local comprehend_client = aws_instance:Comprehend({
+    -- 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,
     })
+    ctx.aws_cm_endpoint = endpoint
+    return ctx.aws_cm_client, endpoint
+end
+
+
+-- 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
 
-    local res, err = comprehend_client:detectToxicContent({
+    -- 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
+        -- the content is not valid 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 +298,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 +366,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 +391,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 +485,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

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` | [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. |
 | `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

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
+  }
+}