diff --git a/apisix/plugins/ai-aws-content-moderation.lua b/apisix/plugins/ai-aws-content-moderation.lua index c40fe84b40b8..d06c9e8a0f69 100644 --- a/apisix/plugins/ai-aws-content-moderation.lua +++ b/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 diff --git a/docs/en/latest/plugins/ai-aws-content-moderation.md b/docs/en/latest/plugins/ai-aws-content-moderation.md index dcddf60ff0d8..fe43939b115d 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` | [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 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..abee14dd1222 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,57 @@ _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 utf8 = require("lua-utf8") + 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 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({ + __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 +909,442 @@ 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") + -- 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") + + 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 +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