From d4dc2c39e7ab3f560cd67d0e974629fc2fab79a6 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 7 Aug 2026 19:44:24 -0700 Subject: [PATCH] fix(guardrails): chunk oversized Bedrock ApplyGuardrail requests instead of failing (#36119) * feat(guardrails): chunk oversized Bedrock ApplyGuardrail requests instead of failing AWS's ApplyGuardrail API rejects requests whose content exceeds the account's per-request "maximum input size in text units" quota with a 400 ValidationException. That cap is account/region/policy-dependent and cannot be predicted from config, so it can only be reacted to. _make_apply_guardrail_request now tries the whole-content call first (no behavior change for requests that already fit). On a too-large ValidationException it bisects the flat content list and retries each half sequentially, recursing until every piece fits or cannot be split further, then merges the per-chunk responses (action, assessments, outputs, usage) into one so callers cannot tell chunking happened. A real guardrail block on any (sub-)chunk still raises immediately. Contextual-grounding requests are never chunked: grounding scores the response holistically against the whole reference source, so fragmenting it would produce misleading scores. Each chunk call also gets a small exponential backoff retry on AWS ThrottlingException (429), since chunking increases the number of per-second API calls and can trade a 400 for a 429. All new state is local to a single request's call stack (no shared cache, no cross-process coordination), so this is safe for single-pod, multi-pod, and cache-less LiteLLM proxy deployments alike. * fix(guardrails): address Bedrock ApplyGuardrail chunking review feedback Fixes three issues flagged in review of the chunking fallback: a single oversized content item couldn't be split (only list-length bisection was supported), a chunked request that got recovered still logged a stray failure telemetry entry alongside the real outcome, and flattening chunk outputs without positional bookkeeping could misalign masked text onto the wrong original message once a chunk had nothing to mask. * test(guardrails): add regression test for multi-level Bedrock guardrail chunking Confirms the too-large bisection recursion isn't capped at a single split: a payload that is still oversized after the first halving keeps splitting until every piece fits, converging on however many chunks it takes rather than only ever producing two. * fix(guardrails): hybrid bin-pack+bisection chunking, whitespace-safe splits Rework Bedrock ApplyGuardrail chunking from pure reactive bisection to a hybrid strategy: bin-pack content into fixed-budget batches up front as the fast path, falling back to the existing recursive bisection only for a batch AWS still rejects as too large. Avoids paying O(log n) round trips on every oversized request when a single pass would do. Also switch single-item text splitting from a raw character midpoint to the nearest whitespace boundary, so a fragment never starts or ends mid-word. Closes the accidental-severing case from review; the residual gap (a multi-word denied phrase deliberately straddling the boundary) is documented as an accepted limitation, since fixing it would require an overlap window reconciled against masked output with no documented length-preservation guarantee from AWS. * chore(ui): regenerate dashboard API types * fix(guardrails): don't retry an oversized Bedrock guardrail call as a throttle AWS reports an ApplyGuardrail request that exceeds the per-request text-unit cap as a ThrottlingException (429), not only as the documented ValidationException (400). Verified against a live guardrail with an active content-filter policy: a 3273-text-unit request comes back as "Input text size (3273 text units) exceeds the maximum allowed (1000 text units) for the content filter policy (Classic tier)". The throttle retry keyed off status 429 alone, so every oversized chunk burned the full backoff-retry budget - each attempt a billed AWS call preceded by a sleep - before the bisection fallback got a chance, at every level of the recursion. A size error is not transient; re-posting the same content can never succeed. It now short-circuits straight to bisection. Also rename _is_input_too_large_validation_error to _is_input_too_large_error (it never keyed off the status code, and the error is not always a ValidationException), correct the docstrings that asserted a 400, and log at warning level when a split happens so the recovery is visible without --detailed_debug. * Revert "chore(ui): regenerate dashboard API types" This reverts commit ebf8ba2fd57f13bccf7aa6c5dfcac41c74db1ed9. * fix(guardrails): group all fragments of one item and stop double-logging Two defects found in review, both invisible to the existing tests. Fragment grouping assumed a split content item always produces exactly two adjacent fragments. That holds for one bisection level but not two: an item split twice yields four fragments, which were regrouped in fixed pairs into two output entries for a single message. Since masking walks the merged outputs by a running index across the original, unchunked message list, that message was written back truncated to its first half and every later message shifted. Fragments now carry the size of the group they belong to, so any number of them collapse back into exactly one output entry. Telemetry was also double-counted. AsyncHTTPHandler.post calls raise_for_status(), so every non-200 from Bedrock reaches _sign_and_post's error path, which logged guardrail_failed_to_respond before re-raising as an HTTPException that the consolidating caller then logged again. A request recovered by chunking reported one failure per rejected attempt plus a success. The ApplyGuardrail path now opts out of that per-attempt logging, since it owns consolidated per-request logging; the connection-level branch still logs, as nothing else records it. The existing tests missed both because their mocks return a non-200 response object, while the real client raises. Added a helper that raises a genuine httpx.HTTPStatusError so these paths are covered the way production hits them, plus a case asserting an unrecoverable failure still logs exactly once rather than zero times. * refactor(guardrails): move Bedrock chunking rationale into docstrings The chunking work explained itself with inline comment blocks, which this repo's conventions do not want. Folded that reasoning into the docstrings of the functions it describes and dropped the comments, including the module-level constant blocks and the test-file banner. No behavior change. The banner also claimed AWS rejects an oversized request with a 400 ValidationException, which live testing disproved, so removing it drops a stale claim as well as an internal ticket reference from a public repo. * feat(guardrails): match AWS default chunk budget and make it configurable ApplyGuardrail's default quota is 25 text units, roughly 25,000 characters, per second. Chunking has to respect that throughput limit rather than just the per-request size, otherwise splitting an oversized request trades a size error for a throttle. The budget now defaults to 25,000 to match that default for every user, up from an arbitrary 20,000. Accounts with raised quotas can spend fewer calls by setting chunk_budget_chars on the guardrail. A value AWS still rejects as too large is bisected automatically, so an over-large setting costs an extra round trip rather than failing the request. * fix(guardrails): never split a Bedrock text into an empty fragment _nearest_whitespace_split_index could return len(text) when the only space at or after the midpoint was the final character, so the first fragment came back identical to the text AWS had just rejected as too large and the second came back empty. AWS rejects the unchanged fragment again, and each retry re-splits it into the same fragment, so an oversized single item shaped like a long unbroken token with one trailing space exhausted the stack with a RecursionError instead of scanning or surfacing Bedrock's error. Candidate boundaries that would leave either side empty are now discarded, and the raw midpoint is used when none remain. The midpoint is always safe because _split_bedrock_content only calls this for text of at least two characters. * style(guardrails): move chunking rationale out of comments and into docstrings * fix(guardrails): raise 500 when Bedrock reports a failure inside a 200 body Also types the credentials parameter on the new chunking helpers and rebuilds fragment grouping without mutating a list or rebinding an index * fix(guardrails): raise 500 when Bedrock reports a failure inside a 200 body Restores the source changes intended for a08e4cf309, which landed with only the test. Also types the credentials parameter on the new chunking helpers and rebuilds fragment grouping without mutating a list or rebinding an index * style(guardrails): sort the constants import into the first-party block * refactor(guardrails): bring the Bedrock chunking path under the LIT lint budgets Annotates never-rebound locals with Final, replaces the retry counter and the two branch-assigned locals with single bindings, and moves the internal chunking chain to Sequence parameters and tuple returns. Collections that reach the logged payload stay lists on purpose: redact_nested_match_and_regex_keys only traverses dict and list, so a tuple would carry PII past redaction. The remaining constructions are contract-bound and carry inline reasons * fix(guardrails): keep the pre-chunking contract for failures reported inside a 200 Reverts the 500 this branch introduced for an AWS 200 whose body carries an Output.__type exception marker: the request proceeds as it did before chunking existed. The logged status is now derived from the merged response instead of being hardcoded to success, so that shape is still reported as guardrail_failed_to_respond. The consolidated failure logger also goes back to logging a dict rather than a bare string, matching both the pre-chunking code and the InvokeGuardrailChecks path in this file * docs(guardrails): correct the docstring for failures reported inside a 200 body The raise was reverted, so the docstring no longer describes the code. Records that the request proceeds by design and points at LIT-5338 for closing the fail-open path behind the existing unreachable_fallback setting --------- Co-authored-by: spencer-burridge <265588760+spencer-burridge@users.noreply.github.com> --- litellm/constants.py | 1 + .../guardrail_hooks/bedrock_guardrails.py | 857 ++++++- .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 10 + .../guardrail_hooks/bedrock_guardrails.py | 3 + .../test_bedrock_guardrails.py | 1989 ++++++++++++----- .../proxy/guardrails/test_init_guardrails.py | 35 + 7 files changed, 2326 insertions(+), 570 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 6f0e9e7afe2..30d3bb1f26e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -280,6 +280,7 @@ TOOL_POLICY_CACHE_TTL_SECONDS: Final = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECO GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS: Final = int( os.getenv("GUARDRAIL_SCANNED_MESSAGES_CACHE_TTL_SECONDS", 24 * 60 * 60) ) +BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS: Final = 25_000 # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE: Final = int(os.getenv("MAX_SIZE_IN_MEMORY_QUEUE", int(LITELLM_ASYNCIO_QUEUE_MAXSIZE * 0.8))) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index e9e729fb118..eecbce57468 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -9,11 +9,14 @@ import os import sys sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path +import asyncio import copy import json +import re import sys -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Mapping, Sequence from datetime import datetime, timezone +from itertools import accumulate, groupby from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast import httpx @@ -23,6 +26,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.exceptions import ModifyResponseException from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys @@ -46,6 +50,7 @@ from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockGuardrailOutput, BedrockGuardrailQualifier, BedrockGuardrailResponse, + BedrockGuardrailUsage, BedrockRequest, BedrockTextContent, ) @@ -53,6 +58,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from botocore.awsrequest import AWSPreparedRequest + from botocore.credentials import Credentials from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -71,6 +77,17 @@ from litellm.types.utils import ( GUARDRAIL_NAME: Final = "bedrock" _BEDROCK_DYNAMIC_BODY_DENYLIST: Final = frozenset({"content", "source"}) +_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = ( + "text unit", + "maximum input size", + "content size", + "too long", + "too large", + "exceeds the maximum", +) +_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3 +_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5 +_BEDROCK_WHITESPACE: Final = re.compile(r"\s") # Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required). _BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke" # InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with @@ -118,6 +135,29 @@ class GuardrailMessageFilterResult(NamedTuple): target_indices: list[int] | None +class BedrockContentChunkResult(NamedTuple): + """One chunk's ApplyGuardrail response, paired with enough bookkeeping to + reconstruct global masked-output positions once every chunk is back. + + `content` is the exact content items this chunk was called with -- needed + so an all-clear chunk (empty `outputs`) can still contribute one unmasked + placeholder per item it covers, keeping every later chunk's masked text + aligned to its original global position. `fragment_group_size` is 1 for an + ordinary chunk, and otherwise the total number of consecutive chunk results + that together make up ONE original content item's own text (split because a + list of length 1 could not be bisected by list length). All of them must be + concatenated back into that one item's masked output rather than treated as + separate items. It is a count rather than a boolean because one item can be + bisected more than once: two levels of splitting produce four fragments for + a single item, not two, and grouping them in fixed pairs would emit two + outputs for one message and shift every later message's masked text. + """ + + response: BedrockGuardrailResponse + content: tuple[BedrockContentItem, ...] + fragment_group_size: int + + class ApplyGuardrailMessageSelection(NamedTuple): """Messages selected for an apply_guardrail scan + write-back metadata.""" @@ -168,12 +208,14 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): content_filter_threshold: float | None = 0.5, prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, + chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" + self.chunk_budget_chars = chunk_budget_chars self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only")) # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks` @@ -759,12 +801,35 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None = None, logging_event_type: GuardrailEventHooks | None = None, ) -> BedrockGuardrailResponse: + """Scan `messages`/`response` with ApplyGuardrail, chunking if it is too large. + + Content is bin-packed into budget-sized batches and each batch posted + sequentially, every batch independently falling back to bisection if AWS + rejects it. The per-batch responses are merged so callers cannot tell whether + chunking happened. + + Content using contextual grounding opts out of chunking entirely: grounding is + scored holistically against the whole reference source, so bisecting it would + fragment that evaluation and yield misleading scores. Such a request keeps the + old behavior of surfacing a too-large error rather than being split. + + `logging_event_type` drives what UI and spend logs report. It is distinct from + Bedrock's `source`, which is INPUT vs OUTPUT for the API body and must not be + confused with the proxy hook (pre_call / during_call / post_call); when omitted, + the legacy source-derived mapping is kept for backward compatibility. + + A guardrail *block* is logged where it happens, in + `_post_apply_guardrail_content`, because chunking stops immediately and there is + no later merged response to log instead. Everything else that fails out of the + chunking flow (an unrecoverable too-large error, a non-size validation error, + exhausted throttle retries) is a genuine end-to-end failure of this one logical + guardrail call and is logged exactly once here. + """ start_time: Final = datetime.now(timezone.utc) credentials, aws_region_name = self._load_credentials() bedrock_request_data: Final[dict] = dict( self.convert_to_bedrock_format(source=source, messages=messages, response=response) ) - bedrock_guardrail_response: BedrockGuardrailResponse = BedrockGuardrailResponse() api_key: str | None = None if request_data: dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data) @@ -778,6 +843,257 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): if request_data.get("api_key") is not None: api_key = request_data["api_key"] + event_type: Final = ( + logging_event_type + if logging_event_type is not None + else (GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call) + ) + + content: Final[tuple[BedrockContentItem, ...]] = tuple(bedrock_request_data.get("content") or ()) + allow_chunking: Final = not self._content_uses_contextual_grounding(content) + + try: + responses: Final = await self._apply_guardrail_content_with_chunking( + content=content, + base_request_data=bedrock_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + except HTTPException as exc: + if not isinstance(exc.detail, dict): + self._log_apply_guardrail_failure( + detail=exc.detail, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + raise + merged_response: Final = self._merge_bedrock_guardrail_responses(responses) + self._log_apply_guardrail_success( + merged_response=merged_response, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + return merged_response + + async def _apply_guardrail_content_with_chunking( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + allow_chunking: bool, + ) -> tuple[BedrockContentChunkResult, ...]: + """Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large. + + Tries `content` as a single call first. AWS's per-request "maximum input + size in text units" quota is account/region/policy-dependent and cannot be + predicted ahead of time, so it is only ever discovered reactively: on an + error whose message indicates the input was too large (a ThrottlingException + in practice, a ValidationException per the docs -- see + ``_is_input_too_large_error``), the content is re-sent in smaller pieces. + + Probing with the whole payload first is what keeps a request AWS would have + accepted at exactly one call. Packing into fixed batches up front instead + would split conversations AWS was happy to take whole, multiplying billed + calls and guardrail latency on traffic that never had a size problem, and + no fixed budget can avoid that because the real cap is unknown here. + + Once a rejection proves the payload is over the cap, a multi-item payload is + re-sent as ``chunk_budget_chars``-sized batches rather than bisected: that + reaches a working size in one step instead of paying an O(log n) ladder of + rejected calls. Bisection remains the fallback for anything bin-packing + cannot make smaller, which is what makes the recursion terminate: a batch + already inside the budget packs back to itself, so it falls through to the + split below. A single oversized + content item (one very long message) is split by its own text instead of + by list length, since a list of length 1 has no items left to bisect -- + the resulting fragments all carry a ``fragment_group_size`` so the merge + step can recombine them into the one content item they came from, rather + than treating each fragment as its own item when reconstructing positions + for masking. That count covers however many fragments the item ended up + split into, not just two, since it can be bisected repeatedly: the + outermost single-item split stamps the total leaf count on every leaf + below it, overwriting any smaller count an inner split had set. A real + guardrail block on any (sub-)chunk raises immediately + -- callers must not lose that signal by continuing to post the remaining + chunks. + """ + try: + response: Final = await self._post_apply_guardrail_content_with_retry( + content=content, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + return ( + BedrockContentChunkResult( + response=response, + content=tuple(content), + fragment_group_size=1, + ), + ) + except HTTPException as exc: + if allow_chunking and self._is_input_too_large_error(exc.detail): + batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars) + if len(batches) > 1: + verbose_proxy_logger.warning( + "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; " + "re-sending as %d batches of at most %d characters", + len(content), + len(batches), + self.chunk_budget_chars, + ) + batch_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below + await self._apply_guardrail_content_with_chunking( + content=batch, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + for batch in batches + ] + return tuple(result for results in batch_results for result in results) + split_content: Final = self._split_bedrock_content(content) + if split_content is None: + raise + first_half, second_half = split_content + is_single_item_text_split: Final = len(content) == 1 + verbose_proxy_logger.warning( + "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; " + "splitting into %d + %d and retrying each", + len(content), + len(first_half), + len(second_half), + ) + first_results: Final = await self._apply_guardrail_content_with_chunking( + content=first_half, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + second_results: Final = await self._apply_guardrail_content_with_chunking( + content=second_half, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + allow_chunking=allow_chunking, + ) + combined_results: Final = tuple(first_results) + tuple(second_results) + if is_single_item_text_split: + return tuple( + result._replace(fragment_group_size=len(combined_results)) for result in combined_results + ) + return combined_results + raise + + async def _post_apply_guardrail_content_with_retry( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> BedrockGuardrailResponse: + """Post one ApplyGuardrail call for `content`, retrying with exponential + backoff on AWS ThrottlingException (HTTP 429). + + Chunking already trades one oversized call for several smaller ones, so + retries here are capped low -- they must not multiply per-request latency + by an order of magnitude when the account's per-second text-unit quota is + the binding constraint rather than the per-request size quota. + + A too-large rejection is deliberately excluded from the retry. AWS reports + it as a ThrottlingException (429), not only as a ValidationException, but + unlike a genuine throttle it is not transient: re-posting the same + oversized content can never succeed. Retrying it would burn every backoff + sleep and every (billed) attempt before the caller's bisection gets a + chance to split the content, at every level of the recursion. + """ + for attempt in range(_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + 1): + try: + return await self._post_apply_guardrail_content( + content=content, + base_request_data=base_request_data, + credentials=credentials, + aws_region_name=aws_region_name, + api_key=api_key, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) + except HTTPException as exc: + if ( + exc.status_code != 429 + or self._is_input_too_large_error(exc.detail) + or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + ): + raise + await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt)) + raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted") + + async def _post_apply_guardrail_content( + self, + content: Sequence[BedrockContentItem], + base_request_data: Mapping[str, Any], + credentials: "Credentials", + aws_region_name: str, + api_key: str | None, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> BedrockGuardrailResponse: + """Make exactly one signed ApplyGuardrail HTTP call for `content` and + parse the result. Raises HTTPException on a guardrail block or any + non-200 response (including 429, handled by the retry wrapper above). + + AWS also reports some failures inside a 200 body, tagging ``Output.__type`` + with an Exception marker. Those deliberately do NOT raise: the request proceeds, + matching the behaviour of this code before chunking existed. The marker survives + the merge, so the one consolidated log entry still records + ``guardrail_failed_to_respond`` rather than a success. Making that path fail + closed is a separate change, tracked apart from this PR, and belongs behind the + existing ``unreachable_fallback`` setting rather than a hardcoded status. + + A block is logged here rather than by the caller: it ends the whole chunking + flow immediately, with no further chunks attempted, so there is no later + merged response for the caller to log instead. + """ + bedrock_request_data: Final = { # mutable-ok: outbound JSON request body + **base_request_data, + "content": content, + } # mutable-ok: outbound JSON request body prepared_request: Final = self._prepare_request( credentials=credentials, data=bedrock_request_data, @@ -792,42 +1108,16 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prepared_request.headers, ) - # UI / spend logs use event_type. Bedrock's `source` is INPUT vs OUTPUT for the API - # body, which must not be confused with the proxy hook (pre_call / during_call / - # post_call). When omitted, keep legacy mapping for backward compatibility. - if logging_event_type is not None: - event_type = logging_event_type - else: - event_type = GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call - httpx_response: Final = await self._sign_and_post( prepared_request=prepared_request, request_data=request_data, event_type=event_type, start_time=start_time, + log_transport_failure=False, ) - ######################################################### - # Add guardrail information to request trace - ######################################################### - _json_response: Final = httpx_response.json() - tracing_detail: Final = self._build_tracing_detail(_json_response) - - # Raw Bedrock JSON is passed here; match/regex redaction runs once inside - # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_provider=self.guardrail_provider, - guardrail_json_response=_json_response, - request_data=request_data or {}, - guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), - start_time=start_time.timestamp(), - end_time=datetime.now(timezone.utc).timestamp(), - duration=(datetime.now(timezone.utc) - start_time).total_seconds(), - event_type=event_type, - tracing_detail=tracing_detail or None, - ) - ######################################################### if httpx_response.status_code == 200: + _json_response: Final = httpx_response.json() # check if the response was flagged verbose_proxy_logger.debug( "Bedrock AI response : %s", @@ -835,19 +1125,462 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response) if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response): + self._log_apply_guardrail_attempt( + httpx_response=httpx_response, + json_response=_json_response, + request_data=request_data, + event_type=event_type, + start_time=start_time, + ) raise self._get_http_exception_for_blocked_guardrail( bedrock_guardrail_response, request_data=request_data ) - else: - status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) - verbose_proxy_logger.error( - "Bedrock AI: error in response. Status code: %s, response: %s", - httpx_response.status_code, - httpx_response.text, - ) - raise HTTPException(status_code=status_code, detail=detail_message) + return bedrock_guardrail_response - return bedrock_guardrail_response + status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response) + verbose_proxy_logger.error( + "Bedrock AI: error in response. Status code: %s, response: %s", + httpx_response.status_code, + httpx_response.text, + ) + raise HTTPException(status_code=status_code, detail=detail_message) + + def _log_apply_guardrail_attempt( + self, + httpx_response: httpx.Response, + json_response: dict, # mutable-ok: raw AWS JSON payload + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log a single ApplyGuardrail HTTP attempt as-is (its own status, + derived from its own response). Used only for the blocked-content + case, which ends the whole chunking flow immediately.""" + tracing_detail: Final = self._build_tracing_detail(BedrockGuardrailResponse(**json_response)) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=json_response, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response), + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + tracing_detail=tracing_detail or None, + ) + + def _log_apply_guardrail_success( + self, + merged_response: BedrockGuardrailResponse, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log one logical ApplyGuardrail call -- possibly several chunk calls + under the hood -- using its final merged response, so a chunked + request produces exactly one telemetry entry, the same as an + unchunked one would. + + AWS can report a failure inside an HTTP 200 body by tagging + ``Output.__type`` with an exception marker. That marker survives the merge, + so the status is derived from the merged response rather than assumed to be + a success, which is what the pre-chunking code reported for that shape.""" + tracing_detail: Final = self._build_tracing_detail(merged_response) + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status=( + "guardrail_failed_to_respond" + if "Exception" in str((merged_response.get("Output") or {}).get("__type", "")) + else "success" + ), + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + tracing_detail=tracing_detail or None, + ) + + def _log_apply_guardrail_failure( + self, + detail: object, + request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper + event_type: GuardrailEventHooks, + start_time: "datetime", + ) -> None: + """Log one logical ApplyGuardrail call that failed end-to-end (an + unrecoverable too-large error, a non-size validation error, or + exhausted throttle retries) as a single failure, rather than logging + every failed attempt chunking made along the way.""" + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={"error": str(detail)}, # mutable-ok: logging helper requires a dict + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + ) + + @staticmethod + def _content_uses_contextual_grounding(content: Sequence[BedrockContentItem]) -> bool: + """True if any content item carries a contextual-grounding qualifier + (``grounding_source``, ``query``, or the ``guard_content`` the response + itself is tagged with once grounding is present).""" + for item in content: + if (item.get("text") or {}).get("qualifiers"): # mutable-ok: read-only empty fallback + return True + return False + + @staticmethod + def _bin_pack_bedrock_content( + content: Sequence[BedrockContentItem], + budget: int, + ) -> tuple[tuple[BedrockContentItem, ...], ...]: + """Pack whole content items, in order, into batches whose combined text + length stays within `budget`, in a single pass that carries the running + total rather than re-summing the open batch per item. + + This is the fast-path half of the hybrid chunking strategy: bin-packing + at a conservative fixed budget keeps the common case at O(n / budget) + ApplyGuardrail calls instead of the O(log n) round trips pure reactive + bisection pays on every oversized request. An item whose own text + already exceeds `budget` is not split here -- it becomes its own + (still oversized) batch and is sent as-is; if AWS rejects that batch as + too large, `_apply_guardrail_content_with_chunking`'s existing + recursive-bisection fallback takes over for that batch only. + + `budget` comes from the guardrail's ``chunk_budget_chars`` setting and + defaults to 25,000, matching ApplyGuardrail's default quota of 25 text + units (roughly 1,000 characters each) per second. Packing to that size and + posting sequentially is what keeps chunking from tripping the rate quota + and trading a size error for a throttle. Accounts with raised quotas can + configure a larger budget to spend fewer calls. + + The budget is not a correctness dependency either way. AWS's effective cap + varies by account, region, and policy, is not a fixed character count, and + cannot be read from config, so any batch it still rejects falls back to + bisection, which self-corrects however wrong the value was. An over-large + budget therefore costs one extra probe-and-bisect round trip rather than + failing the request. + """ + if not content: + return (tuple(content),) + + lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content) + + def assign(carried: tuple[int, int], length: int) -> tuple[int, int]: + batch_index, used = carried + if used + length <= budget: + return batch_index, used + length + return batch_index + 1, length + + batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:]) + return tuple( + tuple(item for _, item in group) + for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0]) + ) + + @staticmethod + def _split_bedrock_content( + content: Sequence[BedrockContentItem], + ) -> tuple[tuple[BedrockContentItem, ...], tuple[BedrockContentItem, ...]] | None: + """Bisect `content` into two roughly-equal, non-empty halves. + + When `content` already holds more than one item, it is split by list + length. When it holds exactly one item, that item's own text is split + instead (a list of length 1 has no items left to bisect, but one very + long message is still a single content item) -- at the whitespace + character nearest the midpoint rather than a raw character index, so + the cut never lands inside a word/token. This is a plain, lossless + cut with no overlap: concatenating the two fragments in order always + reproduces the original text exactly, so merging back at + ``_merge_logical_unit_outputs`` needs no reconciliation step. + + Known, accepted limitation: whitespace splitting only guards against + *accidentally* severing a single token (one denied word, one PII + pattern) across the cut. It does not, and cannot without an overlap + window, stop a *multi-word* denied phrase deliberately positioned to + straddle the boundary -- each fragment can scan clean on its own and + still reassemble into the flagged phrase. AWS's own guidance on this + API acknowledges the same gap for input chunking ("a critical piece of + text could span two (or more) chunks if not carefully divided") with + no documented resolution, and overlap-and-reconcile was evaluated and + rejected for this PR: AWS's masking output has no documented + length-preservation guarantee, so reconciling an overlap region against + masked text is not sound in general. Out of scope for this PR. + + Returns None when there is nothing left to split -- a single item + whose text is too short to halve into two non-empty pieces -- so the + caller can give up and propagate the original too-large error instead + of recursing forever. + """ + if len(content) > 1: + midpoint: Final = max(1, len(content) // 2) + return tuple(content[:midpoint]), tuple(content[midpoint:]) + + text_content: Final = content[0].get("text") or BedrockTextContent() + text: Final = text_content.get("text") or "" + if len(text) < 2: + return None + split_at: Final = BedrockGuardrail._nearest_whitespace_split_index(text) + qualifiers: Final = text_content.get("qualifiers") + + def fragment(piece: str) -> BedrockContentItem: + block: Final = ( + BedrockTextContent(text=piece, qualifiers=qualifiers) if qualifiers else BedrockTextContent(text=piece) + ) + return BedrockContentItem(text=block) + + return (fragment(text[:split_at]),), (fragment(text[split_at:]),) + + @staticmethod + def _nearest_whitespace_split_index(text: str) -> int: + """Return the index nearest `text`'s midpoint that falls on a whitespace + boundary, so splitting `text[:i]` / `text[i:]` there never severs a word. + + Any Unicode whitespace counts, not just an ASCII space. Matching only `" "` + would leave the boundary unguarded for exactly the payloads that get large + enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited, so a deny-listed word sitting at the midpoint of + one would be cut in half, scan clean on both fragments, and reassemble + intact. + + The returned index always leaves both sides non-empty, which is what makes + the caller's recursion terminate. A boundary that would put the split at 0 + or at ``len(text)`` is discarded: it would hand back a fragment identical to + the text just rejected as too large, AWS would reject that again, and each + retry would re-split it into the same unchanged fragment until the stack ran + out. The dangerous shape is a text whose only space at or after the midpoint + is its final character. + + Falls back to the raw midpoint when no usable whitespace boundary exists, either + because `text` has none at all (a single giant token) or because the only + candidates were degenerate. That is still a correct, lossless split, just no + longer guaranteed word-safe for those cases. `text` must be at least two + characters, which `_split_bedrock_content` guarantees, so the midpoint itself + is never degenerate. + """ + midpoint: Final = len(text) // 2 + before: Final = max((found.end() for found in _BEDROCK_WHITESPACE.finditer(text, 0, midpoint)), default=None) + after_match: Final = _BEDROCK_WHITESPACE.search(text, midpoint) + candidates: Final = sorted( + (split for split in (before, after_match.end() if after_match else None) if split is not None), + key=lambda split: abs(split - midpoint), + ) + return next((split for split in candidates if 0 < split < len(text)), midpoint) + + @staticmethod + def _is_input_too_large_error(detail: object) -> bool: + """True if `detail` is an AWS error message for input exceeding the + per-request text-unit quota. + + Matched on the message rather than the status code on purpose: AWS is not + consistent about which error it raises for this. Observed against a live + guardrail with an active content-filter policy, an oversized request comes + back as a *ThrottlingException* (429) reading ``Input text size (3273 text + units) exceeds the maximum allowed (1000 text units) for the content filter + policy (Classic tier)``, while the documented failure mode is a + ValidationException (400). Keying off the message covers both. + + A guardrail *block* is also raised as an HTTPException with status 400, + but its ``detail`` is always a dict (built by + ``_get_http_exception_for_blocked_guardrail``); a non-200 API error's + ``detail`` is always the plain string returned by + ``_parse_bedrock_guardrail_error_response``. Checking ``isinstance(detail, + str)`` is therefore sufficient to never mistake a real block for a + too-large error. + """ + if not isinstance(detail, str): + return False + lowered: Final = detail.lower() + return any(substring in lowered for substring in _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS) + + @staticmethod + def _merge_bedrock_guardrail_responses( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> BedrockGuardrailResponse: + """Merge the per-chunk ApplyGuardrail responses of a chunked request into + one, so a caller cannot tell whether chunking happened. + + Only ever called with responses that all passed (a block raises + immediately from ``_apply_guardrail_content_with_chunking`` and is never + added to this list). ``action`` is only set on the merged response when + at least one chunk's raw response included it, and left absent otherwise + -- mirroring a real single-call response and matching what + ``_build_tracing_detail`` treats as "Bedrock didn't report an action". + + Fields this merge has no opinion on (``actionReason``, ``guardrailCoverage``, + ``blockedResponse``, anything AWS adds later) are carried over from the chunk + responses rather than dropped, so the response and the logged telemetry keep + the shape a single unchunked call returned. The merged keys below win. + + Per AWS's documented ApplyGuardrail contract, a single call's ``outputs`` + is positionally parallel to the ``content`` items *of that call*: an + entry per item when anything in the call was masked, or an empty list + when nothing in the whole call was masked. Downstream masking + (``_apply_masking_to_messages``) walks the merged ``outputs`` by a single + running index across the *original, unchunked* message list, so a later + chunk's masked text must land at the same global position it would have + if chunking had never happened. Naively concatenating each chunk's + ``outputs`` breaks that whenever a chunk had nothing masked (its empty + list would otherwise silently swallow its items' slots, shifting every + later chunk's masked text left onto the wrong message). So every + item -- masked or not -- always contributes exactly one entry here, + falling back to that item's own original (unmasked) text when its + chunk returned no output for it; a wholly-untouched result is then + collapsed back to an empty ``outputs`` list to match a real single-call + no-op response. A chunk that returns a nonzero output count not equal + to its item count is passed through as-is instead of guessed at, since + AWS's docs don't cover partial masking within one multi-item call. + """ + logical_units: Final = BedrockGuardrail._group_fragment_units(chunk_results) + per_unit_outputs: Final = tuple(BedrockGuardrail._merge_logical_unit_outputs(unit) for unit in logical_units) + merged_outputs: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list + output for outputs, _ in per_unit_outputs for output in outputs + ] + any_masked: Final = any(masked for _, masked in per_unit_outputs) + + actions: Final = tuple( + chunk_result.response.get("action") + for chunk_result in chunk_results + if isinstance(chunk_result.response.get("action"), str) + ) + merged_action: Final = ( + "GUARDRAIL_INTERVENED" if "GUARDRAIL_INTERVENED" in actions else (actions[-1] if actions else None) + ) + merged_assessments: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list + assessment + for chunk_result in chunk_results + for assessment in (chunk_result.response.get("assessments") or []) # mutable-ok: logged payload + ] + any_usage_reported: Final = any(chunk_result.response.get("usage") for chunk_result in chunk_results) + + merged: Final[BedrockGuardrailResponse] = cast( # cast-ok: TypedDict assembled from a comprehension + BedrockGuardrailResponse, + { # mutable-ok: builds the TypedDict payload + key: value for chunk_result in chunk_results for key, value in chunk_result.response.items() + }, + ) + if merged_action is not None: + merged["action"] = merged_action + if merged_outputs and any_masked: + merged["outputs"] = merged_outputs + merged["output"] = merged_outputs + if merged_assessments: + merged["assessments"] = merged_assessments + if any_usage_reported: + merged["usage"] = BedrockGuardrail._sum_bedrock_guardrail_usage(chunk_results) + return merged + + @staticmethod + def _sum_bedrock_guardrail_usage( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> BedrockGuardrailUsage: + """Sum each chunk's ``usage`` counters field-by-field into one totals dict. + + Keys are taken from the responses rather than from a fixed list, so a counter + this code does not know about (AWS has added several) is still summed and + reported instead of being silently dropped to zero.""" + chunk_usages: Final = tuple( + chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback + for chunk_result in chunk_results + ) + return cast( # cast-ok: TypedDict assembled from a comprehension + BedrockGuardrailUsage, + { # mutable-ok: builds the TypedDict payload + key: sum(usage.get(key) or 0 for usage in chunk_usages) + for key in dict.fromkeys(key for usage in chunk_usages for key in usage) + }, + ) + + @staticmethod + def _group_fragment_units( + chunk_results: Sequence[BedrockContentChunkResult], + ) -> tuple[tuple[BedrockContentChunkResult, ...], ...]: + """Group consecutive text-fragment chunk results back into the one content + item each group came from, leaving every ordinary chunk result as a unit of + one. + + The group size is read off the results themselves rather than assumed, + because a single content item can be bisected repeatedly: two levels of + splitting yield four fragments for one item, not two. Assuming a fixed pair + here would emit two outputs for one message and shift every later message's + masked text onto the wrong message.""" + + def advance(carried: tuple[int, bool], result: BedrockContentChunkResult) -> tuple[int, bool]: + remaining, _ = carried + if remaining == 0: + return max(1, result.fragment_group_size) - 1, True + return remaining - 1, False + + starts: Final = tuple( + index + for index, (_, starts_unit) in enumerate(tuple(accumulate(chunk_results, advance, initial=(0, False)))[1:]) + if starts_unit + ) + return tuple(tuple(chunk_results[start:end]) for start, end in zip(starts, starts[1:] + (len(chunk_results),))) + + @staticmethod + def _merge_logical_unit_outputs( + unit: tuple[BedrockContentChunkResult, ...], + ) -> tuple[tuple[BedrockGuardrailOutput, ...], bool]: + """Reduce one logical unit (a fragment group of any size, or a single chunk + result) to the ``BedrockGuardrailOutput`` entries it contributes to the + merged response, plus whether any masking actually happened in it. + + Per AWS's documented ApplyGuardrail contract, a single call's + ``outputs`` is positionally parallel to the ``content`` items *of that + call*: an entry per item when anything in the call was masked, or an + empty list when nothing in the whole call was masked. Downstream + masking (``_apply_masking_to_messages``) walks the merged ``outputs`` + by a single running index across the *original, unchunked* message + list, so a later chunk's masked text must land at the same global + position it would have if chunking had never happened. So every item + -- masked or not -- always contributes exactly one entry here, falling + back to that item's own original (unmasked) text when its chunk + returned no output for it. A chunk that returns a nonzero output count + not equal to its item count is passed through as-is instead of guessed + at, since AWS's docs don't cover partial masking within one multi-item + call. + + A unit holding more than one result is a fragment group: every result in it + is one fragment of a single content item's text, so the group collapses to + one entry built from each fragment's masked text (or that fragment's own + original text where it came back unmasked), concatenated in order. This + holds for any group size, not only two. + """ + if len(unit) > 1: + + def fragment_outputs(result: BedrockContentChunkResult) -> tuple[BedrockGuardrailOutput, ...]: + return tuple(result.response.get("outputs") or result.response.get("output") or ()) + + def fragment_text(result: BedrockContentChunkResult) -> str: + source: Final = (result.content[0].get("text") or {}).get( # mutable-ok: read-only fallback + "text" + ) or "" + outputs: Final = fragment_outputs(result) + masked: Final = outputs[0].get("text") if outputs else None + return masked if masked is not None else source + + merged_text: Final = "".join(fragment_text(result) for result in unit) + any_masked: Final = any(fragment_outputs(result) for result in unit) + return (BedrockGuardrailOutput(text=merged_text),), any_masked + + (chunk_result,) = unit + chunk_outputs: Final = chunk_result.response.get("outputs") or chunk_result.response.get("output") or () + if len(chunk_outputs) == len(chunk_result.content): + return tuple(chunk_outputs), bool(chunk_outputs) + if not chunk_outputs: + return tuple( + BedrockGuardrailOutput( + text=(item.get("text") or {}).get("text") or "" # mutable-ok: read-only fallback + ) + for item in chunk_result.content + ), False + return tuple(chunk_outputs), True async def _sign_and_post( self, @@ -855,6 +1588,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): request_data: dict | None, event_type: GuardrailEventHooks, start_time: "datetime", + log_transport_failure: bool = True, ) -> httpx.Response: """POST a signed Bedrock request, logging+raising on network/HTTP errors. @@ -862,6 +1596,20 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): transport-error handling cannot drift. Returns the raw ``httpx.Response`` on success (including non-2xx that httpx did not raise on); the 200-path logging, status and tracing stay with each caller because the two APIs report differently. + + ``log_transport_failure=False`` suppresses the ``guardrail_failed_to_respond`` + entry for a non-200 that is re-raised as an ``HTTPException``, for callers that + own consolidated per-request logging. The ApplyGuardrail path needs this: + ``AsyncHTTPHandler.post`` calls ``raise_for_status()``, so every non-200 lands + in this handler, and one logical request can legitimately produce several of + them (a too-large probe, then each rejected bisection level) while still + succeeding overall. Logging per attempt would report a recovered request as + several failures plus a success. + + The connection-level branch below (timeout, endpoint down) still logs + unconditionally: it re-raises the original exception rather than an + ``HTTPException``, so no consolidating caller catches it, and suppressing it + would drop the only record of the failure. """ try: return await self.async_handler.post( @@ -882,16 +1630,19 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): status_code, detail_message, ) = self._parse_bedrock_guardrail_error_response(err_response) - self.add_standard_logging_guardrail_information_to_request_data( - guardrail_provider=self.guardrail_provider, - guardrail_json_response={"error": detail_message}, - request_data=request_data or {}, - guardrail_status="guardrail_failed_to_respond", - start_time=start_time.timestamp(), - end_time=datetime.now(timezone.utc).timestamp(), - duration=(datetime.now(timezone.utc) - start_time).total_seconds(), - event_type=event_type, - ) + if log_transport_failure: + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_provider=self.guardrail_provider, + guardrail_json_response={ # mutable-ok: logging helper requires a dict + "error": detail_message + }, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict + guardrail_status="guardrail_failed_to_respond", + start_time=start_time.timestamp(), + end_time=datetime.now(timezone.utc).timestamp(), + duration=(datetime.now(timezone.utc) - start_time).total_seconds(), + event_type=event_type, + ) raise HTTPException(status_code=status_code, detail=detail_message) from e except HTTPException: raise @@ -900,7 +1651,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1027,7 +1778,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": detail_message}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1043,7 +1794,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response={"error": str(e)}, - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status="guardrail_failed_to_respond", start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), @@ -1061,7 +1812,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): self.add_standard_logging_guardrail_information_to_request_data( guardrail_provider=self.guardrail_provider, guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response), - request_data=request_data or {}, + request_data=request_data or {}, # mutable-ok: logging helper requires a dict guardrail_status=self._get_invoke_checks_status(bool(violations)), start_time=start_time.timestamp(), end_time=datetime.now(timezone.utc).timestamp(), diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 190d19f3d52..0d23e19f88d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -20,6 +20,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): content_filter_threshold=litellm_params.content_filter_threshold, prompt_attack_threshold=litellm_params.prompt_attack_threshold, pii_confidence_threshold=litellm_params.pii_confidence_threshold, + chunk_budget_chars=litellm_params.chunk_budget_chars, default_on=litellm_params.default_on, disable_exception_on_block=litellm_params.disable_exception_on_block, mask_request_content=litellm_params.mask_request_content, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 6b354a39101..bbb6d758814 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -5,6 +5,7 @@ from typing import Any, Final, Literal from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator from typing_extensions import Required, TypedDict +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.types.proxy.guardrails.guardrail_hooks.akto import ( AktoConfigModel, ) @@ -525,6 +526,15 @@ class BedrockGuardrailConfigModel(BaseModel): description="InvokeGuardrailChecks: block when any sensitiveInformation confidenceScore " ">= this value (scores are in [0,1]). Set to null to make PII detection detect-only.", ) + chunk_budget_chars: int = Field( + default=BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + gt=0, + description="ApplyGuardrail: batch size, in characters, used to re-send content after AWS " + "has rejected a request as too large. Requests AWS accepts are always sent in a single " + "call, so this has no effect until a rejection happens. Defaults to 25,000; a batch AWS " + "still rejects is bisected automatically, so this value only trades round trips against " + "batch size and cannot fail a request on its own.", + ) class LakeraV2GuardrailConfigModel(BaseModel): diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d97bdc3532f..8d66b624341 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -28,6 +28,9 @@ class BedrockGuardrailUsage(TypedDict, total=False): sensitiveInformationPolicyUnits: int | None sensitiveInformationPolicyFreeUnits: int | None contextualGroundingPolicyUnits: int | None + contentPolicyImageUnits: int | None + automatedReasoningPolicyUnits: int | None + automatedReasoningPolicies: int | None class BedrockGuardrailOutput(TypedDict, total=False): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 76a695ce3fd..837fb93d331 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -7,6 +7,7 @@ import os import sys from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException @@ -15,12 +16,18 @@ sys.path.insert(0, os.path.abspath("../../../../../..")) import litellm from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth +from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentChunkResult, BedrockGuardrail, _redact_pii_matches, ) from litellm.proxy.utils import ProxyLogging from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockContentItem, + BedrockTextContent, +) from litellm.types.utils import CallTypes, ModelResponse @@ -53,9 +60,7 @@ async def test__redact_pii_matches_function(): redacted_response = _redact_pii_matches(response_with_pii) # Verify that PII matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" @@ -173,12 +178,8 @@ async def test__redact_pii_matches_multiple_assessments(): redacted_response = _redact_pii_matches(response_multiple_assessments) # Verify all PII in all assessments are redacted - assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] - assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ - "piiEntities" - ] + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"]["piiEntities"] assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" @@ -199,9 +200,7 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII mock_bedrock_response = MagicMock() @@ -239,20 +238,11 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" - ) as mock_debug, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch("litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug") as mock_debug, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method that should log the redacted response @@ -275,37 +265,23 @@ async def test_bedrock_guardrail_logging_uses_redacted_response(): bedrock_response_log_call = call break - assert ( - bedrock_response_log_call is not None - ), "Should have logged Bedrock AI response" + assert bedrock_response_log_call is not None, "Should have logged Bedrock AI response" # Extract the logged response data - logged_response = bedrock_response_log_call[0][ - 1 - ] # Second argument to debug call + logged_response = bedrock_response_log_call[0][1] # Second argument to debug call # Verify that the logged response has redacted PII assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] - == "[REDACTED]" + logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) # Verify other fields are preserved assert logged_response["action"] == "GUARDRAIL_INTERVENED" - assert ( - logged_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["type"] - == "PHONE" - ) + assert logged_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["type"] == "PHONE" slg_list = request_data["metadata"]["standard_logging_guardrail_information"] assert ( - slg_list[0]["guardrail_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + slg_list[0]["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "[REDACTED]" ) @@ -319,9 +295,7 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the Bedrock API response with PII original_response_data = { @@ -361,17 +335,10 @@ async def test_bedrock_guardrail_original_response_not_modified(): # Mock AWS-related methods to ensure test runs without external dependencies with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ) as mock_load_creds, - patch.object( - guardrail, "_prepare_request", return_value=MagicMock() - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")) as mock_load_creds, + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare_request, ): - mock_post.return_value = mock_bedrock_response # Call the method @@ -385,19 +352,12 @@ async def test_bedrock_guardrail_original_response_not_modified(): # (The json() method should return the original data) original_data = mock_bedrock_response.json() assert ( - original_data["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ][0]["match"] + original_data["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" ) # Verify that the returned BedrockGuardrailResponse contains original data - assert ( - result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ - "match" - ] - == "+1 412 555 1212" - ) + assert result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "+1 412 555 1212" print("Original response not modified test passed") @@ -454,18 +414,14 @@ async def test__redact_pii_matches_preserves_non_pii_entities(): redacted_response = _redact_pii_matches(response_with_mixed_data) # Verify that PII entity matches are redacted - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" # Verify that regex matches are also redacted (updated behavior) - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" @@ -496,9 +452,7 @@ async def test_pii_redaction_matches_debug_output_format(): "assessments": [ { "invocationMetrics": { - "guardrailCoverage": { - "textCharacters": {"guarded": 84, "total": 84} - }, + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, "guardrailProcessingLatency": 322, "usage": { "contentPolicyImageUnits": 0, @@ -553,9 +507,7 @@ async def test_pii_redaction_matches_debug_output_format(): redacted_response = _redact_pii_matches(original_response) # Verify the redacted response matches your expected debug output - pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "piiEntities" - ] + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] # All PII matches should be redacted assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" @@ -570,34 +522,19 @@ async def test_pii_redaction_matches_debug_output_format(): assert pii_entities[0]["detected"] == True # Verify that the original response is unchanged - original_pii_entities = original_response["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"] - assert ( - original_pii_entities[0]["match"] == "John Smith" - ), "Original should be unchanged" - assert ( - original_pii_entities[1]["match"] == "324-12-3212" - ), "Original should be unchanged" - assert ( - original_pii_entities[2]["match"] == "607-456-7890" - ), "Original should be unchanged" + original_pii_entities = original_response["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert original_pii_entities[0]["match"] == "John Smith", "Original should be unchanged" + assert original_pii_entities[1]["match"] == "324-12-3212", "Original should be unchanged" + assert original_pii_entities[2]["match"] == "607-456-7890", "Original should be unchanged" # Verify all other metadata is preserved in redacted response assert redacted_response["action"] == "GUARDRAIL_INTERVENED" assert redacted_response["actionReason"] == "Guardrail blocked." assert redacted_response["blockedResponse"] == "Input blocked by PII policy" - assert ( - redacted_response["assessments"][0]["invocationMetrics"][ - "guardrailProcessingLatency" - ] - == 322 - ) + assert redacted_response["assessments"][0]["invocationMetrics"]["guardrailProcessingLatency"] == 322 print("PII redaction matches debug output format test passed") - print( - f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" - ) + print(f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}") print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") @@ -632,14 +569,10 @@ async def test__redact_pii_matches_with_regex_matches(): redacted_response = _redact_pii_matches(response_with_regex) # Verify that regex matches are redacted - regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ - "regexes" - ] + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" - assert ( - regexes[1]["match"] == "[REDACTED]" - ), "Credit card regex match should be redacted" + assert regexes[1]["match"] == "[REDACTED]", "Credit card regex match should be redacted" # Verify other fields are preserved assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" @@ -648,13 +581,9 @@ async def test__redact_pii_matches_with_regex_matches(): assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" # Verify original response is unchanged - original_regexes = response_with_regex["assessments"][0][ - "sensitiveInformationPolicy" - ]["regexes"] + original_regexes = response_with_regex["assessments"][0]["sensitiveInformationPolicy"]["regexes"] assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" - assert ( - original_regexes[1]["match"] == "4111-1111-1111-1111" - ), "Original should be unchanged" + assert original_regexes[1]["match"] == "4111-1111-1111-1111", "Original should be unchanged" print("Regex matches redaction test passed") @@ -690,31 +619,17 @@ async def test__redact_pii_matches_with_custom_words(): # Verify that custom word matches are redacted custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "First custom word match should be redacted" - assert ( - custom_words[1]["match"] == "[REDACTED]" - ), "Second custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "First custom word match should be redacted" + assert custom_words[1]["match"] == "[REDACTED]", "Second custom word match should be redacted" # Verify other fields are preserved - assert ( - custom_words[0]["action"] == "BLOCKED" - ), "Custom word action should be preserved" - assert ( - custom_words[1]["action"] == "ANONYMIZED" - ), "Custom word action should be preserved" + assert custom_words[0]["action"] == "BLOCKED", "Custom word action should be preserved" + assert custom_words[1]["action"] == "ANONYMIZED", "Custom word action should be preserved" # Verify original response is unchanged - original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ - "customWords" - ] - assert ( - original_custom_words[0]["match"] == "confidential_data" - ), "Original should be unchanged" - assert ( - original_custom_words[1]["match"] == "secret_information" - ), "Original should be unchanged" + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"]["customWords"] + assert original_custom_words[0]["match"] == "confidential_data", "Original should be unchanged" + assert original_custom_words[1]["match"] == "secret_information", "Original should be unchanged" print("Custom words redaction test passed") @@ -750,41 +665,21 @@ async def test__redact_pii_matches_with_managed_words(): redacted_response = _redact_pii_matches(response_with_managed_words) # Verify that managed word matches are redacted - managed_words = redacted_response["assessments"][0]["wordPolicy"][ - "managedWordLists" - ] + managed_words = redacted_response["assessments"][0]["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "First managed word match should be redacted" - assert ( - managed_words[1]["match"] == "[REDACTED]" - ), "Second managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "First managed word match should be redacted" + assert managed_words[1]["match"] == "[REDACTED]", "Second managed word match should be redacted" # Verify other fields are preserved - assert ( - managed_words[0]["action"] == "BLOCKED" - ), "Managed word action should be preserved" - assert ( - managed_words[0]["type"] == "PROFANITY" - ), "Managed word type should be preserved" - assert ( - managed_words[1]["action"] == "ANONYMIZED" - ), "Managed word action should be preserved" - assert ( - managed_words[1]["type"] == "HATE_SPEECH" - ), "Managed word type should be preserved" + assert managed_words[0]["action"] == "BLOCKED", "Managed word action should be preserved" + assert managed_words[0]["type"] == "PROFANITY", "Managed word type should be preserved" + assert managed_words[1]["action"] == "ANONYMIZED", "Managed word action should be preserved" + assert managed_words[1]["type"] == "HATE_SPEECH", "Managed word type should be preserved" # Verify original response is unchanged - original_managed_words = response_with_managed_words["assessments"][0][ - "wordPolicy" - ]["managedWordLists"] - assert ( - original_managed_words[0]["match"] == "inappropriate_word" - ), "Original should be unchanged" - assert ( - original_managed_words[1]["match"] == "offensive_term" - ), "Original should be unchanged" + original_managed_words = response_with_managed_words["assessments"][0]["wordPolicy"]["managedWordLists"] + assert original_managed_words[0]["match"] == "inappropriate_word", "Original should be unchanged" + assert original_managed_words[1]["match"] == "offensive_term", "Original should be unchanged" print("Managed words redaction test passed") @@ -841,9 +736,7 @@ async def test__redact_pii_matches_comprehensive_coverage(): # PII entities pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] - assert ( - pii_entities[0]["match"] == "[REDACTED]" - ), "PII entity match should be redacted" + assert pii_entities[0]["match"] == "[REDACTED]", "PII entity match should be redacted" # Regex matches regexes = assessment["sensitiveInformationPolicy"]["regexes"] @@ -851,15 +744,11 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Custom words custom_words = assessment["wordPolicy"]["customWords"] - assert ( - custom_words[0]["match"] == "[REDACTED]" - ), "Custom word match should be redacted" + assert custom_words[0]["match"] == "[REDACTED]", "Custom word match should be redacted" # Managed words managed_words = assessment["wordPolicy"]["managedWordLists"] - assert ( - managed_words[0]["match"] == "[REDACTED]" - ), "Managed word match should be redacted" + assert managed_words[0]["match"] == "[REDACTED]", "Managed word match should be redacted" # Verify all other fields are preserved assert pii_entities[0]["type"] == "EMAIL" @@ -868,21 +757,10 @@ async def test__redact_pii_matches_comprehensive_coverage(): # Verify original response is unchanged original_assessment = comprehensive_response["assessments"][0] - assert ( - original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] - == "user@example.com" - ) - assert ( - original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] - == "555-123-4567" - ) - assert ( - original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" - ) - assert ( - original_assessment["wordPolicy"]["managedWordLists"][0]["match"] - == "inappropriate" - ) + assert original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] == "user@example.com" + assert original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] == "555-123-4567" + assert original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + assert original_assessment["wordPolicy"]["managedWordLists"][0]["match"] == "inappropriate" print("Comprehensive coverage redaction test passed") @@ -914,9 +792,7 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method to avoid actual AWS credential loading - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -926,10 +802,12 @@ async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + ) print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") @@ -944,9 +822,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) # Create guardrail without explicit aws_bedrock_runtime_endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -960,9 +836,7 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -972,10 +846,10 @@ async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): ) # Verify that the custom endpoint from environment is used in the URL - expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + expected_url = ( + f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, f"Expected URL to contain env endpoint. Got: {prepped_request.url}" print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") @@ -988,9 +862,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) # Create guardrail without any custom endpoint - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock credentials mock_credentials = MagicMock() @@ -1004,9 +876,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey aws_region_name = "us-west-2" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1017,9 +887,7 @@ async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkey # Verify that the default endpoint is used expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected default URL. Got: {prepped_request.url}" + assert prepped_request.url == expected_url, f"Expected default URL. Got: {prepped_request.url}" print(f"Default endpoint test passed. URL: {prepped_request.url}") @@ -1057,9 +925,7 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch aws_region_name = "us-east-1" # Mock the _load_credentials method - with patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) - ): + with patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name)): # Call _prepare_request which internally calls get_runtime_endpoint prepped_request = guardrail._prepare_request( credentials=mock_credentials, @@ -1069,10 +935,12 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch ) # Verify that the parameter takes precedence over environment variable - expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" - assert ( - prepped_request.url == expected_url - ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + expected_url = ( + f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + ) + assert prepped_request.url == expected_url, ( + f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + ) print(f"Parameter precedence test passed. URL: {prepped_request.url}") @@ -1081,14 +949,10 @@ async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" # Create a BedrockGuardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Mock the make_bedrock_api_request method - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api_request: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api_request: # Test the apply_guardrail method with tool_calls in response inputs = { "texts": [], @@ -1115,14 +979,9 @@ async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): assert guardrailed_inputs is not None assert "tool_calls" in guardrailed_inputs assert len(guardrailed_inputs["tool_calls"]) == 1 - assert ( - guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" - ) + assert guardrailed_inputs["tool_calls"][0]["id"] == "call_eFSCWFsyL7MclHYnzKrcQnMK" assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" - assert ( - guardrailed_inputs["tool_calls"][0]["function"]["arguments"] - == '{"location":"São Paulo"}' - ) + assert guardrailed_inputs["tool_calls"][0]["function"]["arguments"] == '{"location":"São Paulo"}' # Verify that the Bedrock API was NOT called since there's no text to process mock_api_request.assert_not_called() print("✅ apply_guardrail with tool_calls test passed - no API call made") @@ -1136,14 +995,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): policies (e.g. PII on model output) then returned action=NONE for non-streaming completions that go through unified_guardrail -> process_output_response. """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1168,14 +1023,10 @@ async def test_bedrock_apply_guardrail_response_uses_OUTPUT_source(): @pytest.mark.asyncio async def test_bedrock_apply_guardrail_request_uses_INPUT_source(): """input_type='request' must call Bedrock with source=INPUT and user messages.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") bedrock_none = {"action": "NONE", "output": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = bedrock_none await guardrail.apply_guardrail( @@ -1258,12 +1109,8 @@ async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): # Mock AWS-related methods with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -1431,9 +1278,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" def _create_guardrail(self) -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") @pytest.mark.asyncio async def test_should_handle_all_null_policy_sub_lists(self): @@ -1554,9 +1399,7 @@ class TestShouldRaiseGuardrailBlockedExceptionNullSafety: { "sensitiveInformationPolicy": { "piiEntities": None, - "regexes": [ - {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} - ], + "regexes": [{"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"}], }, } ], @@ -1611,18 +1454,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_none_texts_in_inputs(self): """inputs[\"texts\"] is explicitly None — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {"texts": None} # Explicit None mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1645,18 +1484,14 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_should_handle_missing_texts_key(self): """inputs has no \"texts\" key at all — should not crash.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") inputs = {} # No "texts" key mock_credentials = MagicMock() with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, patch.object( guardrail, "_load_credentials", @@ -1677,9 +1512,7 @@ class TestApplyGuardrailNullSafety: @pytest.mark.asyncio async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") # Test 1: ANONYMIZED action should NOT raise exception anonymized_response = { @@ -1700,9 +1533,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): ], } - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) + should_raise = guardrail._should_raise_guardrail_blocked_exception(anonymized_response) assert should_raise is False, "ANONYMIZED actions should not raise exceptions" # Test 2: BLOCKED action should raise exception @@ -1710,13 +1541,7 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "I can't provide that information."}], "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } + {"topicPolicy": {"topics": [{"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"}]}} ], } @@ -1738,19 +1563,13 @@ async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): } ] }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"}]}, } ], } should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" + assert should_raise is True, "Mixed actions with any BLOCKED should raise exceptions" # Test 4: NONE action should not raise exception none_response = { @@ -1782,9 +1601,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): When logging_event_type is set, it must be forwarded to standard guardrail logging. When omitted, INPUT maps to pre_call (legacy). """ - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1795,13 +1612,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): mock_bedrock_response.json.return_value = { "action": "NONE", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "match": "GG", "action": "BLOCKED"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "match": "GG", "action": "BLOCKED"}]}} ], } @@ -1811,12 +1622,8 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -1831,15 +1638,13 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): request_data=request_data, logging_event_type=GuardrailEventHooks.during_call, ) - assert ( - mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call - ) + assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call # Raw Bedrock JSON is forwarded; redaction runs once in # CustomGuardrail.add_standard_logging_guardrail_information_to_request_data. assert ( - mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][ - "sensitiveInformationPolicy" - ]["piiEntities"][0]["match"] + mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] == "GG" ) @@ -1855,9 +1660,7 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs(): @pytest.mark.asyncio async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "test-access-key" mock_credentials.secret_key = "test-secret-key" @@ -1873,15 +1676,9 @@ async def test_make_bedrock_api_request_filters_dynamic_evaluation_overrides(): prepared_request.headers = {} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), - patch.object( - guardrail, "_prepare_request", return_value=prepared_request - ) as mock_prepare_request, + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=prepared_request) as mock_prepare_request, patch.object( guardrail, "get_guardrail_dynamic_request_body_params", @@ -1933,9 +1730,7 @@ async def test_during_call_hook_invokes_bedrock_async_moderation_hook(): "model": "gpt-4", "messages": [{"role": "user", "content": "test"}], }, - user_api_key_dict=UserAPIKeyAuth( - api_key="test_key", user_id="test_user" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test_key", user_id="test_user"), call_type="completion", ) finally: @@ -1991,11 +1786,7 @@ def test_extract_blocked_assessments_multiple_policies(): "action": "GUARDRAIL_INTERVENED", "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Investment", "type": "DENY", "action": "BLOCKED"} - ] - }, + "topicPolicy": {"topics": [{"name": "Investment", "type": "DENY", "action": "BLOCKED"}]}, "contentPolicy": { "filters": [ { @@ -2006,9 +1797,7 @@ def test_extract_blocked_assessments_multiple_policies(): } ] }, - "wordPolicy": { - "customWords": [{"match": "forbidden", "action": "BLOCKED"}] - }, + "wordPolicy": {"customWords": [{"match": "forbidden", "action": "BLOCKED"}]}, } ], } @@ -2023,13 +1812,7 @@ def test_extract_blocked_assessments_only_anonymized_returns_empty(): response = { "action": "GUARDRAIL_INTERVENED", "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } assert g._extract_blocked_assessments(response) == [] @@ -2049,23 +1832,14 @@ def test_get_http_exception_includes_assessments_and_identifier(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "BLOCKED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "BLOCKED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) assert isinstance(exc, HTTPException) assert exc.status_code == 400 assert exc.detail["error"] == "Violated guardrail policy" - assert ( - exc.detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question." - ) + assert exc.detail["bedrock_guardrail_response"] == "Sorry, the model cannot answer this question." assert exc.detail["guardrailIdentifier"] == "amgllac6xf3r" assert exc.detail["guardrailVersion"] == "1" assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy" @@ -2088,15 +1862,11 @@ def test_extract_violation_category_names_mixed_policies(): {"name": "Tax Advice", "action": "BLOCKED"}, ] }, - "contentPolicy": { - "filters": [{"type": "VIOLENCE", "action": "BLOCKED"}] - }, + "contentPolicy": {"filters": [{"type": "VIOLENCE", "action": "BLOCKED"}]}, "wordPolicy": { "managedWordLists": [{"type": "PROFANITY", "action": "BLOCKED"}], }, - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}] - }, + "sensitiveInformationPolicy": {"piiEntities": [{"type": "EMAIL", "action": "BLOCKED"}]}, } ], } @@ -2120,13 +1890,9 @@ def test_extract_violation_category_names_does_not_leak_user_input(): "assessments": [ { "wordPolicy": { - "customWords": [ - {"match": "secret-codeword-abc-123", "action": "BLOCKED"} - ], - }, - "sensitiveInformationPolicy": { - "regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}] + "customWords": [{"match": "secret-codeword-abc-123", "action": "BLOCKED"}], }, + "sensitiveInformationPolicy": {"regexes": [{"match": "4111-1111-1111-1111", "action": "BLOCKED"}]}, } ], } @@ -2166,13 +1932,7 @@ def test_extract_violation_category_names_skips_anonymized(): g = _make_guardrail() response = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}] - } - } - ], + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED"}]}}], } assert g._extract_violation_category_names(response) == [] @@ -2190,9 +1950,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): ``tracing_detail`` so downstream loggers (OTEL, ...) can surface the raw provider verdict as a queryable attribute without re-parsing the redacted guardrail_response blob.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2202,13 +1960,7 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): mock_bedrock_response.status_code = 200 mock_bedrock_response.json.return_value = { "action": "GUARDRAIL_INTERVENED", - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Fiduciary Advice", "action": "BLOCKED"}]}}], } request_data = { @@ -2217,12 +1969,8 @@ async def test_make_bedrock_api_request_forwards_guardrail_action(): } with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2253,9 +2001,7 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): """If the Bedrock response omits ``action`` (older / partial payloads), the field must be left off ``tracing_detail`` rather than written as ``None`` — downstream code expects strings or absence, not nulls.""" - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + guardrail = BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") mock_credentials = MagicMock() mock_credentials.access_key = "k" mock_credentials.secret_key = "s" @@ -2266,12 +2012,8 @@ async def test_make_bedrock_api_request_omits_guardrail_action_when_missing(): mock_bedrock_response.json.return_value = {"assessments": []} with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), patch.object( guardrail, @@ -2300,13 +2042,7 @@ def test_get_http_exception_no_blocked_assessments_omits_field(): "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "blocked"}], "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - {"type": "NAME", "action": "ANONYMIZED", "match": "Jack"} - ] - } - } + {"sensitiveInformationPolicy": {"piiEntities": [{"type": "NAME", "action": "ANONYMIZED", "match": "Jack"}]}} ], } exc = g._get_http_exception_for_blocked_guardrail(response) @@ -2370,9 +2106,7 @@ async def test_streaming_post_call_only_runs_output_scan(): yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: out = [] async for chunk in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), @@ -2382,18 +2116,11 @@ async def test_streaming_post_call_only_runs_output_scan(): out.append(chunk) assert len(out) >= 1 - output_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT" - ] + output_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "OUTPUT"] assert len(output_calls) == 1 assert output_calls[0].kwargs.get("request_data") is request_data - assert ( - output_calls[0].kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) - input_calls = [ - c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT" - ] + assert output_calls[0].kwargs.get("logging_event_type") == GuardrailEventHooks.post_call + input_calls = [c for c in mock_make.call_args_list if c.kwargs.get("source") == "INPUT"] assert len(input_calls) == 0 @@ -2432,9 +2159,7 @@ async def test_streaming_post_call_output_only_path_passes_request_data_to_make_ yield c minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: async for _ in guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=UserAPIKeyAuth(), response=mock_stream(), @@ -2488,9 +2213,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): ) minimal = {"action": "NONE", "assessments": [], "outputs": []} - with patch.object( - guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal) - ) as mock_make: + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(return_value=minimal)) as mock_make: await guardrail.async_post_call_success_hook( data=request_data, user_api_key_dict=UserAPIKeyAuth(), @@ -2499,10 +2222,7 @@ async def test_post_call_success_hook_only_runs_output_scan(): sources = [c.kwargs.get("source") for c in mock_make.call_args_list] assert sources == ["OUTPUT"] - assert ( - mock_make.call_args.kwargs.get("logging_event_type") - == GuardrailEventHooks.post_call - ) + assert mock_make.call_args.kwargs.get("logging_event_type") == GuardrailEventHooks.post_call # --------------------------------------------------------------------------- @@ -2522,9 +2242,7 @@ _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo." def _grounding_guardrail() -> BedrockGuardrail: - return BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) + return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT") def _grounding_messages() -> list: @@ -2556,27 +2274,19 @@ def _model_response(content: str) -> ModelResponse: # Expected OUTPUT content blocks, keyed by their grounding qualifier, so the # per-test assertions read as the block sequence they expect. -_GROUNDING_SOURCE_BLOCK = { - "text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]} -} +_GROUNDING_SOURCE_BLOCK = {"text": {"text": _GROUNDING_SOURCE_TEXT, "qualifiers": ["grounding_source"]}} _QUERY_BLOCK = {"text": {"text": _GROUNDING_QUERY_TEXT, "qualifiers": ["query"]}} -_GUARD_BLOCK = { - "text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]} -} +_GUARD_BLOCK = {"text": {"text": _GROUNDING_RESPONSE_TEXT, "qualifiers": ["guard_content"]}} def _input_request(messages: list) -> dict: """Arrange a guardrail and act: build the Bedrock INPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="INPUT", messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages) def _output_request(messages: list, response=None) -> dict: """Arrange a guardrail and act: build the Bedrock OUTPUT payload.""" - return _grounding_guardrail().convert_to_bedrock_format( - source="OUTPUT", response=response, messages=messages - ) + return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages) def test_grounding_input_strips_grounding_and_query_qualifiers(): @@ -2600,9 +2310,7 @@ def test_grounding_input_leaves_existing_guarded_text_unqualified(): """An existing guarded_text input block keeps its legacy unqualified payload.""" expected_request = {"source": "INPUT", "content": [{"text": {"text": "policy"}}]} - actual_request = _input_request( - [{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}] - ) + actual_request = _input_request([{"role": "user", "content": [{"type": "guarded_text", "text": "policy"}]}]) assert actual_request == expected_request @@ -2615,9 +2323,7 @@ def test_grounding_output_assembles_source_query_and_response(): "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], } - actual_request = _output_request( - _grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(_grounding_messages(), _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2629,9 +2335,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): "content": [{"text": {"text": "Hi there."}}], } - actual_request = _output_request( - [{"role": "user", "content": "hello"}], _model_response("Hi there.") - ) + actual_request = _output_request([{"role": "user", "content": "hello"}], _model_response("Hi there.")) assert actual_request == expected_request @@ -2639,9 +2343,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." - uk_source_block = { - "text": {"text": uk_source_text, "qualifiers": ["grounding_source"]} - } + uk_source_block = {"text": {"text": uk_source_text, "qualifiers": ["grounding_source"]}} messages = [ { "role": "system", @@ -2662,9 +2364,7 @@ def test_grounding_output_combines_multiple_sources(): ], } - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == expected_request @@ -2709,9 +2409,7 @@ def test_grounding_source_trusted_only_from_app_roles(role, is_trusted): if is_trusted: expected_content = [_GROUNDING_SOURCE_BLOCK, *expected_content] - actual_request = _output_request( - messages, _model_response(_GROUNDING_RESPONSE_TEXT) - ) + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) assert actual_request == {"source": "OUTPUT", "content": expected_content} @@ -2748,12 +2446,8 @@ async def test_grounding_output_blocked_raises_400(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = mock_bedrock_response @@ -2791,13 +2485,7 @@ def _blocked_bedrock_httpx_response() -> MagicMock: response.json.return_value = { "action": "GUARDRAIL_INTERVENED", "outputs": [{"text": "Sorry, the model cannot answer this question."}], - "assessments": [ - { - "topicPolicy": { - "topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}] - } - } - ], + "assessments": [{"topicPolicy": {"topics": [{"name": "Denied", "type": "DENY", "action": "BLOCKED"}]}}], } return response @@ -2820,12 +2508,8 @@ async def test_make_bedrock_api_request_block_raises_modify_response_when_flag_s mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2857,12 +2541,8 @@ async def test_make_bedrock_api_request_block_raises_http_400_when_flag_unset(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2903,12 +2583,8 @@ async def test_async_pre_call_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2953,12 +2629,8 @@ async def test_async_moderation_hook_propagates_modify_response_on_block(): mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -2998,12 +2670,8 @@ async def test_async_post_call_success_hook_attaches_original_response_on_block( mock_credentials.token = None with ( - patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post, - patch.object( - guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") - ), + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), patch.object(guardrail, "_prepare_request", return_value=MagicMock()), ): mock_post.return_value = _blocked_bedrock_httpx_response() @@ -3032,9 +2700,7 @@ async def test_apply_guardrail_propagates_modify_response_on_block(): disable_exception_on_block=True, ) - with patch.object( - guardrail, "make_bedrock_api_request", new_callable=AsyncMock - ) as mock_api: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.side_effect = ModifyResponseException( message="Sorry, the model cannot answer this question.", model="bedrock-nova-micro", @@ -3276,6 +2942,1160 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n assert response is not None +def _too_large_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = { + "message": "Input is too long. Content size exceeds the maximum input size in text units." + } + response.text = json.dumps(response.json.return_value) + return response + + +def _other_validation_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 400 + response.json.return_value = {"message": "guardrailIdentifier is not valid"} + response.text = json.dumps(response.json.return_value) + return response + + +def _throttling_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 429 + response.json.return_value = {"message": "Rate exceeded"} + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _too_large_throttling_httpx_response() -> MagicMock: + """The shape AWS actually returns for an oversized ApplyGuardrail request when + the guardrail has an active content-filter policy: a 429 ThrottlingException, + not the documented 400 ValidationException. Message taken from a live call.""" + response = MagicMock() + response.status_code = 429 + response.json.return_value = { + "message": ( + "Input text size (3273 text units) exceeds the maximum allowed " + "(1000 text units) for the content filter policy (Classic tier)." + ) + } + response.text = json.dumps(response.json.return_value) + response.headers = {} + return response + + +def _passing_bedrock_httpx_response(marker: str) -> MagicMock: + """A successful ApplyGuardrail response tagged with `marker` so tests can + verify which chunk produced which output/usage after merging.""" + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "NONE", + "outputs": [{"text": marker}], + "assessments": [], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _blocking_bedrock_httpx_response(marker: str) -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": marker}], + "assessments": [{"topicPolicy": {"topics": [{"name": marker, "type": "DENY", "action": "BLOCKED"}]}}], + "usage": {"contentPolicyUnits": 1}, + } + return response + + +def _bedrock_guardrail_for_chunk_tests() -> "BedrockGuardrail": + return BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + ) + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunks_on_too_large_validation_error(): + """A too-large 400 on the whole-content call must trigger a bisect-and-retry, + and the two chunk responses must be merged (assessments concatenated, usage + summed, outputs concatenated) rather than losing either half's result.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first half of a very long message"}, + {"role": "user", "content": "second half of a very long message"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _blocking_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + detail = exc_info.value.detail + assert exc_info.value.status_code == 400 + assert "chunk-2" in detail["bedrock_guardrail_response"] + assert detail["assessments"][0]["matches"][0]["name"] == "chunk-2" + + +@pytest.mark.asyncio +async def test_apply_guardrail_merges_usage_and_outputs_across_chunks_when_both_pass(): + """When both chunks pass clean, the merged response must still carry both + chunks' outputs/usage forward (needed for accurate logging/telemetry) and + must not itself raise.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + assert result.get("usage", {}).get("contentPolicyUnits") == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_recurses_past_first_bisection_into_four_chunks(): + """A payload that is still too large after one bisection must keep splitting + -- chunking is not capped at two pieces. Four messages where both the + whole-content call AND both first-level halves are too large must recurse + one level deeper into four chunks that all fit, not give up after the + first split.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "message one"}, + {"role": "user", "content": "message two"}, + {"role": "user", "content": "message three"}, + {"role": "user", "content": "message four"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole content: [1,2,3,4] + _too_large_validation_httpx_response(), # first half: [1,2] + _passing_bedrock_httpx_response("message one"), + _passing_bedrock_httpx_response("message two"), + _too_large_validation_httpx_response(), # second half: [3,4] + _passing_bedrock_httpx_response("message three"), + _passing_bedrock_httpx_response("message four"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["message one", "message two", "message three", "message four"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_when_grounding_present(): + """Contextual-grounding requests are scored holistically against the whole + source; chunking them would silently produce misleading grounding scores. + A too-large error on a grounded request must propagate unchanged, not be + bisected.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_does_not_chunk_on_non_size_validation_error(): + """A 400 for an unrelated validation problem (e.g. a bad guardrail id) must + not trigger chunking -- retrying a bad-config error split into pieces would + just fail twice more and mask the real problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _other_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + assert "not valid" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_unsplittable_text_propagates_original_error(): + """A too-large error on content that has been bisected down to text too + short to split further (< 2 characters) must propagate the original error + rather than looping or crashing.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "a"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 1 + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_on_single_item_splits_by_text_and_succeeds(): + """A too-large error on content that is already down to a single content + item must be bisected by that item's own text (not abandoned), so an + oversized single message can still be scanned successfully in halves.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "one giant single block of text"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + response = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 3 + assert response.get("action") == "NONE" + + +def _raised_bedrock_error(status_code: int, message: str) -> httpx.HTTPStatusError: + """A non-200 the way `AsyncHTTPHandler.post` actually surfaces it. + + That handler calls `response.raise_for_status()`, so in production a non-200 from + Bedrock arrives as a raised `httpx.HTTPStatusError` carrying the response, never + as a returned response object. Tests that return the response instead exercise a + branch real traffic never reaches. A real `httpx.Response` is used rather than a + MagicMock because the transport helper branches on + `isinstance(err_response, httpx.Response)`.""" + response = httpx.Response( + status_code=status_code, + json={"message": message}, + request=httpx.Request("POST", "https://bedrock-runtime.us-east-1.amazonaws.com/guardrail"), + ) + return httpx.HTTPStatusError(message, request=response.request, response=response) + + +_TOO_LARGE_MESSAGE = "Input is too long. Content size exceeds the maximum input size in text units." + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_once_when_client_raises_for_status(): + """The too-large attempt recovered by chunking must still produce exactly one + telemetry entry when the HTTP client raises for status, which is what really + happens: `AsyncHTTPHandler.post` calls `raise_for_status()`. + + Regression for per-attempt `guardrail_failed_to_respond` entries leaking out of + the transport helper on a request that ultimately succeeded, which made a + recovered request look like several failures plus a success.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + assert result.get("action") == "NONE" + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["success"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_still_logs_once_when_client_raises(): + """Suppressing the transport helper's per-attempt logging must not swallow the only + record of a genuine failure: an unsplittable too-large request still has to produce + exactly one `guardrail_failed_to_respond` entry, not zero.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "x"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _post_side_effect(*_args, **_kwargs): + raise _raised_bedrock_error(400, _TOO_LARGE_MESSAGE) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + statuses = [call.kwargs.get("guardrail_status") for call in mock_log.call_args_list] + assert statuses == ["guardrail_failed_to_respond"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_single_item_split_twice_still_yields_one_output_per_item(): + """One oversized content item that needs two levels of text bisection ends up + as four text fragments, and all four must still collapse back into exactly + ONE output entry, because they all came from one original content item. + + Downstream masking (`_apply_masking_to_messages`) walks the merged outputs by + a running index across the original, unchunked message list, so emitting more + than one entry for a single message shifts every later message's masked text + onto the wrong message and drops the surplus. Regression for fragment + grouping assuming fragments only ever arrive as adjacent sibling *pairs*, + which holds for one bisection level but not for two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [{"role": "user", "content": "aaaa bbbb cccc dddd eeee ffff gggg hhhh"}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + responses = [ + _too_large_validation_httpx_response(), # whole single item + _too_large_validation_httpx_response(), # first half + _passing_bedrock_httpx_response("q1"), + _passing_bedrock_httpx_response("q2"), + _too_large_validation_httpx_response(), # second half + _passing_bedrock_httpx_response("q3"), + _passing_bedrock_httpx_response("q4"), + ] + + async def _post_side_effect(*_args, **_kwargs): + return responses.pop(0) + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert mock_post.await_count == 7 + assert not responses + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["q1q2q3q4"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_retries_after_throttling_then_succeeds(): + """A chunk call throttled with a 429 must be retried with backoff and + eventually succeed, rather than surfacing the 429 to the caller.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _throttling_httpx_response() + if call_count == 3: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 4 + mock_sleep.assert_awaited() + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-1", "chunk-2"] + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunking_logs_exactly_once_as_success(): + """A too-large 400 that is recovered by chunking must not leave behind a + 'guardrail_failed_to_respond' telemetry entry for the initial oversized + attempt: the whole logical request (1 too-large attempt + 2 chunk + attempts here) must produce exactly one standard-logging entry, and it + must reflect the eventual success, not the transient too-large failure.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _passing_bedrock_httpx_response("chunk-1") + return _passing_bedrock_httpx_response("chunk-2") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.side_effect = _post_side_effect + + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_apply_guardrail_unrecoverable_failure_logs_exactly_once_as_failed(): + """A too-large error that cannot be recovered (chunking disabled by + contextual grounding) must still log exactly once, as a failure -- not be + silently dropped by the chunking telemetry consolidation.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + { + "role": "system", + "content": [{"type": "grounding_source", "text": "reference source text"}], + }, + {"role": "user", "content": "what does the source say?"}, + ] + model_response = ModelResponse() + model_response.choices = [litellm.Choices(message=litellm.Message(content="a grounded answer", role="assistant"))] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object( + guardrail, + "add_standard_logging_guardrail_information_to_request_data", + ) as mock_log, + ): + mock_post.return_value = _too_large_validation_httpx_response() + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="OUTPUT", + messages=messages, + response=model_response, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_chunk_merge_preserves_masking_position(): + """An earlier chunk that comes back clean (empty `outputs`) must not + shift a later chunk's masked text onto the wrong message. Regression for: + flattening outputs without positional metadata let a later chunk's PII + redaction get applied to the first message while the actual PII-bearing + message (in a later chunk) was forwarded unmasked.""" + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + mask_request_content=True, + ) + + request_data = { + "model": "bedrock-nova-micro", + "messages": [ + {"role": "user", "content": "clean chunk with nothing to mask"}, + {"role": "user", "content": "chunk with PII: John Doe"}, + ], + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + def _clean_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = {"action": "NONE", "assessments": []} + return response + + def _masked_httpx_response() -> MagicMock: + response = MagicMock() + response.status_code = 200 + response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "chunk with PII: [NAME]"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "NAME", "match": "John Doe", "action": "ANONYMIZED"}] + } + } + ], + } + return response + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_validation_httpx_response() + if call_count == 2: + return _clean_httpx_response() + return _masked_httpx_response() + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(), + data=request_data, + call_type="acompletion", + ) + + assert call_count == 3 + updated_messages = request_data["messages"] + assert updated_messages[0]["content"] == "clean chunk with nothing to mask" + assert updated_messages[1]["content"] == "chunk with PII: [NAME]" + + +@pytest.mark.asyncio +async def test_apply_guardrail_accepted_content_costs_exactly_one_call(): + """Content AWS accepts must cost exactly one ApplyGuardrail call, however far over + the chunk budget it is. Chunking is a recovery path, not something every request + pays for. Regression for: bin-packing eagerly on every request, which split + conversations AWS was happy to take whole and multiplied billed calls and guardrail + latency on traffic that never had a size problem.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [{"role": "user", "content": item_text} for _ in range(3)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + return _passing_bedrock_httpx_response(f"batch-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 1 + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_small_content_makes_exactly_one_call(): + """Content that fits entirely within the budget in a single batch must + make exactly one ApplyGuardrail call -- confirms bin-packing does not + introduce an extra probe call for the common (small-request) case.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "short message one"}, + {"role": "user", "content": "short message two"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.return_value = _passing_bedrock_httpx_response("single-batch") + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_post.assert_awaited_once() + assert result.get("action") == "NONE" + + +@pytest.mark.asyncio +async def test_apply_guardrail_batch_under_budget_still_rejected_falls_back_to_bisection(): + """A batch that fits the budget guess but is still rejected by AWS as too large + (a lower real per-account/region/policy cap) must fall back to bisection for that + batch only, and any other batch from the same request that AWS already accepted + must not be re-sent. + + Three half-budget items pack into two batches once the whole-payload probe is + rejected, so the sequence is probe, batch one (rejected), its two halves, batch + two.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + item_text = "x" * (BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2) + messages = [ + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + {"role": "user", "content": item_text}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count in (1, 2): + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response(f"chunk-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 5 + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["chunk-3", "chunk-4", "chunk-5"] + + +def test_split_bedrock_content_single_item_splits_on_whitespace_not_mid_word(): + """A single content item whose raw character midpoint would fall inside a + word must instead split at the nearest whitespace, so neither fragment + ends or begins mid-token. Regression for the Veria AI review finding: a + denied word/PII pattern straddling a raw character-midpoint cut could be + truncated on both fragments and scan clean on each, then reassemble into + the original unmasked text -- a detection bypass.""" + text = ("a" * 20) + " " + ("b" * 30) + raw_midpoint = len(text) // 2 + assert text[raw_midpoint] == "b" + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + + assert first_text + second_text == text + assert first_text == ("a" * 20) + " " + assert second_text == "b" * 30 + + +def test_split_bedrock_content_single_item_with_no_whitespace_falls_back_to_midpoint(): + """A single giant token with no whitespace anywhere has no safe split + point, so the split must fall back to the raw character midpoint rather + than failing or looping.""" + text = "a" * 40 + content = [BedrockContentItem(text=BedrockTextContent(text=text))] + + split_content = BedrockGuardrail._split_bedrock_content(content) + assert split_content is not None + first_half, second_half = split_content + + first_text = first_half[0]["text"]["text"] + second_text = second_half[0]["text"]["text"] + assert first_text + second_text == text + assert len(first_text) == 20 + assert len(second_text) == 20 + + +def test_bin_pack_bedrock_content_packs_minimal_batches_within_budget(): + """Many medium items should pack into the minimal number of in-order + batches that each stay within budget, not one batch per item.""" + items = [BedrockContentItem(text=BedrockTextContent(text="x" * 30)) for _ in range(10)] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert sum(len(batch) for batch in batches) == 10 + for batch in batches: + combined_len = sum(len(item["text"]["text"]) for item in batch) + assert combined_len <= 100 + assert len(batches) == 4 + + +def test_bin_pack_bedrock_content_oversized_single_item_becomes_its_own_batch(): + """An item whose own text already exceeds the budget must not be + pre-split here -- it becomes its own oversized batch, and only the + reactive bisection fallback (on an AWS rejection) may split it later.""" + small_item = BedrockContentItem(text=BedrockTextContent(text="short")) + oversized_item = BedrockContentItem(text=BedrockTextContent(text="x" * 200)) + items = [small_item, oversized_item, small_item] + + batches = BedrockGuardrail._bin_pack_bedrock_content(items, budget=100) + + assert batches == ((small_item,), (oversized_item,), (small_item,)) + + +def test_bin_pack_bedrock_content_empty_content_makes_exactly_one_empty_batch(): + """Empty content must still pack into exactly one (empty) batch, matching + pre-bin-packing behavior of sending the content list as-is in one call -- + bin-packing must not turn an empty request into zero ApplyGuardrail calls.""" + assert BedrockGuardrail._bin_pack_bedrock_content([], budget=100) == ((),) + + +@pytest.mark.asyncio +async def test_apply_guardrail_too_large_reported_as_429_bisects_without_burning_retries(): + """AWS reports an oversized ApplyGuardrail request as a 429 ThrottlingException + (not the documented 400 ValidationException) when the guardrail has an active + content-filter policy. That is not a transient throttle -- re-posting the same + oversized content can never succeed -- so it must bisect immediately instead of + consuming the exponential-backoff retry budget first. + + Regression for a bug found against a live guardrail: because the throttle retry + only keyed off status 429, every oversized chunk burned all + _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES attempts (each a billed AWS call, + each preceded by a backoff sleep) before bisection got a chance, at every level + of the recursion.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + messages = [ + {"role": "user", "content": "chunk one text"}, + {"role": "user", "content": "chunk two text"}, + ] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + call_count = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal call_count + call_count += 1 + if call_count == 1: + return _too_large_throttling_httpx_response() + return _passing_bedrock_httpx_response(f"half-{call_count}") + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.asyncio.sleep", + new_callable=AsyncMock, + ) as mock_sleep, + ): + mock_post.side_effect = _post_side_effect + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert call_count == 3 + mock_sleep.assert_not_awaited() + assert result.get("action") == "NONE" + output_texts = [o.get("text") for o in result.get("outputs") or []] + assert output_texts == ["half-2", "half-3"] + + +def test_chunk_budget_defaults_to_apply_guardrail_per_second_quota(): + """The default budget must track ApplyGuardrail's default quota of 25 text units + (about 1,000 characters each) per second. Packing to that size and posting + sequentially is what stops chunking from trading a size error for a throttle, so + this default is a deliberate match to AWS behaviour rather than an arbitrary + number.""" + assert BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS == 25_000 + assert BedrockGuardrail(guardrailIdentifier="g", guardrailVersion="DRAFT").chunk_budget_chars == 25_000 + + +@pytest.mark.asyncio +async def test_configured_chunk_budget_changes_how_content_is_packed(): + """An account with raised quotas can set a larger `chunk_budget_chars` and have it + actually drive packing once AWS has rejected a payload, spending fewer + ApplyGuardrail calls for the same content instead of being pinned to the + conservative default. + + Four 20,000-character messages are 80,000 characters total, and every call here is + preceded by the one whole-payload probe AWS rejects. At the 25,000 default only one + message fits per batch, so it is the probe plus four; at 50,000 two fit per batch, + so it is the probe plus two.""" + messages = [{"role": "user", "content": "x" * 20_000} for _ in range(4)] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + async def _calls_made_with_budget(budget: int) -> int: + guardrail = BedrockGuardrail( + guardrail_name="test-bedrock-guard", + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + disable_exception_on_block=False, + chunk_budget_chars=budget, + ) + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + posted = 0 + + async def _post_side_effect(*_args, **_kwargs): + nonlocal posted + posted += 1 + if posted == 1: + return _too_large_validation_httpx_response() + return _passing_bedrock_httpx_response("ok") + + mock_post.side_effect = _post_side_effect + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + return mock_post.await_count + + assert await _calls_made_with_budget(BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS) == 5 + assert await _calls_made_with_budget(50_000) == 3 + + +def test_split_index_never_produces_an_empty_fragment(): + """Both fragments must be non-empty for every splittable text, so bisection always + makes progress. + + A text whose only qualifying whitespace is its final character is the dangerous + shape: taking that boundary puts the split at len(text), leaving the first fragment + identical to the input that was just rejected and the second empty. The recursion + would then resubmit the unchanged fragment forever and exhaust the stack instead of + scanning or surfacing Bedrock's error.""" + for text in ("ab ", "xxxx ", ("x" * 40) + " ", " ab", "a b", "ab", " "): + split_at = BedrockGuardrail._nearest_whitespace_split_index(text) + assert 0 < split_at < len(text), f"degenerate split {split_at} for {text!r}" + assert text[:split_at] and text[split_at:], f"empty fragment for {text!r}" + assert text[:split_at] + text[split_at:] == text + + +@pytest.mark.asyncio +async def test_oversized_single_item_with_trailing_space_gives_up_instead_of_recursing(): + """An oversized single item whose only space is trailing must bottom out and + surface Bedrock's error, not recurse forever. + + AWS is modelled the way it really behaves, rejecting every attempt, because the + danger is a fragment identical to the input that was just rejected: AWS would + reject it again, and each retry would split it into the same unchanged fragment. + A split that always shrinks the text terminates and re-raises; one that can return + the whole text raises RecursionError instead. The call-count bound is generous: + halving 41 characters down to unsplittable is a handful of attempts, nowhere near + a stack limit.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + messages = [{"role": "user", "content": ("x" * 40) + " "}] + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + ): + mock_post.side_effect = lambda *_a, **_k: _too_large_validation_httpx_response() + + with pytest.raises(HTTPException) as excinfo: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=messages, + request_data={"model": "bedrock-nova-micro"}, + ) + + assert excinfo.value.status_code == 400 + assert mock_post.await_count < 200 + + class TestBedrockOnlyScanNewMessages: """Bedrock apply_guardrail honors only_scan_new_messages: scans only the per-session diff. @@ -3559,14 +4379,10 @@ class TestBedrockIncrementalFlagInteractions: session = {"litellm_session_id": "sess-flags-mask"} with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: mock_api.return_value = {"action": "NONE", "output": [], "outputs": []} - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1 mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 1, "masking mode must re-scan every turn, exactly once" @pytest.mark.asyncio @@ -3586,9 +4402,7 @@ class TestBedrockIncrementalFlagInteractions: assert mock_api.call_count == 2, "incremental attempt + full-scan fallback" assert result["texts"] == ["MASKED q1"], "masked content must be applied" mock_api.reset_mock() - await guardrail.apply_guardrail( - inputs={"texts": ["q1"]}, request_data=session, input_type="request" - ) + await guardrail.apply_guardrail(inputs={"texts": ["q1"]}, request_data=session, input_type="request") assert mock_api.call_count == 2, "no hashes persisted, so the double scan repeats" @pytest.mark.asyncio @@ -3647,9 +4461,7 @@ async def test_moderation_hook_honors_the_mcp_event_type(mode, call_type, should } with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: - mock_api.return_value = MagicMock( - action="NONE", output=[], outputs=[], assessments=[] - ) + mock_api.return_value = MagicMock(action="NONE", output=[], outputs=[], assessments=[]) await guardrail.async_moderation_hook( data=data, user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", user_id="u"), @@ -3707,3 +4519,146 @@ class TestScanOnlyToolResultsWithLatestRoleFilter: assert result["texts"] == ["TOOL-RESULT"] warning_text = " ".join(str(arg) for c in mock_warning.call_args_list for arg in c.args) assert "scan_only_tool_results" in warning_text + + +@pytest.mark.parametrize("separator", ["\n", "\t", "\r\n", "　"]) +def test_split_bedrock_content_splits_on_any_whitespace_not_just_space(separator): + """Regression: the midpoint split must land on any Unicode whitespace, not only an + ASCII space. + + Matching only " " left the boundary unguarded for exactly the payloads that grow + large enough to need splitting: JSON lines, source code, logs and transcripts are + newline or tab delimited. A deny-listed word sitting at the midpoint of one was cut + in half, scanned clean on both fragments, and reassembled intact, which is the + single-token detection bypass the whitespace split exists to close.""" + text = separator.join(["aaaaaaa"] * 4) + separator + "BADWORDXYZ" + separator + separator.join(["bbbbbbb"] * 4) + + first, second = BedrockGuardrail._split_bedrock_content([BedrockContentItem(text=BedrockTextContent(text=text))]) + + first_text = first[0]["text"]["text"] + second_text = second[0]["text"]["text"] + assert first_text + second_text == text, "split must stay lossless" + assert "BADWORDXYZ" in first_text or "BADWORDXYZ" in second_text, "split severed the token" + + +def test_merge_bedrock_responses_preserves_fields_the_merge_has_no_opinion_on(): + """Regression: merging must not drop AWS response fields it does not itself merge. + + The merged response used to be rebuilt from an empty dict holding only action, + outputs, assessments and usage, so actionReason, guardrailCoverage and anything AWS + adds later vanished from the guardrail_json_response the Admin UI renders, on every + ApplyGuardrail request rather than only chunked ones.""" + chunk = BedrockContentChunkResult( + response={ + "action": "NONE", + "actionReason": "No action.", + "guardrailCoverage": {"textCharacters": {"guarded": 41, "total": 41}}, + "usage": {"contentPolicyUnits": 1}, + }, + content=[BedrockContentItem(text=BedrockTextContent(text="hello"))], + fragment_group_size=1, + ) + + merged = BedrockGuardrail._merge_bedrock_guardrail_responses([chunk]) + + assert merged["actionReason"] == "No action." + assert merged["guardrailCoverage"] == {"textCharacters": {"guarded": 41, "total": 41}} + + +def test_merge_bedrock_usage_sums_counters_not_on_the_known_list(): + """Regression: usage counters were summed from a hardcoded list of six keys, so the + ones AWS also returns (contentPolicyImageUnits, the automatedReasoning pair) were + reported as absent no matter what the chunks actually used.""" + chunks = [ + BedrockContentChunkResult( + response={"action": "NONE", "usage": {"contentPolicyImageUnits": units, "contentPolicyUnits": 1}}, + content=[BedrockContentItem(text=BedrockTextContent(text="x"))], + fragment_group_size=1, + ) + for units in (3, 4) + ] + + usage = BedrockGuardrail._merge_bedrock_guardrail_responses(chunks)["usage"] + + assert usage["contentPolicyImageUnits"] == 7 + assert usage["contentPolicyUnits"] == 2 + + +@pytest.mark.asyncio +async def test_apply_guardrail_exception_inside_200_logs_failure_and_proceeds(): + """Regression: AWS can report a failure inside an HTTP 200 body via Output.__type, + and that must be logged as guardrail_failed_to_respond rather than success. + + Real AWS does this: an unrecognised operation path on bedrock-runtime answers + HTTP 200 with {"Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}}. + Consolidating telemetry had replaced the derived status with a hardcoded "success", + which reported a failed scan as a clean one. The request itself still proceeds, which + is the behaviour of the code before chunking existed.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + exception_response = MagicMock() + exception_response.status_code = 200 + exception_response.json.return_value = { + "Output": {"__type": "com.amazon.coral.service#UnknownOperationException"}, + "Version": "1.0", + } + exception_response.text = json.dumps(exception_response.json.return_value) + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.return_value = exception_response + + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + assert result is not None, "the request proceeds, as it did before chunking existed" + mock_log.assert_called_once() + assert mock_log.call_args.kwargs["guardrail_status"] == "guardrail_failed_to_respond" + + +@pytest.mark.asyncio +async def test_apply_guardrail_failure_logs_a_dict_not_a_bare_string(): + """Regression: the consolidated failure logger must log guardrail_json_response as a + dict, the shape the pre-chunking code and the InvokeGuardrailChecks path both use. + + Consolidating telemetry had changed it to a bare string on the ApplyGuardrail path + only, which breaks any consumer that reads it as a mapping and leaves the two paths + in this file inconsistent.""" + guardrail = _bedrock_guardrail_for_chunk_tests() + + mock_credentials = MagicMock() + mock_credentials.access_key = "k" + mock_credentials.secret_key = "s" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()), + patch.object(guardrail, "add_standard_logging_guardrail_information_to_request_data") as mock_log, + ): + mock_post.side_effect = _raised_bedrock_error(400, "guardrailIdentifier is not valid") + + with pytest.raises(HTTPException): + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=[{"role": "user", "content": "hello"}], + request_data={"model": "bedrock-nova-micro"}, + ) + + mock_log.assert_called_once() + logged = mock_log.call_args.kwargs["guardrail_json_response"] + assert isinstance(logged, dict), f"expected a dict, got {type(logged).__name__}" + assert "error" in logged diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index 71e775842e3..8edb56ce25e 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py @@ -38,6 +38,41 @@ def test_initialize_presidio_guardrail(): assert result["litellm_params"].mode == "pre_call" +def test_initialize_bedrock_forwards_chunk_budget_chars(): + """Regression: `chunk_budget_chars` set in config.yaml must reach the guardrail. + + The field lives on BedrockGuardrailConfigModel, so LitellmParams parsed it and the + Admin UI rendered it, but initialize_bedrock enumerates its kwargs explicitly and + dropped it. The setting validated and then silently did nothing. Asserting through + initialize_guardrail rather than the constructor is the point: constructing + BedrockGuardrail directly bypasses the only path a user can actually reach. + """ + import litellm + from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail + + test_guardrail = { + "guardrail_name": "test_bedrock_chunk_budget", + "litellm_params": { + "guardrail": SupportedGuardrailIntegrations.BEDROCK.value, + "mode": "pre_call", + "guardrailIdentifier": "test-guardrail", + "guardrailVersion": "DRAFT", + "chunk_budget_chars": 60_000, + }, + } + + guardrail_handler = InMemoryGuardrailHandler() + guardrail_handler.initialize_guardrail(guardrail=test_guardrail) + + initialized = [ + callback + for callback in litellm.callbacks + if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_chunk_budget" + ] + assert initialized, "bedrock guardrail was not registered as a callback" + assert initialized[-1].chunk_budget_chars == 60_000 + + def test_initialize_guardrail_preserves_guardrail_info(): """ Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the