scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-10 03:28:53 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 85

refactor(guardrails): move Bedrock chunking rationale into docstrings

Browse source 
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.
This commit is contained in:
spencer-burridge 2026-07-29 13:48:48 -05:00
parent 2b449f0cd0
commit c64a02c4a3
2 changed files with 47 additions and 102 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -84,22 +84,7 @@ from litellm.types.utils import (
)
GUARDRAIL_NAME = "bedrock"
# KNOWN LIMITATION (chunking, below): splitting an oversized message's text on
# a whitespace boundary (see `_nearest_whitespace_split_index`) prevents
# accidentally severing a single token -- one denied word, one PII pattern --
# across a chunk boundary. It does not stop a multi-word denied phrase
# deliberately positioned to straddle that boundary, since each fragment can
# scan clean independently. AWS's own guidance for this API acknowledges the
# same gap for input chunking with no documented resolution; closing it would
# require an overlap window reconciled against masked output, which AWS does
# not guarantee to be length-preserving. Accepted as out of scope.
_BEDROCK_DYNAMIC_BODY_DENYLIST = frozenset({"content", "source"})
# ApplyGuardrail's per-request "maximum input size in text units" quota is
# region/account/policy-dependent and cannot be predicted from config, so it is
# only ever discovered reactively: AWS rejects an over-quota call with a 400
# ValidationException whose message contains one of these substrings (matched
# case-insensitively against the parsed AWS error message). On a match the
# content is bisected and each half retried, rather than surfacing the 400.
_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS = (
"text unit",
"maximum input size",
@ -108,21 +93,7 @@ _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS = (
"too large",
"exceeds the maximum",
)
# Conservative starting guess for how much content (by character count) to send
# in one ApplyGuardrail call, used to pre-bin-pack content instead of always
# starting from the whole payload. This is NOT a correctness dependency: it only
# sets how many calls the common case takes. Any bin AWS still rejects as too
# large (because the real per-request text-unit cap for this account/region/
# policy is lower than this guess -- that cap is not knowable ahead of time and
# is not a fixed character count) falls back to the recursive bisection below,
# which self-corrects regardless of how wrong this guess was. So a too-generous
# guess here costs the same one extra probe-and-bisect round trip that pure
# reactive bisection would have paid anyway, while a well-tuned guess makes the
# common case a single pass instead of O(log n) round trips per request.
_BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS = 20_000
# Exponential backoff for a chunk call throttled with ThrottlingException (429).
# Kept small: chunking already trades one oversized call for several smaller
# ones, so retries must not multiply per-request latency by an order of magnitude.
_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES = 3
_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS = 0.5
# Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required).
@ -826,6 +797,30 @@ 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 = datetime.now(timezone.utc)
credentials, aws_region_name = self._load_credentials()
bedrock_request_data: dict = dict(
@ -844,18 +839,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
if request_data.get("api_key") is not None:
api_key = request_data["api_key"]
# 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
content: list[BedrockContentItem] = bedrock_request_data.get("content") or []
# Contextual grounding scores the response holistically against the whole
# reference source; bisecting it would fragment that evaluation and produce
# misleading grounding scores, so a too-large error is never chunked here.
allow_chunking = not self._content_uses_contextual_grounding(content)
batches = (
self._bin_pack_bedrock_content(content, budget=_BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS)
@ -880,11 +869,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
]
except HTTPException as exc:
# A block is logged where it happens, inside _post_apply_guardrail_content,
# since chunking stops immediately and there is no later merged response to
# log instead. Anything else reaching here (unrecoverable too-large error,
# a non-size validation error, exhausted throttle retries) is a genuine
# end-to-end failure of this logical guardrail call and is logged once here.
if not isinstance(exc.detail, dict):
self._log_apply_guardrail_failure(
detail=exc.detail,
@ -930,7 +914,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
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. A real
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.
@ -991,11 +977,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
combined_results = first_results + second_results
if is_single_item_text_split:
# Every leaf below this point came from one content item's own
# text, however many levels deep the splitting went. Stamping
# the total count on all of them (overwriting any smaller count
# an inner split set) is what lets the merge step regroup them
# into exactly one output entry for that one item.
return [result._replace(fragment_group_size=len(combined_results)) for result in combined_results]
return combined_results
raise
@ -1064,6 +1045,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
"""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).
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 = {**base_request_data, "content": content}
prepared_request = self._prepare_request(
@ -1097,9 +1082,6 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
)
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
# A block ends the whole chunking flow immediately (no further
# chunks are attempted), so it is logged here rather than by the
# caller -- there is no later "final merged response" to log instead.
self._log_apply_guardrail_attempt(
httpx_response=httpx_response,
json_response=_json_response,
@ -1216,6 +1198,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
(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` (`_BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS` at the call site)
is a conservative starting guess, not a correctness dependency. It only
sets how many calls the common case takes. AWS's real per-request
text-unit cap depends on 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
guess was. A too-generous guess therefore costs one extra probe-and-bisect
round trip, the same one pure reactive bisection would have paid anyway.
"""
if not content:
return [content]
@ -1457,12 +1448,15 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
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:
# Every result here is one fragment of a single content item's text, so
# the whole group collapses to one entry: each fragment's masked text
# (or its own original text when that fragment came back unmasked),
# concatenated in order. Holds for any group size, not just two.
def fragment_outputs(result: BedrockContentChunkResult) -> list[BedrockGuardrailOutput]:
return list(result.response.get("outputs") or result.response.get("output") or [])

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

@ -3282,17 +3282,6 @@ async def test_chat_completion_modify_response_exception_streaming_logging_obj_n
assert response is not None
###############################################################################
# AIKMG-278: chunk oversized ApplyGuardrail requests instead of failing.
#
# AWS rejects an ApplyGuardrail call whose content exceeds the account's
# "maximum input size in text units" quota with a 400 ValidationException.
# Rather than surfacing that 400 to the caller, split the content in half and
# retry each half; merge the per-half responses back into one. Grounded
# (contextual-grounding) requests are never chunked -- grounding scores the
# response holistically against the whole source, so fragmenting it would
# produce misleading scores.
###############################################################################
def _too_large_validation_httpx_response() -> MagicMock:
@ -3399,7 +3388,6 @@ async def test_apply_guardrail_chunks_on_too_large_validation_error():
nonlocal call_count
call_count += 1
if call_count == 1:
# Whole-content call: too large.
return _too_large_validation_httpx_response()
if call_count == 2:
return _passing_bedrock_httpx_response("chunk-1")
@ -3419,13 +3407,9 @@ async def test_apply_guardrail_chunks_on_too_large_validation_error():
request_data={"model": "bedrock-nova-micro"},
)
# 1 whole-content attempt + 2 chunk attempts.
assert call_count == 3
detail = exc_info.value.detail
assert exc_info.value.status_code == 400
# The merged response must retain the block signal from chunk-2 even
# though chunk-1 passed clean -- losing it would silently bypass a
# guardrail hit.
assert "chunk-2" in detail["bedrock_guardrail_response"]
assert detail["assessments"][0]["matches"][0]["name"] == "chunk-2"
@ -3499,9 +3483,6 @@ async def test_apply_guardrail_recurses_past_first_bisection_into_four_chunks():
mock_credentials.secret_key = "s"
mock_credentials.token = None
# Depth-first recursion order for 4 items bisected down to single items:
# whole -> [1,2] -> [1] -> [2] -> [3,4] -> [3] -> [4]. Only single-item
# calls (positions 3, 4, 6, 7) fit; every multi-item call is too large.
responses = [
_too_large_validation_httpx_response(), # whole content: [1,2,3,4]
_too_large_validation_httpx_response(), # first half: [1,2]
@ -3528,7 +3509,6 @@ async def test_apply_guardrail_recurses_past_first_bisection_into_four_chunks():
request_data={"model": "bedrock-nova-micro"},
)
# 1 whole-content + 2 first-level halves + 4 second-level quarters = 7.
assert mock_post.await_count == 7
assert not responses
assert result.get("action") == "NONE"
@ -3576,7 +3556,6 @@ async def test_apply_guardrail_does_not_chunk_when_grounding_present():
request_data={"model": "bedrock-nova-micro"},
)
# Exactly one call: no bisect-and-retry for a grounded request.
assert mock_post.await_count == 1
assert exc_info.value.status_code == 400
@ -3769,8 +3748,6 @@ async def test_apply_guardrail_unrecoverable_failure_still_logs_once_when_client
exactly one `guardrail_failed_to_respond` entry, not zero."""
guardrail = _bedrock_guardrail_for_chunk_tests()
# Single character: `_split_bedrock_content` cannot halve this into two non-empty
# pieces, so chunking gives up and the original error propagates.
messages = [{"role": "user", "content": "x"}]
mock_credentials = MagicMock()
@ -3824,8 +3801,6 @@ async def test_apply_guardrail_single_item_split_twice_still_yields_one_output_p
mock_credentials.secret_key = "s"
mock_credentials.token = None
# whole item -> [first half] -> [q1] [q2] -> [second half] -> [q3] [q4].
# Only the four quarters fit; the whole item and both halves are too large.
responses = [
_too_large_validation_httpx_response(), # whole single item
_too_large_validation_httpx_response(), # first half
@ -3856,8 +3831,6 @@ async def test_apply_guardrail_single_item_split_twice_still_yields_one_output_p
assert not responses
assert result.get("action") == "NONE"
output_texts = [o.get("text") for o in result.get("outputs") or []]
# One original content item in, so exactly one output entry out, carrying all
# four fragments' text in order.
assert output_texts == ["q1q2q3q4"]
@ -3885,7 +3858,6 @@ async def test_apply_guardrail_chunk_retries_after_throttling_then_succeeds():
if call_count == 1:
return _too_large_validation_httpx_response()
if call_count == 2:
# First chunk throttled once, then succeeds on retry.
return _throttling_httpx_response()
if call_count == 3:
return _passing_bedrock_httpx_response("chunk-1")
@ -4105,9 +4077,6 @@ async def test_apply_guardrail_bin_packs_under_budget_content_with_no_probe_call
attempt."""
guardrail = _bedrock_guardrail_for_chunk_tests()
# Five items, individually tiny, whose combined length exceeds the budget
# only when summed -- proves this triggers packing into multiple batches
# by SIZE, not by falling back to per-item chunking.
item_text = "x" * (_BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS // 2)
messages = [{"role": "user", "content": item_text} for _ in range(3)]
@ -4136,8 +4105,6 @@ async def test_apply_guardrail_bin_packs_under_budget_content_with_no_probe_call
request_data={"model": "bedrock-nova-micro"},
)
# Three items at budget/2 each pack two-per-batch (2 + 1), never a single
# oversized call and never a wasted whole-content probe: exactly 2 calls.
assert call_count == 2
assert result.get("action") == "NONE"
@ -4202,12 +4169,8 @@ async def test_apply_guardrail_batch_under_budget_still_rejected_falls_back_to_b
nonlocal call_count
call_count += 1
if call_count == 1:
# First pre-packed batch (items 1+2): accepted immediately.
return _passing_bedrock_httpx_response("batch-1")
if call_count == 2:
# Second pre-packed batch (item 3 alone): rejected as too large
# despite fitting the fixed budget guess -- simulates a lower
# real-world per-account cap.
return _too_large_validation_httpx_response()
return _passing_bedrock_httpx_response(f"batch-2-bisected-{call_count}")
@ -4224,12 +4187,9 @@ async def test_apply_guardrail_batch_under_budget_still_rejected_falls_back_to_b
request_data={"model": "bedrock-nova-micro"},
)
# batch-1 (1 call, accepted) + batch-2 (1 rejected + 2 bisected halves) = 4.
assert call_count == 4
assert result.get("action") == "NONE"
output_texts = [o.get("text") for o in result.get("outputs") or []]
# batch-2's two bisected text fragments came from the same original
# content item, so they are merged back into one combined output entry.
assert output_texts == ["batch-1", "batch-2-bisected-3batch-2-bisected-4"]
@ -4240,8 +4200,6 @@ def test_split_bedrock_content_single_item_splits_on_whitespace_not_mid_word():
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."""
# 20 'a's + space + 30 'b's: the raw character midpoint (25) falls inside
# the run of 'b's, proving the split must move off it to the nearest space.
text = ("a" * 20) + " " + ("b" * 30)
raw_midpoint = len(text) // 2
assert text[raw_midpoint] == "b"
@ -4254,10 +4212,7 @@ def test_split_bedrock_content_single_item_splits_on_whitespace_not_mid_word():
first_text = first_half[0]["text"]["text"]
second_text = second_half[0]["text"]["text"]
# Lossless: concatenating the two fragments reproduces the original exactly.
assert first_text + second_text == text
# Word-safe: the split lands exactly on the whitespace boundary, not
# inside either the "a" or "b" run.
assert first_text == ("a" * 20) + " "
assert second_text == "b" * 30
@ -4291,8 +4246,6 @@ def test_bin_pack_bedrock_content_packs_minimal_batches_within_budget():
for batch in batches:
combined_len = sum(len(item["text"]["text"]) for item in batch)
assert combined_len <= 100
# 10 items * 30 chars = 300 chars at a 100-char budget packs into 3 batches
# of 3 items (90 chars) plus 1 batch of 1 item -- never one batch per item.
assert len(batches) == 4
@ -4367,9 +4320,7 @@ async def test_apply_guardrail_too_large_reported_as_429_bisects_without_burning
request_data={"model": "bedrock-nova-micro"},
)
# One rejected whole-content call, then exactly one call per bisected half.
assert call_count == 3
# No backoff sleep: a size error must not be treated as a transient throttle.
mock_sleep.assert_not_awaited()
assert result.get("action") == "NONE"
output_texts = [o.get("text") for o in result.get("outputs") or []]