scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

fix(guardrails): chunk oversized Bedrock ApplyGuardrail requests instead of failing (#36119)

Browse source 
* 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 ebf8ba2fd5.

* 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>
This commit is contained in:
yucheng-berri 2026-08-07 19:44:24 -07:00 • committed by GitHub
parent 0a606cb258
commit d4dc2c39e7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 2326 additions and 570 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/constants.py
View file

@ -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)))

857
litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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(),

1
litellm/proxy/guardrails/guardrail_initializers.py
View file

@ -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,

10
litellm/types/guardrails.py
View file

@ -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):

3
litellm/types/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py
View file

@ -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):

1989
tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py
View file

File diff suppressed because it is too large Load diff

35
tests/test_litellm/proxy/guardrails/test_init_guardrails.py
View file

@ -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