Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py: 10%
1125 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-10 12:01 +0000
1# +-------------------------------------------------------------+
2#
3# Use Bedrock Guardrails for your LLM calls
4#
5# +-------------------------------------------------------------+
6# Thank you users! We ❤️ you! - Krrish & Ishaan
8import os
9import sys
11sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path
12import asyncio
13import copy
14import json
15import re
16import sys
17import time
18from collections.abc import AsyncGenerator, Mapping, Sequence
19from datetime import datetime, timezone
20from itertools import accumulate, groupby
21from types import MappingProxyType
22from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, NamedTuple, Optional, cast
24import httpx
25from fastapi import HTTPException
26from pydantic import TypeAdapter, ValidationError
28import litellm
29from litellm._logging import verbose_proxy_logger
30from litellm.caching import DualCache
31from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
32from litellm.exceptions import ModifyResponseException
33from litellm.integrations.custom_guardrail import CustomGuardrail
34from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route
35from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
36from litellm.litellm_core_utils.litellm_logging import (
37 _get_masked_values, # pyright: ignore[reportPrivateUsage] # the shared header-masking helper has no public name
38)
39from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import (
40 bedrock_guardrail_cost_by_unit,
41 guardrail_cost_total,
42)
43from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicMessagesHandler
44from litellm.llms.base_llm.guardrail_translation.utils import (
45 effective_scan_only_tool_results_for_guardrail,
46)
47from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token, run_aws_signing
48from litellm.llms.custom_httpx.http_handler import (
49 get_async_httpx_client,
50 httpxSpecialProvider,
51)
52from litellm.proxy._types import UserAPIKeyAuth
53from litellm.proxy.common_request_processing import serialize_http_exception_detail
54from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired
55from litellm.proxy.guardrails.anthropic_sse import (
56 anthropic_sse_chunks_from_response,
57 anthropic_sse_error_frames,
58 assemble_anthropic_sse_stream,
59 is_raw_sse_stream,
60 model_response_text,
61)
62from litellm.types.guardrails import (
63 BedrockChecksConfigModel,
64 BedrockGuardrailStreamingParams,
65 GuardrailEventHooks,
66 LitellmParams,
67)
68from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
69from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import (
70 BedrockChecksMessage,
71 BedrockChecksViolation,
72 BedrockContentItem,
73 BedrockGuardrailChecksResponse,
74 BedrockGuardrailOutput,
75 BedrockGuardrailQualifier,
76 BedrockGuardrailResponse,
77 BedrockGuardrailUsage,
78 BedrockRequest,
79 BedrockTextContent,
80)
81from litellm.types.utils import GenericGuardrailAPIInputs
83if TYPE_CHECKING: 83 ↛ 84line 83 didn't jump to line 84 because the condition on line 83 was never true
84 from botocore.awsrequest import AWSPreparedRequest
85 from botocore.credentials import Credentials
87 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
89from litellm.types.utils import (
90 CallTypes,
91 CallTypesLiteral,
92 Choices,
93 GuardrailStatus,
94 GuardrailTracingDetail,
95 Message,
96 ModelResponse,
97 ModelResponseStream,
98 StreamingChoices,
99 TextChoices,
100)
102GUARDRAIL_NAME: Final = "bedrock"
103_BEDROCK_DYNAMIC_BODY_DENYLIST: Final = frozenset({"content", "source"})
104_BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS: Final = (
105 "text unit",
106 "maximum input size",
107 "content size",
108 "too long",
109 "too large",
110 "exceeds the maximum",
111)
112_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES: Final = 3
113_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS: Final = 0.5
114_BEDROCK_WHITESPACE: Final = re.compile(r"\s")
115_NO_TRACING_DETAIL: Final[GuardrailTracingDetail] = {}
116# Resource-less, detect-only InvokeGuardrailChecks API (no guardrail resource required).
117_BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH: Final = "/guardrail-checks/invoke"
118# InvokeGuardrailChecks accepts at most 10 content blocks per message. A message with
119# more text blocks is split across multiple messages so ALL content is scanned --
120# never truncated (truncation would let a user hide content past the limit).
121_BEDROCK_CHECKS_MAX_CONTENT_BLOCKS: Final = 10
122_BEDROCK_CHECKS_KNOWN_KEYS: Final = frozenset({"contentFilter", "promptAttack", "sensitiveInformation"})
123# Keys in a sensitiveInformation result that pinpoint the PII location. They are
124# stripped before the response is handed to standard logging / telemetry so the
125# detected PII span cannot be reconstructed from logs.
126_BEDROCK_CHECKS_PII_LOCATION_KEYS: Final = (
127 "beginOffset",
128 "endOffset",
129 "messageIndex",
130 "contentIndex",
131)
133# Maps an OpenAI message content-block ``type`` to the Bedrock guardrail qualifier
134# it represents, so callers can drive contextual grounding by tagging their content.
135# The model response is qualified as ``guard_content`` directly by the OUTPUT builder;
136# the existing ``guarded_text`` marker is intentionally left unmapped here so its
137# guardrail-hook payload is unchanged by this feature.
138_CONTENT_TYPE_TO_QUALIFIER: Final[dict[str, BedrockGuardrailQualifier]] = {
139 "grounding_source": "grounding_source",
140 "query": "query",
141}
143# Roles whose ``grounding_source`` blocks are trusted as reference material for the
144# contextual-grounding check. Only app-authored roles qualify: ``tool``/``function``
145# results and ``user`` content can carry caller- or externally-influenced text, which
146# must not be graded against as if it were the application's own source material.
147_GROUNDING_SOURCE_TRUSTED_ROLES: Final = frozenset({"system", "developer"})
150class QualifiedTextBlock(NamedTuple):
151 """A piece of message text paired with its Bedrock grounding qualifier (if any)."""
153 text: str
154 qualifier: BedrockGuardrailQualifier | None
157class GuardrailMessageFilterResult(NamedTuple):
158 payload_messages: list[AllMessageValues] | None
159 original_messages: list[AllMessageValues] | None
160 target_indices: list[int] | None
163class BedrockContentChunkResult(NamedTuple):
164 """One chunk's ApplyGuardrail response, paired with enough bookkeeping to
165 reconstruct global masked-output positions once every chunk is back.
167 `content` is the exact content items this chunk was called with -- needed
168 so an all-clear chunk (empty `outputs`) can still contribute one unmasked
169 placeholder per item it covers, keeping every later chunk's masked text
170 aligned to its original global position. `fragment_group_size` is 1 for an
171 ordinary chunk, and otherwise the total number of consecutive chunk results
172 that together make up ONE original content item's own text (split because a
173 list of length 1 could not be bisected by list length). All of them must be
174 concatenated back into that one item's masked output rather than treated as
175 separate items. It is a count rather than a boolean because one item can be
176 bisected more than once: two levels of splitting produce four fragments for
177 a single item, not two, and grouping them in fixed pairs would emit two
178 outputs for one message and shift every later message's masked text.
179 """
181 response: BedrockGuardrailResponse
182 content: tuple[BedrockContentItem, ...]
183 fragment_group_size: int
186class ApplyGuardrailMessageSelection(NamedTuple):
187 """Messages selected for an apply_guardrail scan + write-back metadata."""
189 filtered_messages: list[AllMessageValues] | None
190 # Slice of the flat `texts` list actually scanned (offset, length),
191 # used to write masked content back to the right positions. None = whole list.
192 scanned_slice: tuple[int, int] | None
193 # True when messages were selected by their original role.
194 scanned_role_subset: bool
195 # True when there is nothing to scan (e.g. no user-role message).
196 skip_scan: bool = False
199def _redact_pii_matches(response_json: dict) -> dict:
200 """
201 Redact match-like fields from a Bedrock ApplyGuardrail JSON payload.
203 Delegates to :func:`redact_nested_match_and_regex_keys` (same rules as spend
204 logging). Kept as a Bedrock-module entry point for existing unit tests.
205 """
206 redacted: Final = redact_nested_match_and_regex_keys(response_json)
207 return redacted if isinstance(redacted, dict) else response_json
210def _redact_assessment_match_fields(assessments: list[dict]) -> list[dict]:
211 """
212 Redact sensitive match-like fields from blocked assessment summaries.
214 This is used for customer-visible error payloads (HTTPException.detail) where
215 we want to preserve policy/type/action metadata without echoing raw matched
216 content.
217 """
218 redacted: Final = redact_nested_match_and_regex_keys(assessments)
219 return redacted if isinstance(redacted, list) else assessments
222_RESPONSES_API_CALL_TYPES: Final = frozenset({CallTypes.responses, CallTypes.aresponses})
225def _is_responses_api_route(request_route: str | None) -> bool:
226 if request_route is None:
227 return False
228 call_types: Final = get_call_types_for_route(request_route)
229 return call_types is not None and any(call_type in _RESPONSES_API_CALL_TYPES for call_type in call_types)
232class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
233 # During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
234 # OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
235 use_native_during_call_hook: ClassVar[bool] = True
237 def __init__(
238 self,
239 guardrailIdentifier: str | None = None,
240 guardrailVersion: str | None = None,
241 disable_exception_on_block: bool | None = False,
242 checks: BedrockChecksConfigModel | Mapping[str, object] | None = None,
243 content_filter_threshold: float | None = 0.5,
244 prompt_attack_threshold: float | None = 0.5,
245 pii_confidence_threshold: float | None = 0.5,
246 chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
247 contextual_grounding_from_messages: bool = False,
248 streaming_buffer_until_moderated: bool | None = None,
249 streaming_sampling_rate: int | None = None,
250 streaming_end_of_stream_only: bool | None = None,
251 streaming_buffer_release_on_scan: bool | None = None,
252 **kwargs,
253 ):
254 self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback)
255 self._set_streaming_params(
256 BedrockGuardrailStreamingParams.from_extras(
257 MappingProxyType(
258 {
259 "streaming_buffer_until_moderated": streaming_buffer_until_moderated,
260 "streaming_sampling_rate": streaming_sampling_rate,
261 "streaming_end_of_stream_only": streaming_end_of_stream_only,
262 "streaming_buffer_release_on_scan": streaming_buffer_release_on_scan,
263 }
264 )
265 )
266 )
267 self.guardrailIdentifier = guardrailIdentifier
268 self.guardrailVersion = guardrailVersion
269 self.guardrail_provider = "bedrock"
270 self.chunk_budget_chars = chunk_budget_chars
271 self.contextual_grounding_from_messages = contextual_grounding_from_messages
272 self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
274 # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
275 # routes the guardrail to InvokeGuardrailChecks; absent => ApplyGuardrail.
276 self.checks: dict[str, object] | None = self._normalize_checks(checks)
277 # Per-check block thresholds; a score >= threshold blocks. None => the
278 # check is detect-only (logged, never blocks).
279 self.content_filter_threshold = content_filter_threshold
280 self.prompt_attack_threshold = prompt_attack_threshold
281 self.pii_confidence_threshold = pii_confidence_threshold
283 # store kwargs as optional_params
284 self.optional_params = kwargs
286 self.disable_exception_on_block: bool = disable_exception_on_block or False
287 """
288 If True, will not raise an exception when the guardrail is blocked.
289 """
291 # `checks` (InvokeGuardrailChecks) and `guardrailIdentifier`/`guardrailVersion`
292 # (ApplyGuardrail) are two different APIs; configuring both is ambiguous.
293 if self.checks is not None and (self.guardrailIdentifier is not None or self.guardrailVersion is not None):
294 raise ValueError(
295 "Bedrock guardrail accepts either 'guardrailIdentifier'/'guardrailVersion' (ApplyGuardrail) "
296 "or 'checks' (InvokeGuardrailChecks), not both."
297 )
299 # Set supported event hooks to include MCP hooks
300 kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks()))
302 super().__init__(**kwargs)
303 BaseAWSLLM.__init__(self)
305 # InvokeGuardrailChecks is detect-only: it never returns rewritten content,
306 # so masking has no effect in checks mode.
307 if self.checks is not None and (
308 getattr(self, "mask_request_content", False) or getattr(self, "mask_response_content", False)
309 ):
310 verbose_proxy_logger.warning(
311 "Bedrock Guardrail: mask_request_content/mask_response_content have no "
312 "effect with 'checks' (InvokeGuardrailChecks is detect-only)."
313 )
315 verbose_proxy_logger.debug(
316 "Bedrock Guardrail initialized with guardrailIdentifier: %s, guardrailVersion: %s, checks: %s",
317 self.guardrailIdentifier,
318 self.guardrailVersion,
319 list(self.checks.keys()) if self.checks else None,
320 )
322 def _set_streaming_params(self, streaming_params: BedrockGuardrailStreamingParams) -> None:
323 self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated
324 self.streaming_sampling_rate = streaming_params.streaming_sampling_rate
325 self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only
326 self.streaming_buffer_release_on_scan = streaming_params.streaming_buffer_release_on_scan
328 def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None:
329 super().update_in_memory_litellm_params(litellm_params)
330 self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra))
332 def _streams_incrementally(self) -> bool:
333 if self.mask_response_content:
334 return False
335 if not self.streaming_buffer_until_moderated:
336 return True
337 return self.streaming_buffer_release_on_scan and not self.streaming_end_of_stream_only
339 @classmethod
340 def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]:
341 return [
342 GuardrailEventHooks.pre_call,
343 GuardrailEventHooks.post_call,
344 GuardrailEventHooks.during_call,
345 GuardrailEventHooks.pre_mcp_call,
346 GuardrailEventHooks.during_mcp_call,
347 ]
349 @staticmethod
350 def _normalize_checks(checks: BedrockChecksConfigModel | Mapping[str, object] | None) -> dict[str, object] | None:
351 """Normalize the configured `checks` into a plain dict for the API body.
353 Accepts a pydantic ``BedrockChecksConfigModel`` or a raw dict; drops None /
354 unknown keys. Returns None when no usable check is configured (=> ApplyGuardrail).
355 """
356 if checks is None:
357 return None
358 raw = checks.model_dump(exclude_none=True) if isinstance(checks, BedrockChecksConfigModel) else dict(checks)
359 unknown_keys: Final = set(raw.keys()) - _BEDROCK_CHECKS_KNOWN_KEYS
360 if unknown_keys:
361 verbose_proxy_logger.warning(
362 "BedrockGuardrail: unrecognized check key(s) %s will be ignored; "
363 "recognized keys will still be used for InvokeGuardrailChecks. "
364 "Known keys: %s.",
365 sorted(unknown_keys),
366 sorted(_BEDROCK_CHECKS_KNOWN_KEYS),
367 )
368 cleaned = {key: value for key, value in raw.items() if key in _BEDROCK_CHECKS_KNOWN_KEYS and value is not None}
369 if not cleaned and raw:
370 raise ValueError(
371 f"BedrockGuardrail: 'checks' block contained only unrecognized or empty keys {sorted(raw.keys())}. "
372 f"Known keys: {sorted(_BEDROCK_CHECKS_KNOWN_KEYS)}. "
373 "Fix the guardrail config or remove the 'checks' block to use ApplyGuardrail mode."
374 )
375 return cleaned or None
377 def _create_bedrock_input_content_request(self, messages: list[AllMessageValues] | None) -> BedrockRequest:
378 """
379 Create a bedrock request for the input content - the LLM request.
380 """
381 bedrock_request: Final[BedrockRequest] = BedrockRequest(source="INPUT")
382 bedrock_request_content: Final[list[BedrockContentItem]] = []
383 if messages is None:
384 return bedrock_request
385 for message in messages:
386 blocks = self.get_content_items_for_message(message=message)
387 if blocks is None:
388 continue
389 for block in blocks:
390 # INPUT scans send plain text only. Grounding qualifiers are attached
391 # exclusively when assembling the OUTPUT request, so a caller cannot use
392 # a grounding_source/query tag to change how input-safety policies treat
393 # their content (which would be an input-guardrail bypass).
394 bedrock_request_content.append(BedrockContentItem(text=BedrockTextContent(text=block.text)))
396 bedrock_request["content"] = bedrock_request_content
397 return bedrock_request
399 def _create_bedrock_output_content_request(
400 self,
401 response: object,
402 messages: list[AllMessageValues] | None = None,
403 ) -> BedrockRequest:
404 """
405 Create a bedrock request for the output content - the LLM response.
407 Contextual grounding grades the response against the reference source and
408 the user query from the request. When the request tagged any
409 ``grounding_source``/``query`` blocks, they are emitted first and the
410 response is qualified as ``guard_content`` so Bedrock can score grounding.
411 Without such tags the payload is the legacy single response block.
412 """
413 bedrock_request: Final[BedrockRequest] = BedrockRequest(source="OUTPUT")
414 grounding_blocks: Final = self._collect_grounding_blocks(messages)
415 bedrock_request_content: Final[list[BedrockContentItem]] = [
416 self._build_content_item(block) for block in grounding_blocks
417 ]
418 has_grounding: Final = len(bedrock_request_content) > 0
419 # Append the response (the content to guard) after any grounding blocks; assign
420 # unconditionally so harvested grounding blocks survive a non-ModelResponse input.
421 bedrock_request_content.extend(self._build_response_content_items(response, has_grounding=has_grounding))
422 bedrock_request["content"] = bedrock_request_content
423 return bedrock_request
425 def _build_response_content_items(self, response: object, has_grounding: bool) -> list[BedrockContentItem]:
426 """Build content item(s) from the model response. When the request supplied
427 grounding, the response is qualified ``guard_content`` so Bedrock can score it.
428 """
429 items: Final[list[BedrockContentItem]] = []
430 if not isinstance(response, litellm.ModelResponse):
431 return items
432 for choice in response.choices:
433 if (
434 isinstance(choice, litellm.Choices)
435 and isinstance(choice.message.content, str)
436 and choice.message.content
437 ):
438 block = QualifiedTextBlock(
439 text=choice.message.content,
440 qualifier="guard_content" if has_grounding else None,
441 )
442 items.append(self._build_content_item(block))
443 return items
445 def convert_to_bedrock_format(
446 self,
447 source: Literal["INPUT", "OUTPUT"],
448 messages: list[AllMessageValues] | None = None,
449 response: object | None = None,
450 ) -> BedrockRequest:
451 """
452 Convert the litellm messages/response to the bedrock request format.
454 If source is "INPUT", then messages is required.
455 If source is "OUTPUT", then response is required.
457 Returns:
458 BedrockRequest: The bedrock request object.
459 """
460 bedrock_request: BedrockRequest = BedrockRequest(source=source)
461 if source == "INPUT":
462 bedrock_request = self._create_bedrock_input_content_request(messages=messages)
463 elif source == "OUTPUT":
464 bedrock_request = self._create_bedrock_output_content_request(response=response, messages=messages)
465 return bedrock_request
467 def get_content_items_for_message(self, message: AllMessageValues) -> list[QualifiedTextBlock] | None:
468 """
469 Flatten a message into text blocks, preserving any contextual-grounding
470 qualifier carried by the content-block ``type`` (grounding_source / query).
471 Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to
472 derive grounding qualifiers from it.
473 """
474 content: Final = message.get("content")
475 if content is None:
476 return None
477 blocks: Final[list[QualifiedTextBlock]] = []
478 if isinstance(content, str):
479 blocks.append(QualifiedTextBlock(text=content, qualifier=None))
480 elif isinstance(content, list):
481 for item in content:
482 if isinstance(item, dict) and "text" in item:
483 qualifier = _CONTENT_TYPE_TO_QUALIFIER.get(item.get("type", ""))
484 blocks.append(QualifiedTextBlock(text=item["text"], qualifier=qualifier))
485 elif isinstance(item, str):
486 blocks.append(QualifiedTextBlock(text=item, qualifier=None))
487 return blocks
489 def _build_content_item(self, block: QualifiedTextBlock) -> BedrockContentItem:
490 """Build a Bedrock content item, attaching qualifiers only when present."""
491 text_content: Final = BedrockTextContent(text=block.text)
492 if block.qualifier is not None:
493 text_content["qualifiers"] = [block.qualifier]
494 return BedrockContentItem(text=text_content)
496 def _collect_grounding_blocks(self, messages: list[AllMessageValues] | None) -> list[QualifiedTextBlock]:
497 """Harvest grounding_source/query blocks from the request for an OUTPUT scan.
499 ``grounding_source`` is honored only from app-authored roles (system /
500 developer). A grounding_source tag on a ``user``, ``tool`` or ``function``
501 message is ignored, so neither a forwarded end-user message nor a tool/function
502 result carrying externally-influenced content can supply fake evidence for the
503 contextual-grounding check to grade the response against. ``query`` is accepted
504 from any role (it is the user's question).
506 With ``contextual_grounding_from_messages`` on, a request with no tagged blocks
507 falls back to the plain messages: system / developer text is the grounding
508 source and the latest user message is the query.
509 """
510 grounding: Final[list[QualifiedTextBlock]] = []
511 for message in messages or []:
512 role = message.get("role")
513 for block in self.get_content_items_for_message(message=message) or []:
514 if (
515 block.qualifier == "query"
516 or block.qualifier == "grounding_source"
517 and role in _GROUNDING_SOURCE_TRUSTED_ROLES
518 ):
519 grounding.append(block)
520 if grounding or not self.contextual_grounding_from_messages:
521 return grounding
522 return self._derive_grounding_blocks_from_plain_messages(messages)
524 def _derive_grounding_blocks_from_plain_messages(
525 self, messages: list[AllMessageValues] | None
526 ) -> list[QualifiedTextBlock]:
527 if not messages:
528 return []
529 latest_user_index: Final = self._find_latest_message_index(messages, target_role="user")
530 if latest_user_index is None:
531 return []
532 sources: Final = tuple(
533 QualifiedTextBlock(text=block.text, qualifier="grounding_source")
534 for message in messages
535 if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES
536 for block in self.get_content_items_for_message(message=message) or []
537 if block.text
538 )
539 queries: Final = tuple(
540 QualifiedTextBlock(text=block.text, qualifier="query")
541 for block in self.get_content_items_for_message(message=messages[latest_user_index]) or []
542 if block.text
543 )
544 if not sources or not queries:
545 return []
546 return [*sources, *queries]
548 def supports_scan_only_tool_results(self) -> bool:
549 return self.experimental_use_latest_role_message_only is not True
551 def _prepare_guardrail_messages_for_role(
552 self,
553 messages: list[AllMessageValues] | None,
554 ) -> GuardrailMessageFilterResult:
555 """Return payload + merge metadata for the latest user message."""
556 # NOTE: This logic probably belongs in CustomGuardrail once other guardrails adopt the feature.
558 if messages is None:
559 return GuardrailMessageFilterResult(None, None, None)
561 if self.experimental_use_latest_role_message_only is not True:
562 return GuardrailMessageFilterResult(messages, None, None)
564 latest_index: Final = self._find_latest_message_index(messages, target_role="user")
565 if latest_index is None:
566 return GuardrailMessageFilterResult(None, None, None)
568 original_messages: Final = list(messages)
569 payload_messages: Final = [messages[latest_index]]
570 return GuardrailMessageFilterResult(
571 payload_messages=payload_messages,
572 original_messages=original_messages,
573 target_indices=[latest_index],
574 )
576 def _find_latest_message_index(self, messages: list[AllMessageValues], target_role: str) -> int | None:
577 for index in range(len(messages) - 1, -1, -1):
578 if messages[index].get("role", None) == target_role:
579 return index
580 return None
582 @staticmethod
583 def _count_message_texts(message: AllMessageValues) -> int:
584 """Count the text segments the guardrail translation layer extracts from a message."""
585 content: Final = message.get("content")
586 if isinstance(content, str):
587 return 1
588 if isinstance(content, list):
589 return sum(1 for item in content if isinstance(item, dict) and item.get("text") is not None)
590 return 0
592 def _locate_message_texts_slice(
593 self,
594 structured_messages: list[AllMessageValues],
595 target_index: int,
596 texts: list[str],
597 ) -> tuple[int, int] | None:
598 """
599 Map one message's text segments to their (offset, length) slice in the
600 flat `texts` list built by the guardrail translation handler.
602 Returns None when the reconstruction does not line up with `texts`
603 (the caller must then avoid positional write-back).
604 """
605 offset = 0
606 total = 0
607 target_count = 0
608 for index, message in enumerate(structured_messages):
609 count = self._count_message_texts(message)
610 if index < target_index:
611 offset += count
612 elif index == target_index:
613 target_count = count
614 total += count
615 if total != len(texts) or target_count == 0:
616 return None
617 return offset, target_count
619 def _select_messages_for_apply_guardrail(
620 self,
621 texts: list[str],
622 inputs: "GenericGuardrailAPIInputs",
623 request_data: dict,
624 input_type: Literal["request", "response"],
625 ) -> ApplyGuardrailMessageSelection:
626 """
627 Decide which messages an apply_guardrail scan should cover.
629 With ``experimental_use_latest_role_message_only`` enabled, request
630 scans must select by the ORIGINAL message roles. The flat `texts` list
631 has no role information, and wrapping it in role="user" mock messages
632 makes the latest-user filter degenerate to "latest text of any role",
633 leaking tool/assistant content to the INPUT scan
634 (https://github.com/BerriAI/litellm/issues/23476).
635 """
636 mock_messages: list[AllMessageValues] = [ChatCompletionUserMessage(role="user", content=text) for text in texts]
638 if self.experimental_use_latest_role_message_only is not True:
639 return ApplyGuardrailMessageSelection(
640 filtered_messages=mock_messages,
641 scanned_slice=None,
642 scanned_role_subset=False,
643 )
645 # Prefer inputs["structured_messages"]: it is built alongside `texts` by
646 # the translation handler and stays aligned with it even when
647 # skip_system_message_in_guardrail / skip_tool_message_in_guardrail drop
648 # messages. The fallback to request_data["messages"] is the *unfiltered*
649 # list, so it only lines up with `texts` when no skip flags are active.
650 # When a skip flag is set and we land on this fallback (direct
651 # apply_guardrail callers with no structured_messages),
652 # _locate_message_texts_slice will detect the length mismatch and return
653 # None, and the write-back guard below safely skips masking rather than
654 # corrupting positions.
655 structured_messages: Final = cast(
656 list[AllMessageValues] | None,
657 inputs.get("structured_messages") or request_data.get("messages"),
658 )
659 if input_type != "request" or not structured_messages:
660 # No role information available (e.g. raw-text callers like
661 # /guardrails/apply_guardrail) — keep the legacy behavior of
662 # scanning the latest text only.
663 filter_result: Final = self._prepare_guardrail_messages_for_role(messages=mock_messages)
664 return ApplyGuardrailMessageSelection(
665 filtered_messages=filter_result.payload_messages or mock_messages,
666 scanned_slice=None,
667 scanned_role_subset=False,
668 )
670 latest_user_index: Final = self._find_latest_message_index(structured_messages, target_role="user")
671 if latest_user_index is None:
672 if effective_scan_only_tool_results_for_guardrail(self):
673 verbose_proxy_logger.warning(
674 "Bedrock Guardrail: experimental_use_latest_role_message_only scans only the latest "
675 "user message, so scan_only_tool_results leaves nothing to scan for this request"
676 )
677 verbose_proxy_logger.debug("Bedrock Guardrail: no user-role message in request, skipping INPUT scan")
678 return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True)
680 selected_message: Final = structured_messages[latest_user_index]
681 if self._count_message_texts(selected_message) == 0:
682 verbose_proxy_logger.debug(
683 "Bedrock Guardrail: latest user message has no text content, skipping INPUT scan"
684 )
685 return ApplyGuardrailMessageSelection(None, None, True, skip_scan=True)
687 return ApplyGuardrailMessageSelection(
688 filtered_messages=[selected_message],
689 scanned_slice=self._locate_message_texts_slice(
690 structured_messages=structured_messages,
691 target_index=latest_user_index,
692 texts=texts,
693 ),
694 scanned_role_subset=True,
695 )
697 def _merge_masked_texts(
698 self,
699 masked_texts: list,
700 texts: list,
701 scanned_slice: tuple[int, int] | None,
702 scanned_role_subset: bool,
703 ) -> list:
704 """
705 Reconcile the guardrail's masked output with the flat `texts` list.
707 - No masked output: keep the originals (guardrail allowed content as-is).
708 - A slice was scanned: write masked content back to those positions only,
709 keeping the list aligned with the caller's message↔text mappings.
710 - A role-selected subset was scanned but could not be mapped back to
711 flat-text positions (scanned_slice is None): keep the originals rather
712 than misapply masked content to the wrong message. Guarding on
713 scanned_slice rather than a length comparison also covers the case
714 where the masked subset happens to match len(texts) (e.g. both length
715 1).
716 - Otherwise (whole list scanned): use the masked output as-is.
717 """
718 if not masked_texts:
719 return texts
720 if scanned_slice is not None:
721 offset, length = scanned_slice
722 merged_texts: Final = list(texts)
723 for masked_index, masked_text in enumerate(masked_texts[:length]):
724 merged_texts[offset + masked_index] = masked_text
725 return merged_texts
726 if scanned_role_subset:
727 verbose_proxy_logger.warning(
728 "Bedrock Guardrail: could not align masked texts with request texts, skipping masking write-back"
729 )
730 return texts
731 return masked_texts
733 def _merge_filtered_messages(
734 self,
735 original_messages: list[AllMessageValues] | None,
736 updated_target_messages: list[AllMessageValues],
737 target_indices: list[int] | None,
738 ) -> list[AllMessageValues]:
739 if not target_indices:
740 return updated_target_messages
742 if not original_messages:
743 original_messages = []
745 merged_messages = list(original_messages)
746 if not merged_messages:
747 merged_messages = list(updated_target_messages)
748 for replacement_index, updated_message in zip(target_indices, updated_target_messages):
749 if replacement_index < len(merged_messages):
750 merged_messages[replacement_index] = updated_message
752 return merged_messages
754 # NOTE: Consider moving these helpers to CustomGuardrail when the filtering
755 # logic becomes shared across providers.
757 #### CALL HOOKS - proxy only ####
758 def _load_credentials(self, bearer_token: str | None = None):
759 try:
760 from botocore.credentials import Credentials
761 except ImportError:
762 raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
763 ## CREDENTIALS ##
764 aws_secret_access_key: Final = self.optional_params.get("aws_secret_access_key", None)
765 aws_access_key_id: Final = self.optional_params.get("aws_access_key_id", None)
766 aws_session_token: Final = self.optional_params.get("aws_session_token", None)
767 aws_region_name = self.optional_params.get("aws_region_name", None)
768 aws_role_name: Final = self.optional_params.get("aws_role_name", None)
769 aws_session_name: Final = self.optional_params.get("aws_session_name", None)
770 aws_profile_name: Final = self.optional_params.get("aws_profile_name", None)
771 aws_web_identity_token: Final = self.optional_params.get("aws_web_identity_token", None)
772 aws_sts_endpoint: Final = self.optional_params.get("aws_sts_endpoint", None)
773 aws_external_id: Final = self.optional_params.get("aws_external_id", None)
775 ### SET REGION NAME ###
776 aws_region_name = self.get_aws_region_name_for_non_llm_api_calls(
777 aws_region_name=aws_region_name,
778 )
780 credentials: Final[Credentials | None] = (
781 None
782 if bearer_token is not None
783 else self.get_credentials(
784 aws_access_key_id=aws_access_key_id,
785 aws_secret_access_key=aws_secret_access_key,
786 aws_session_token=aws_session_token,
787 aws_region_name=aws_region_name,
788 aws_session_name=aws_session_name,
789 aws_profile_name=aws_profile_name,
790 aws_role_name=aws_role_name,
791 aws_web_identity_token=aws_web_identity_token,
792 aws_sts_endpoint=aws_sts_endpoint,
793 aws_external_id=aws_external_id,
794 )
795 )
796 return credentials, aws_region_name
798 def _prepare_request(
799 self,
800 credentials,
801 data: dict,
802 optional_params: dict,
803 aws_region_name: str,
804 api_key: str | None = None,
805 extra_headers: dict | None = None,
806 request_path: str | None = None,
807 ):
808 headers = {"Content-Type": "application/json"}
809 if extra_headers is not None:
810 headers = {"Content-Type": "application/json", **extra_headers}
812 aws_bedrock_runtime_endpoint: Final = self.optional_params.get("aws_bedrock_runtime_endpoint", None)
813 _, proxy_endpoint_url = self.get_runtime_endpoint(
814 api_base=None,
815 aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
816 aws_region_name=aws_region_name,
817 )
818 # Default to the ApplyGuardrail resource path. Callers pass an explicit
819 # request_path for the resource-less InvokeGuardrailChecks endpoint (where
820 # guardrailIdentifier/guardrailVersion are None and must not be interpolated).
821 if request_path is None:
822 request_path = f"/guardrail/{self.guardrailIdentifier}/version/{self.guardrailVersion}/apply"
823 proxy_endpoint_url = f"{proxy_endpoint_url}{request_path}"
824 encoded_data: Final = json.dumps(data).encode("utf-8")
826 aws_bearer_token: Final = bedrock_bearer_token(api_key)
828 if aws_bearer_token is not None:
829 try:
830 from botocore.awsrequest import AWSRequest
831 except ImportError:
832 raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
833 headers["Authorization"] = f"Bearer {aws_bearer_token}"
834 request = AWSRequest(
835 method="POST",
836 url=proxy_endpoint_url,
837 data=encoded_data,
838 headers=headers,
839 )
840 else:
841 try:
842 from botocore.auth import SigV4Auth
843 from botocore.awsrequest import AWSRequest
844 except ImportError:
845 raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
847 sigv4: Final = SigV4Auth(credentials, "bedrock", aws_region_name)
848 request = AWSRequest(
849 method="POST",
850 url=proxy_endpoint_url,
851 data=encoded_data,
852 headers=headers,
853 )
854 sigv4.add_auth(request)
855 if (
856 extra_headers is not None and "Authorization" in extra_headers
857 ): # prevent sigv4 from overwriting the auth header
858 request.headers["Authorization"] = extra_headers["Authorization"]
859 prepped_request: Final = request.prepare()
861 return prepped_request
863 async def make_bedrock_api_request(
864 self,
865 source: Literal["INPUT", "OUTPUT"],
866 messages: list[AllMessageValues] | None = None,
867 response: litellm.ModelResponse | None = None,
868 request_data: dict | None = None,
869 logging_event_type: GuardrailEventHooks | None = None,
870 ) -> BedrockGuardrailResponse:
871 """Dispatch to the configured Bedrock guardrail API.
873 ``checks`` selects the resource-less, detect-only InvokeGuardrailChecks API;
874 otherwise the ApplyGuardrail API is used. Both return a ``BedrockGuardrailResponse``
875 (the checks path returns an empty one on a pass, which downstream masking treats
876 as a no-op) and raise on a blocked request.
877 """
878 if self.checks is not None:
879 return await self._make_invoke_guardrail_checks_request(
880 source=source,
881 messages=messages,
882 response=response,
883 request_data=request_data,
884 logging_event_type=logging_event_type,
885 )
886 return await self._make_apply_guardrail_request(
887 source=source,
888 messages=messages,
889 response=response,
890 request_data=request_data,
891 logging_event_type=logging_event_type,
892 )
894 async def _make_apply_guardrail_request(
895 self,
896 source: Literal["INPUT", "OUTPUT"],
897 messages: list[AllMessageValues] | None = None,
898 response: litellm.ModelResponse | None = None,
899 request_data: dict | None = None,
900 logging_event_type: GuardrailEventHooks | None = None,
901 ) -> BedrockGuardrailResponse:
902 """Scan `messages`/`response` with ApplyGuardrail, chunking if it is too large.
904 Content is bin-packed into budget-sized batches and each batch posted
905 sequentially, every batch independently falling back to bisection if AWS
906 rejects it. The per-batch responses are merged so callers cannot tell whether
907 chunking happened.
909 Content using contextual grounding opts out of chunking entirely: grounding is
910 scored holistically against the whole reference source, so bisecting it would
911 fragment that evaluation and yield misleading scores. Such a request keeps the
912 old behavior of surfacing a too-large error rather than being split.
914 `logging_event_type` drives what UI and spend logs report. It is distinct from
915 Bedrock's `source`, which is INPUT vs OUTPUT for the API body and must not be
916 confused with the proxy hook (pre_call / during_call / post_call); when omitted,
917 the legacy source-derived mapping is kept for backward compatibility.
919 A guardrail *block* is logged where it happens, in
920 `_post_apply_guardrail_content`, because chunking stops immediately and there is
921 no later merged response to log instead. Everything else that fails out of the
922 chunking flow (an unrecoverable too-large error, a non-size validation error,
923 exhausted throttle retries) is a genuine end-to-end failure of this one logical
924 guardrail call and is logged exactly once here.
925 """
926 start_time: Final = datetime.now(timezone.utc)
927 bedrock_request_data: Final[dict] = dict(
928 self.convert_to_bedrock_format(source=source, messages=messages, response=response)
929 )
930 api_key: str | None = None
931 if request_data:
932 dynamic_request_body_params = self.get_guardrail_dynamic_request_body_params(request_data=request_data)
933 bedrock_request_data.update(
934 {
935 key: value
936 for key, value in dynamic_request_body_params.items()
937 if key not in _BEDROCK_DYNAMIC_BODY_DENYLIST
938 }
939 )
940 if request_data.get("api_key") is not None:
941 api_key = request_data["api_key"]
943 event_type: Final = (
944 logging_event_type
945 if logging_event_type is not None
946 else (GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call)
947 )
949 content: Final[tuple[BedrockContentItem, ...]] = tuple(bedrock_request_data.get("content") or ())
950 if not content:
951 # ApplyGuardrail rejects an empty content list with a 400, so a turn this extractor
952 # found no text in is skipped rather than turned into a failed request
953 verbose_proxy_logger.debug(
954 "Bedrock Guardrail %s: no %s content to scan, skipping ApplyGuardrail",
955 self.guardrail_name,
956 source,
957 )
958 return BedrockGuardrailResponse()
959 credentials, aws_region_name = await run_aws_signing(
960 self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
961 )
962 allow_chunking: Final = not self._content_uses_contextual_grounding(content)
964 completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
965 try:
966 responses: Final = await self._apply_guardrail_content_with_chunking(
967 content=content,
968 base_request_data=bedrock_request_data,
969 credentials=credentials,
970 aws_region_name=aws_region_name,
971 api_key=api_key,
972 request_data=request_data,
973 event_type=event_type,
974 start_time=start_time,
975 allow_chunking=allow_chunking,
976 completed_chunk_usages=completed_chunk_usages,
977 )
978 except HTTPException as exc:
979 if not isinstance(exc.detail, dict):
980 self._log_apply_guardrail_failure(
981 detail=exc.detail,
982 request_data=request_data,
983 event_type=event_type,
984 start_time=start_time,
985 aws_region_name=aws_region_name,
986 completed_chunk_usages=completed_chunk_usages,
987 )
988 raise
989 merged_response: Final = self._merge_bedrock_guardrail_responses(responses)
990 self._log_apply_guardrail_success(
991 merged_response=merged_response,
992 request_data=request_data,
993 event_type=event_type,
994 start_time=start_time,
995 aws_region_name=aws_region_name,
996 )
997 return merged_response
999 async def _apply_guardrail_content_with_chunking(
1000 self,
1001 content: Sequence[BedrockContentItem],
1002 base_request_data: Mapping[str, object],
1003 credentials: "Credentials | None",
1004 aws_region_name: str,
1005 api_key: str | None,
1006 request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
1007 event_type: GuardrailEventHooks,
1008 start_time: "datetime",
1009 allow_chunking: bool,
1010 completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: billed-chunk usage accumulator
1011 ) -> tuple[BedrockContentChunkResult, ...]:
1012 """Post `content` to ApplyGuardrail, chunking only if AWS rejects it as too large.
1014 Tries `content` as a single call first. AWS's per-request "maximum input
1015 size in text units" quota is account/region/policy-dependent and cannot be
1016 predicted ahead of time, so it is only ever discovered reactively: on an
1017 error whose message indicates the input was too large (a ThrottlingException
1018 in practice, a ValidationException per the docs -- see
1019 ``_is_input_too_large_error``), the content is re-sent in smaller pieces.
1021 Probing with the whole payload first is what keeps a request AWS would have
1022 accepted at exactly one call. Packing into fixed batches up front instead
1023 would split conversations AWS was happy to take whole, multiplying billed
1024 calls and guardrail latency on traffic that never had a size problem, and
1025 no fixed budget can avoid that because the real cap is unknown here.
1027 Once a rejection proves the payload is over the cap, a multi-item payload is
1028 re-sent as ``chunk_budget_chars``-sized batches rather than bisected: that
1029 reaches a working size in one step instead of paying an O(log n) ladder of
1030 rejected calls. Bisection remains the fallback for anything bin-packing
1031 cannot make smaller, which is what makes the recursion terminate: a batch
1032 already inside the budget packs back to itself, so it falls through to the
1033 split below. A single oversized
1034 content item (one very long message) is split by its own text instead of
1035 by list length, since a list of length 1 has no items left to bisect --
1036 the resulting fragments all carry a ``fragment_group_size`` so the merge
1037 step can recombine them into the one content item they came from, rather
1038 than treating each fragment as its own item when reconstructing positions
1039 for masking. That count covers however many fragments the item ended up
1040 split into, not just two, since it can be bisected repeatedly: the
1041 outermost single-item split stamps the total leaf count on every leaf
1042 below it, overwriting any smaller count an inner split had set. A real
1043 guardrail block on any (sub-)chunk raises immediately
1044 -- callers must not lose that signal by continuing to post the remaining
1045 chunks.
1046 """
1047 try:
1048 response: Final = await self._post_apply_guardrail_content_with_retry(
1049 content=content,
1050 base_request_data=base_request_data,
1051 credentials=credentials,
1052 aws_region_name=aws_region_name,
1053 api_key=api_key,
1054 request_data=request_data,
1055 event_type=event_type,
1056 start_time=start_time,
1057 completed_chunk_usages=completed_chunk_usages,
1058 )
1059 return (
1060 BedrockContentChunkResult(
1061 response=response,
1062 content=tuple(content),
1063 fragment_group_size=1,
1064 ),
1065 )
1066 except HTTPException as exc:
1067 if allow_chunking and self._is_input_too_large_error(exc.detail):
1068 batches: Final = self._bin_pack_bedrock_content(content, budget=self.chunk_budget_chars)
1069 if len(batches) > 1:
1070 verbose_proxy_logger.warning(
1071 "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; "
1072 "re-sending as %d batches of at most %d characters",
1073 len(content),
1074 len(batches),
1075 self.chunk_budget_chars,
1076 )
1077 batch_results: Final = [ # mutable-ok: await needs a list comprehension; frozen to a tuple below
1078 await self._apply_guardrail_content_with_chunking(
1079 content=batch,
1080 base_request_data=base_request_data,
1081 credentials=credentials,
1082 aws_region_name=aws_region_name,
1083 api_key=api_key,
1084 request_data=request_data,
1085 event_type=event_type,
1086 start_time=start_time,
1087 allow_chunking=allow_chunking,
1088 completed_chunk_usages=completed_chunk_usages,
1089 )
1090 for batch in batches
1091 ]
1092 return tuple(result for results in batch_results for result in results)
1093 split_content: Final = self._split_bedrock_content(content)
1094 if split_content is None:
1095 raise
1096 first_half, second_half = split_content
1097 is_single_item_text_split: Final = len(content) == 1
1098 verbose_proxy_logger.warning(
1099 "Bedrock Guardrail: ApplyGuardrail rejected %d content item(s) as too large; "
1100 "splitting into %d + %d and retrying each",
1101 len(content),
1102 len(first_half),
1103 len(second_half),
1104 )
1105 first_results: Final = await self._apply_guardrail_content_with_chunking(
1106 content=first_half,
1107 base_request_data=base_request_data,
1108 credentials=credentials,
1109 aws_region_name=aws_region_name,
1110 api_key=api_key,
1111 request_data=request_data,
1112 event_type=event_type,
1113 start_time=start_time,
1114 allow_chunking=allow_chunking,
1115 completed_chunk_usages=completed_chunk_usages,
1116 )
1117 second_results: Final = await self._apply_guardrail_content_with_chunking(
1118 content=second_half,
1119 base_request_data=base_request_data,
1120 credentials=credentials,
1121 aws_region_name=aws_region_name,
1122 api_key=api_key,
1123 request_data=request_data,
1124 event_type=event_type,
1125 start_time=start_time,
1126 allow_chunking=allow_chunking,
1127 completed_chunk_usages=completed_chunk_usages,
1128 )
1129 combined_results: Final = tuple(first_results) + tuple(second_results)
1130 if is_single_item_text_split:
1131 return tuple(
1132 result._replace(fragment_group_size=len(combined_results)) for result in combined_results
1133 )
1134 return combined_results
1135 raise
1137 async def _post_apply_guardrail_content_with_retry(
1138 self,
1139 content: Sequence[BedrockContentItem],
1140 base_request_data: Mapping[str, object],
1141 credentials: "Credentials | None",
1142 aws_region_name: str,
1143 api_key: str | None,
1144 request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
1145 event_type: GuardrailEventHooks,
1146 start_time: "datetime",
1147 completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: passed through to the single-call layer
1148 ) -> BedrockGuardrailResponse:
1149 """Post one ApplyGuardrail call for `content`, retrying with exponential
1150 backoff on AWS ThrottlingException (HTTP 429).
1152 Chunking already trades one oversized call for several smaller ones, so
1153 retries here are capped low -- they must not multiply per-request latency
1154 by an order of magnitude when the account's per-second text-unit quota is
1155 the binding constraint rather than the per-request size quota.
1157 A too-large rejection is deliberately excluded from the retry. AWS reports
1158 it as a ThrottlingException (429), not only as a ValidationException, but
1159 unlike a genuine throttle it is not transient: re-posting the same
1160 oversized content can never succeed. Retrying it would burn every backoff
1161 sleep and every (billed) attempt before the caller's bisection gets a
1162 chance to split the content, at every level of the recursion.
1163 """
1164 for attempt in range(_BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES + 1):
1165 try:
1166 return await self._post_apply_guardrail_content(
1167 content=content,
1168 base_request_data=base_request_data,
1169 credentials=credentials,
1170 aws_region_name=aws_region_name,
1171 api_key=api_key,
1172 request_data=request_data,
1173 event_type=event_type,
1174 start_time=start_time,
1175 completed_chunk_usages=completed_chunk_usages,
1176 )
1177 except HTTPException as exc:
1178 if (
1179 exc.status_code != 429
1180 or self._is_input_too_large_error(exc.detail)
1181 or attempt >= _BEDROCK_APPLY_GUARDRAIL_MAX_THROTTLE_RETRIES
1182 ):
1183 raise
1184 await asyncio.sleep(_BEDROCK_APPLY_GUARDRAIL_BASE_BACKOFF_SECONDS * (2**attempt))
1185 raise HTTPException(status_code=500, detail="Bedrock guardrail throttle retries exhausted")
1187 async def _post_apply_guardrail_content(
1188 self,
1189 content: Sequence[BedrockContentItem],
1190 base_request_data: Mapping[str, object],
1191 credentials: "Credentials | None",
1192 aws_region_name: str,
1193 api_key: str | None,
1194 request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
1195 event_type: GuardrailEventHooks,
1196 start_time: "datetime",
1197 completed_chunk_usages: list[BedrockGuardrailUsage], # mutable-ok: billed-chunk usage accumulator
1198 ) -> BedrockGuardrailResponse:
1199 """Make exactly one signed ApplyGuardrail HTTP call for `content` and
1200 parse the result. Raises HTTPException on a guardrail block or any
1201 non-200 response (including 429, handled by the retry wrapper above).
1203 AWS also reports some failures inside a 200 body, tagging ``Output.__type``
1204 with an Exception marker. Those deliberately do NOT raise: the request proceeds,
1205 matching the behaviour of this code before chunking existed. The marker survives
1206 the merge, so the one consolidated log entry still records
1207 ``guardrail_failed_to_respond`` rather than a success. Making that path fail
1208 closed is a separate change, tracked apart from this PR, and belongs behind the
1209 existing ``unreachable_fallback`` setting rather than a hardcoded status.
1211 A block is logged here rather than by the caller: it ends the whole chunking
1212 flow immediately, with no further chunks attempted, so there is no later
1213 merged response for the caller to log instead. The logged usage still spans
1214 the whole logical request: chunks that passed before the block appended what
1215 AWS billed them to ``completed_chunk_usages``, and the attempt log sums those
1216 with the blocking call's own usage.
1217 """
1218 bedrock_request_data: Final = { # mutable-ok: outbound JSON request body
1219 **base_request_data,
1220 "content": content,
1221 }
1222 prepared_request: Final = await run_aws_signing(
1223 self._prepare_request,
1224 credentials=credentials,
1225 data=bedrock_request_data,
1226 optional_params=self.optional_params,
1227 aws_region_name=aws_region_name,
1228 api_key=api_key,
1229 )
1230 headers_dict: Final = dict(prepared_request.headers) # mutable-ok: the masking helper requires a dict
1231 verbose_proxy_logger.debug(
1232 "Bedrock AI request body: %s, url %s, headers: %s",
1233 bedrock_request_data,
1234 prepared_request.url,
1235 _get_masked_values(headers_dict),
1236 )
1238 httpx_response: Final = await self._sign_and_post(
1239 prepared_request=prepared_request,
1240 request_data=request_data,
1241 event_type=event_type,
1242 start_time=start_time,
1243 log_transport_failure=False,
1244 )
1246 if httpx_response.status_code == 200:
1247 _json_response: Final = httpx_response.json()
1248 # check if the response was flagged
1249 verbose_proxy_logger.debug(
1250 "Bedrock AI response : %s",
1251 redact_nested_match_and_regex_keys(_json_response),
1252 )
1253 bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
1254 if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
1255 self._log_apply_guardrail_attempt(
1256 httpx_response=httpx_response,
1257 json_response=_json_response,
1258 request_data=request_data,
1259 event_type=event_type,
1260 start_time=start_time,
1261 aws_region_name=aws_region_name,
1262 completed_chunk_usages=completed_chunk_usages,
1263 )
1264 raise self._get_http_exception_for_blocked_guardrail(
1265 bedrock_guardrail_response, request_data=request_data
1266 )
1267 response_usage: Final = bedrock_guardrail_response.get("usage")
1268 if isinstance(response_usage, dict):
1269 completed_chunk_usages.append(response_usage)
1270 return bedrock_guardrail_response
1272 status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
1273 verbose_proxy_logger.error(
1274 "Bedrock AI: error in response. Status code: %s, response: %s",
1275 httpx_response.status_code,
1276 httpx_response.text,
1277 )
1278 raise HTTPException(status_code=status_code, detail=detail_message)
1280 def _log_apply_guardrail_attempt(
1281 self,
1282 httpx_response: httpx.Response,
1283 json_response: dict, # mutable-ok: raw AWS JSON payload
1284 request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
1285 event_type: GuardrailEventHooks,
1286 start_time: "datetime",
1287 aws_region_name: str | None,
1288 completed_chunk_usages: Sequence[BedrockGuardrailUsage],
1289 ) -> None:
1290 """Log the blocking ApplyGuardrail attempt, which ends the whole chunking
1291 flow immediately. Its status derives from its own response, but its usage
1292 (and so its cost) spans every billed call of the logical request: the
1293 chunks that passed before the block plus the blocking call itself."""
1294 blocking_usage: Final = json_response.get("usage")
1295 billed_usages: Final[tuple[BedrockGuardrailUsage, ...]] = tuple(completed_chunk_usages) + (
1296 (blocking_usage,) if isinstance(blocking_usage, dict) else ()
1297 )
1298 logged_json_response: Final = (
1299 { # mutable-ok: raw AWS JSON payload carrying the total billed usage
1300 **json_response,
1301 "usage": self._sum_usage_counters(billed_usages),
1302 }
1303 if completed_chunk_usages
1304 else json_response
1305 )
1306 tracing_detail: Final = self._build_tracing_detail(
1307 BedrockGuardrailResponse(**logged_json_response), aws_region_name=aws_region_name
1308 )
1309 self.add_standard_logging_guardrail_information_to_request_data(
1310 guardrail_provider=self.guardrail_provider,
1311 guardrail_json_response=logged_json_response,
1312 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1313 guardrail_status=self._get_bedrock_guardrail_response_status(response=httpx_response),
1314 start_time=start_time.timestamp(),
1315 end_time=datetime.now(timezone.utc).timestamp(),
1316 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1317 event_type=event_type,
1318 tracing_detail=tracing_detail or None,
1319 )
1321 def _log_apply_guardrail_success(
1322 self,
1323 merged_response: BedrockGuardrailResponse,
1324 request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
1325 event_type: GuardrailEventHooks,
1326 start_time: "datetime",
1327 aws_region_name: str | None,
1328 ) -> None:
1329 """Log one logical ApplyGuardrail call -- possibly several chunk calls
1330 under the hood -- using its final merged response, so a chunked
1331 request produces exactly one telemetry entry, the same as an
1332 unchunked one would.
1334 AWS can report a failure inside an HTTP 200 body by tagging
1335 ``Output.__type`` with an exception marker. That marker survives the merge,
1336 so the status is derived from the merged response rather than assumed to be
1337 a success, which is what the pre-chunking code reported for that shape."""
1338 tracing_detail: Final = self._build_tracing_detail(merged_response, aws_region_name=aws_region_name)
1339 self.add_standard_logging_guardrail_information_to_request_data(
1340 guardrail_provider=self.guardrail_provider,
1341 guardrail_json_response=dict(merged_response), # mutable-ok: logging helper requires a dict
1342 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1343 guardrail_status=(
1344 "guardrail_failed_to_respond"
1345 if "Exception" in str((merged_response.get("Output") or {}).get("__type", ""))
1346 else "success"
1347 ),
1348 start_time=start_time.timestamp(),
1349 end_time=datetime.now(timezone.utc).timestamp(),
1350 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1351 event_type=event_type,
1352 tracing_detail=tracing_detail or None,
1353 )
1355 def _log_apply_guardrail_failure(
1356 self,
1357 detail: object,
1358 request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
1359 event_type: GuardrailEventHooks,
1360 start_time: "datetime",
1361 aws_region_name: str | None,
1362 completed_chunk_usages: Sequence[BedrockGuardrailUsage],
1363 ) -> None:
1364 """Log one logical ApplyGuardrail call that failed end-to-end (an
1365 unrecoverable too-large error, a non-size validation error, or
1366 exhausted throttle retries) as a single failure, rather than logging
1367 every failed attempt chunking made along the way. Chunk calls AWS
1368 billed before the failure still carry their usage and cost."""
1369 billed_usage: Final = self._sum_usage_counters(completed_chunk_usages) if completed_chunk_usages else None
1370 error_payload: Final = {"error": str(detail)} # mutable-ok: logging helper requires a dict
1371 json_response: Final = (
1372 {**error_payload, "usage": billed_usage} # mutable-ok: logging helper requires a dict
1373 if billed_usage is not None
1374 else error_payload
1375 )
1376 tracing_detail: Final = (
1377 self._build_tracing_detail(BedrockGuardrailResponse(usage=billed_usage), aws_region_name=aws_region_name)
1378 if billed_usage is not None
1379 else None
1380 )
1381 self.add_standard_logging_guardrail_information_to_request_data(
1382 guardrail_provider=self.guardrail_provider,
1383 guardrail_json_response=json_response,
1384 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1385 guardrail_status="guardrail_failed_to_respond",
1386 start_time=start_time.timestamp(),
1387 end_time=datetime.now(timezone.utc).timestamp(),
1388 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1389 event_type=event_type,
1390 tracing_detail=tracing_detail or None,
1391 )
1393 @staticmethod
1394 def _content_uses_contextual_grounding(content: Sequence[BedrockContentItem]) -> bool:
1395 """True if any content item carries a contextual-grounding qualifier
1396 (``grounding_source``, ``query``, or the ``guard_content`` the response
1397 itself is tagged with once grounding is present)."""
1398 for item in content:
1399 if (item.get("text") or {}).get("qualifiers"): # mutable-ok: read-only empty fallback
1400 return True
1401 return False
1403 @staticmethod
1404 def _bin_pack_bedrock_content(
1405 content: Sequence[BedrockContentItem],
1406 budget: int,
1407 ) -> tuple[tuple[BedrockContentItem, ...], ...]:
1408 """Pack whole content items, in order, into batches whose combined text
1409 length stays within `budget`, in a single pass that carries the running
1410 total rather than re-summing the open batch per item.
1412 This is the fast-path half of the hybrid chunking strategy: bin-packing
1413 at a conservative fixed budget keeps the common case at O(n / budget)
1414 ApplyGuardrail calls instead of the O(log n) round trips pure reactive
1415 bisection pays on every oversized request. An item whose own text
1416 already exceeds `budget` is not split here -- it becomes its own
1417 (still oversized) batch and is sent as-is; if AWS rejects that batch as
1418 too large, `_apply_guardrail_content_with_chunking`'s existing
1419 recursive-bisection fallback takes over for that batch only.
1421 `budget` comes from the guardrail's ``chunk_budget_chars`` setting and
1422 defaults to 25,000, matching ApplyGuardrail's default quota of 25 text
1423 units (roughly 1,000 characters each) per second. Packing to that size and
1424 posting sequentially is what keeps chunking from tripping the rate quota
1425 and trading a size error for a throttle. Accounts with raised quotas can
1426 configure a larger budget to spend fewer calls.
1428 The budget is not a correctness dependency either way. AWS's effective cap
1429 varies by account, region, and policy, is not a fixed character count, and
1430 cannot be read from config, so any batch it still rejects falls back to
1431 bisection, which self-corrects however wrong the value was. An over-large
1432 budget therefore costs one extra probe-and-bisect round trip rather than
1433 failing the request.
1434 """
1435 if not content:
1436 return (tuple(content),)
1438 lengths: Final = tuple(len((item.get("text") or BedrockTextContent()).get("text") or "") for item in content)
1440 def assign(carried: tuple[int, int], length: int) -> tuple[int, int]:
1441 batch_index, used = carried
1442 if used + length <= budget:
1443 return batch_index, used + length
1444 return batch_index + 1, length
1446 batch_numbers: Final = (index for index, _ in tuple(accumulate(lengths, assign, initial=(0, 0)))[1:])
1447 return tuple(
1448 tuple(item for _, item in group)
1449 for _, group in groupby(zip(batch_numbers, content), key=lambda pair: pair[0])
1450 )
1452 @staticmethod
1453 def _split_bedrock_content(
1454 content: Sequence[BedrockContentItem],
1455 ) -> tuple[tuple[BedrockContentItem, ...], tuple[BedrockContentItem, ...]] | None:
1456 """Bisect `content` into two roughly-equal, non-empty halves.
1458 When `content` already holds more than one item, it is split by list
1459 length. When it holds exactly one item, that item's own text is split
1460 instead (a list of length 1 has no items left to bisect, but one very
1461 long message is still a single content item) -- at the whitespace
1462 character nearest the midpoint rather than a raw character index, so
1463 the cut never lands inside a word/token. This is a plain, lossless
1464 cut with no overlap: concatenating the two fragments in order always
1465 reproduces the original text exactly, so merging back at
1466 ``_merge_logical_unit_outputs`` needs no reconciliation step.
1468 Known, accepted limitation: whitespace splitting only guards against
1469 *accidentally* severing a single token (one denied word, one PII
1470 pattern) across the cut. It does not, and cannot without an overlap
1471 window, stop a *multi-word* denied phrase deliberately positioned to
1472 straddle the boundary -- each fragment can scan clean on its own and
1473 still reassemble into the flagged phrase. AWS's own guidance on this
1474 API acknowledges the same gap for input chunking ("a critical piece of
1475 text could span two (or more) chunks if not carefully divided") with
1476 no documented resolution, and overlap-and-reconcile was evaluated and
1477 rejected for this PR: AWS's masking output has no documented
1478 length-preservation guarantee, so reconciling an overlap region against
1479 masked text is not sound in general. Out of scope for this PR.
1481 Returns None when there is nothing left to split -- a single item
1482 whose text is too short to halve into two non-empty pieces -- so the
1483 caller can give up and propagate the original too-large error instead
1484 of recursing forever.
1485 """
1486 if len(content) > 1:
1487 midpoint: Final = max(1, len(content) // 2)
1488 return tuple(content[:midpoint]), tuple(content[midpoint:])
1490 text_content: Final = content[0].get("text") or BedrockTextContent()
1491 text: Final = text_content.get("text") or ""
1492 if len(text) < 2:
1493 return None
1494 split_at: Final = BedrockGuardrail._nearest_whitespace_split_index(text)
1495 qualifiers: Final = text_content.get("qualifiers")
1497 def fragment(piece: str) -> BedrockContentItem:
1498 block: Final = (
1499 BedrockTextContent(text=piece, qualifiers=qualifiers) if qualifiers else BedrockTextContent(text=piece)
1500 )
1501 return BedrockContentItem(text=block)
1503 return (fragment(text[:split_at]),), (fragment(text[split_at:]),)
1505 @staticmethod
1506 def _nearest_whitespace_split_index(text: str) -> int:
1507 """Return the index nearest `text`'s midpoint that falls on a whitespace
1508 boundary, so splitting `text[:i]` / `text[i:]` there never severs a word.
1510 Any Unicode whitespace counts, not just an ASCII space. Matching only `" "`
1511 would leave the boundary unguarded for exactly the payloads that get large
1512 enough to need splitting: JSON lines, source code, logs and transcripts are
1513 newline or tab delimited, so a deny-listed word sitting at the midpoint of
1514 one would be cut in half, scan clean on both fragments, and reassemble
1515 intact.
1517 The returned index always leaves both sides non-empty, which is what makes
1518 the caller's recursion terminate. A boundary that would put the split at 0
1519 or at ``len(text)`` is discarded: it would hand back a fragment identical to
1520 the text just rejected as too large, AWS would reject that again, and each
1521 retry would re-split it into the same unchanged fragment until the stack ran
1522 out. The dangerous shape is a text whose only space at or after the midpoint
1523 is its final character.
1525 Falls back to the raw midpoint when no usable whitespace boundary exists, either
1526 because `text` has none at all (a single giant token) or because the only
1527 candidates were degenerate. That is still a correct, lossless split, just no
1528 longer guaranteed word-safe for those cases. `text` must be at least two
1529 characters, which `_split_bedrock_content` guarantees, so the midpoint itself
1530 is never degenerate.
1531 """
1532 midpoint: Final = len(text) // 2
1533 before: Final = max((found.end() for found in _BEDROCK_WHITESPACE.finditer(text, 0, midpoint)), default=None)
1534 after_match: Final = _BEDROCK_WHITESPACE.search(text, midpoint)
1535 candidates: Final = sorted(
1536 (split for split in (before, after_match.end() if after_match else None) if split is not None),
1537 key=lambda split: abs(split - midpoint),
1538 )
1539 return next((split for split in candidates if 0 < split < len(text)), midpoint)
1541 @staticmethod
1542 def _is_input_too_large_error(detail: object) -> bool:
1543 """True if `detail` is an AWS error message for input exceeding the
1544 per-request text-unit quota.
1546 Matched on the message rather than the status code on purpose: AWS is not
1547 consistent about which error it raises for this. Observed against a live
1548 guardrail with an active content-filter policy, an oversized request comes
1549 back as a *ThrottlingException* (429) reading ``Input text size (3273 text
1550 units) exceeds the maximum allowed (1000 text units) for the content filter
1551 policy (Classic tier)``, while the documented failure mode is a
1552 ValidationException (400). Keying off the message covers both.
1554 A guardrail *block* is also raised as an HTTPException with status 400,
1555 but its ``detail`` is always a dict (built by
1556 ``_get_http_exception_for_blocked_guardrail``); a non-200 API error's
1557 ``detail`` is always the plain string returned by
1558 ``_parse_bedrock_guardrail_error_response``. Checking ``isinstance(detail,
1559 str)`` is therefore sufficient to never mistake a real block for a
1560 too-large error.
1561 """
1562 if not isinstance(detail, str):
1563 return False
1564 lowered: Final = detail.lower()
1565 return any(substring in lowered for substring in _BEDROCK_TOO_LARGE_ERROR_SUBSTRINGS)
1567 @staticmethod
1568 def _merge_bedrock_guardrail_responses(
1569 chunk_results: Sequence[BedrockContentChunkResult],
1570 ) -> BedrockGuardrailResponse:
1571 """Merge the per-chunk ApplyGuardrail responses of a chunked request into
1572 one, so a caller cannot tell whether chunking happened.
1574 Only ever called with responses that all passed (a block raises
1575 immediately from ``_apply_guardrail_content_with_chunking`` and is never
1576 added to this list). ``action`` is only set on the merged response when
1577 at least one chunk's raw response included it, and left absent otherwise
1578 -- mirroring a real single-call response and matching what
1579 ``_build_tracing_detail`` treats as "Bedrock didn't report an action".
1581 Fields this merge has no opinion on (``actionReason``, ``guardrailCoverage``,
1582 ``blockedResponse``, anything AWS adds later) are carried over from the chunk
1583 responses rather than dropped, so the response and the logged telemetry keep
1584 the shape a single unchunked call returned. The merged keys below win.
1586 Per AWS's documented ApplyGuardrail contract, a single call's ``outputs``
1587 is positionally parallel to the ``content`` items *of that call*: an
1588 entry per item when anything in the call was masked, or an empty list
1589 when nothing in the whole call was masked. Downstream masking
1590 (``_apply_masking_to_messages``) walks the merged ``outputs`` by a single
1591 running index across the *original, unchunked* message list, so a later
1592 chunk's masked text must land at the same global position it would have
1593 if chunking had never happened. Naively concatenating each chunk's
1594 ``outputs`` breaks that whenever a chunk had nothing masked (its empty
1595 list would otherwise silently swallow its items' slots, shifting every
1596 later chunk's masked text left onto the wrong message). So every
1597 item -- masked or not -- always contributes exactly one entry here,
1598 falling back to that item's own original (unmasked) text when its
1599 chunk returned no output for it; a wholly-untouched result is then
1600 collapsed back to an empty ``outputs`` list to match a real single-call
1601 no-op response. A chunk that returns a nonzero output count not equal
1602 to its item count is passed through as-is instead of guessed at, since
1603 AWS's docs don't cover partial masking within one multi-item call.
1604 """
1605 logical_units: Final = BedrockGuardrail._group_fragment_units(chunk_results)
1606 per_unit_outputs: Final = tuple(BedrockGuardrail._merge_logical_unit_outputs(unit) for unit in logical_units)
1607 merged_outputs: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list
1608 output for outputs, _ in per_unit_outputs for output in outputs
1609 ]
1610 any_masked: Final = any(masked for _, masked in per_unit_outputs)
1612 actions: Final = tuple(
1613 chunk_result.response.get("action")
1614 for chunk_result in chunk_results
1615 if isinstance(chunk_result.response.get("action"), str)
1616 )
1617 merged_action: Final = (
1618 "GUARDRAIL_INTERVENED" if "GUARDRAIL_INTERVENED" in actions else (actions[-1] if actions else None)
1619 )
1620 merged_assessments: Final = [ # mutable-ok: logged payload; redaction only traverses dict/list
1621 assessment
1622 for chunk_result in chunk_results
1623 for assessment in (chunk_result.response.get("assessments") or []) # mutable-ok: logged payload
1624 ]
1625 any_usage_reported: Final = any(chunk_result.response.get("usage") for chunk_result in chunk_results)
1627 merged: Final[BedrockGuardrailResponse] = cast( # cast-ok: TypedDict assembled from a comprehension
1628 BedrockGuardrailResponse,
1629 { # mutable-ok: builds the TypedDict payload
1630 key: value for chunk_result in chunk_results for key, value in chunk_result.response.items()
1631 },
1632 )
1633 if merged_action is not None:
1634 merged["action"] = merged_action
1635 if merged_outputs and any_masked:
1636 merged["outputs"] = merged_outputs
1637 merged["output"] = merged_outputs
1638 if merged_assessments:
1639 merged["assessments"] = merged_assessments
1640 if any_usage_reported:
1641 merged["usage"] = BedrockGuardrail._sum_bedrock_guardrail_usage(chunk_results)
1642 return merged
1644 @staticmethod
1645 def _sum_bedrock_guardrail_usage(
1646 chunk_results: Sequence[BedrockContentChunkResult],
1647 ) -> BedrockGuardrailUsage:
1648 """Sum each chunk's ``usage`` counters field-by-field into one totals dict.
1650 Keys are taken from the responses rather than from a fixed list, so a counter
1651 this code does not know about (AWS has added several) is still summed and
1652 reported instead of being silently dropped to zero."""
1653 return BedrockGuardrail._sum_usage_counters(
1654 tuple(
1655 chunk_result.response.get("usage") or {} # mutable-ok: read-only empty fallback
1656 for chunk_result in chunk_results
1657 )
1658 )
1660 @staticmethod
1661 def _sum_usage_counters(usages: Sequence[BedrockGuardrailUsage]) -> BedrockGuardrailUsage:
1662 return cast( # cast-ok: TypedDict assembled from a comprehension
1663 BedrockGuardrailUsage,
1664 { # mutable-ok: builds the TypedDict payload
1665 key: sum(usage.get(key) or 0 for usage in usages)
1666 for key in dict.fromkeys(key for usage in usages for key in usage)
1667 },
1668 )
1670 @staticmethod
1671 def _group_fragment_units(
1672 chunk_results: Sequence[BedrockContentChunkResult],
1673 ) -> tuple[tuple[BedrockContentChunkResult, ...], ...]:
1674 """Group consecutive text-fragment chunk results back into the one content
1675 item each group came from, leaving every ordinary chunk result as a unit of
1676 one.
1678 The group size is read off the results themselves rather than assumed,
1679 because a single content item can be bisected repeatedly: two levels of
1680 splitting yield four fragments for one item, not two. Assuming a fixed pair
1681 here would emit two outputs for one message and shift every later message's
1682 masked text onto the wrong message."""
1684 def advance(carried: tuple[int, bool], result: BedrockContentChunkResult) -> tuple[int, bool]:
1685 remaining, _ = carried
1686 if remaining == 0:
1687 return max(1, result.fragment_group_size) - 1, True
1688 return remaining - 1, False
1690 starts: Final = tuple(
1691 index
1692 for index, (_, starts_unit) in enumerate(tuple(accumulate(chunk_results, advance, initial=(0, False)))[1:])
1693 if starts_unit
1694 )
1695 return tuple(tuple(chunk_results[start:end]) for start, end in zip(starts, starts[1:] + (len(chunk_results),)))
1697 @staticmethod
1698 def _merge_logical_unit_outputs(
1699 unit: tuple[BedrockContentChunkResult, ...],
1700 ) -> tuple[tuple[BedrockGuardrailOutput, ...], bool]:
1701 """Reduce one logical unit (a fragment group of any size, or a single chunk
1702 result) to the ``BedrockGuardrailOutput`` entries it contributes to the
1703 merged response, plus whether any masking actually happened in it.
1705 Per AWS's documented ApplyGuardrail contract, a single call's
1706 ``outputs`` is positionally parallel to the ``content`` items *of that
1707 call*: an entry per item when anything in the call was masked, or an
1708 empty list when nothing in the whole call was masked. Downstream
1709 masking (``_apply_masking_to_messages``) walks the merged ``outputs``
1710 by a single running index across the *original, unchunked* message
1711 list, so a later chunk's masked text must land at the same global
1712 position it would have if chunking had never happened. So every item
1713 -- masked or not -- always contributes exactly one entry here, falling
1714 back to that item's own original (unmasked) text when its chunk
1715 returned no output for it. A chunk that returns a nonzero output count
1716 not equal to its item count is passed through as-is instead of guessed
1717 at, since AWS's docs don't cover partial masking within one multi-item
1718 call.
1720 A unit holding more than one result is a fragment group: every result in it
1721 is one fragment of a single content item's text, so the group collapses to
1722 one entry built from each fragment's masked text (or that fragment's own
1723 original text where it came back unmasked), concatenated in order. This
1724 holds for any group size, not only two.
1725 """
1726 if len(unit) > 1:
1728 def fragment_outputs(result: BedrockContentChunkResult) -> tuple[BedrockGuardrailOutput, ...]:
1729 return tuple(result.response.get("outputs") or result.response.get("output") or ())
1731 def fragment_text(result: BedrockContentChunkResult) -> str:
1732 source: Final = (result.content[0].get("text") or {}).get( # mutable-ok: read-only fallback
1733 "text"
1734 ) or ""
1735 outputs: Final = fragment_outputs(result)
1736 masked: Final = outputs[0].get("text") if outputs else None
1737 return masked if masked is not None else source
1739 merged_text: Final = "".join(fragment_text(result) for result in unit)
1740 any_masked: Final = any(fragment_outputs(result) for result in unit)
1741 return (BedrockGuardrailOutput(text=merged_text),), any_masked
1743 (chunk_result,) = unit
1744 chunk_outputs: Final = chunk_result.response.get("outputs") or chunk_result.response.get("output") or ()
1745 if len(chunk_outputs) == len(chunk_result.content):
1746 return tuple(chunk_outputs), bool(chunk_outputs)
1747 if not chunk_outputs:
1748 return tuple(
1749 BedrockGuardrailOutput(
1750 text=(item.get("text") or {}).get("text") or "" # mutable-ok: read-only fallback
1751 )
1752 for item in chunk_result.content
1753 ), False
1754 return tuple(chunk_outputs), True
1756 async def _sign_and_post(
1757 self,
1758 prepared_request: "AWSPreparedRequest",
1759 request_data: dict | None,
1760 event_type: GuardrailEventHooks,
1761 start_time: "datetime",
1762 log_transport_failure: bool = True,
1763 ) -> httpx.Response:
1764 """POST a signed Bedrock request, logging+raising on network/HTTP errors.
1766 Shared by both the ApplyGuardrail and InvokeGuardrailChecks paths so their
1767 transport-error handling cannot drift. Returns the raw ``httpx.Response`` on
1768 success (including non-2xx that httpx did not raise on); the 200-path logging,
1769 status and tracing stay with each caller because the two APIs report differently.
1771 ``log_transport_failure=False`` suppresses the ``guardrail_failed_to_respond``
1772 entry for a non-200 that is re-raised as an ``HTTPException``, for callers that
1773 own consolidated per-request logging. The ApplyGuardrail path needs this:
1774 ``AsyncHTTPHandler.post`` calls ``raise_for_status()``, so every non-200 lands
1775 in this handler, and one logical request can legitimately produce several of
1776 them (a too-large probe, then each rejected bisection level) while still
1777 succeeding overall. Logging per attempt would report a recovered request as
1778 several failures plus a success.
1780 The connection-level branch below (timeout, endpoint down) still logs
1781 unconditionally: it re-raises the original exception rather than an
1782 ``HTTPException``, so no consolidating caller catches it, and suppressing it
1783 would drop the only record of the failure.
1784 """
1785 try:
1786 return await self.async_handler.post(
1787 url=prepared_request.url,
1788 data=prepared_request.body,
1789 headers=prepared_request.headers,
1790 )
1791 except HTTPException:
1792 # Propagate HTTPException (e.g. from non-200 path) as-is
1793 raise
1794 except Exception as e:
1795 # If this is an HTTP error with a response body (e.g. httpx.HTTPStatusError),
1796 # extract the AWS error message and propagate it
1797 err_response: Final = getattr(e, "response", None)
1798 if isinstance(err_response, httpx.Response):
1799 try:
1800 (
1801 status_code,
1802 detail_message,
1803 ) = self._parse_bedrock_guardrail_error_response(err_response)
1804 if log_transport_failure:
1805 self.add_standard_logging_guardrail_information_to_request_data(
1806 guardrail_provider=self.guardrail_provider,
1807 guardrail_json_response={ # mutable-ok: logging helper requires a dict
1808 "error": detail_message
1809 },
1810 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1811 guardrail_status="guardrail_failed_to_respond",
1812 start_time=start_time.timestamp(),
1813 end_time=datetime.now(timezone.utc).timestamp(),
1814 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1815 event_type=event_type,
1816 )
1817 raise HTTPException(status_code=status_code, detail=detail_message) from e
1818 except HTTPException:
1819 raise
1820 # Endpoint down, timeout, or other HTTP/network errors
1821 verbose_proxy_logger.error("Bedrock AI: failed to make guardrail request: %s", str(e))
1822 self.add_standard_logging_guardrail_information_to_request_data(
1823 guardrail_provider=self.guardrail_provider,
1824 guardrail_json_response={"error": str(e)},
1825 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1826 guardrail_status="guardrail_failed_to_respond",
1827 start_time=start_time.timestamp(),
1828 end_time=datetime.now(timezone.utc).timestamp(),
1829 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1830 event_type=event_type,
1831 )
1832 raise
1834 ########### InvokeGuardrailChecks (resource-less, detect-only) ############
1836 @staticmethod
1837 def _chunk_texts_into_checks_messages(
1838 role: Literal["user", "assistant", "system"], texts: list[str]
1839 ) -> list[BedrockChecksMessage]:
1840 """Group ``texts`` into role-tagged messages of <= the API content-block cap.
1842 A source message with more text blocks than the per-message limit is split
1843 across multiple messages so EVERY block is scanned. Truncating instead would
1844 let a user hide prohibited content past the limit (guardrail bypass).
1845 """
1846 cap: Final = _BEDROCK_CHECKS_MAX_CONTENT_BLOCKS
1847 return [
1848 BedrockChecksMessage(
1849 role=role,
1850 content=[{"text": text} for text in texts[start : start + cap]],
1851 )
1852 for start in range(0, len(texts), cap)
1853 ]
1855 def _build_invoke_guardrail_checks_messages(
1856 self,
1857 source: Literal["INPUT", "OUTPUT"],
1858 messages: list[AllMessageValues] | None = None,
1859 response: litellm.ModelResponse | None = None,
1860 ) -> list[BedrockChecksMessage]:
1861 """Build the role-tagged `messages` array for InvokeGuardrailChecks.
1863 INPUT scans the request messages, OUTPUT scans the model response as an
1864 ``assistant`` turn. Every non-empty text block of every message is scanned;
1865 messages exceeding the per-message content-block cap are split into multiple
1866 messages rather than truncated.
1868 INPUT content is tagged ``user`` regardless of the caller-supplied role.
1869 Bedrock excludes ``system`` content from prompt-attack evaluation, so
1870 trusting a caller's ``system``/``developer`` label would let an injection
1871 avoid the promptAttack check. At the proxy every INPUT message is
1872 caller-controlled, so all of it is treated as untrusted user input, matching
1873 AWS guidance to tag untrusted content as user input.
1874 """
1875 if source == "OUTPUT":
1876 # Reuse the ApplyGuardrail output extractor (single source of truth for
1877 # pulling assistant text out of a ModelResponse), then re-tag as an
1878 # assistant turn for the role-based InvokeGuardrailChecks payload.
1879 output_request: Final = self._create_bedrock_output_content_request(response=response)
1880 output_texts: Final = [
1881 text for item in output_request.get("content") or [] if (text := (item.get("text") or {}).get("text"))
1882 ]
1883 return self._chunk_texts_into_checks_messages("assistant", output_texts)
1885 return [
1886 checks_message
1887 for message in messages or []
1888 for checks_message in self._chunk_texts_into_checks_messages(
1889 "user",
1890 [block.text for block in self.get_content_items_for_message(message) or [] if block.text],
1891 )
1892 ]
1894 async def _make_invoke_guardrail_checks_request(
1895 self,
1896 source: Literal["INPUT", "OUTPUT"],
1897 messages: list[AllMessageValues] | None = None,
1898 response: litellm.ModelResponse | None = None,
1899 request_data: dict | None = None,
1900 logging_event_type: GuardrailEventHooks | None = None,
1901 ) -> BedrockGuardrailResponse:
1902 """Run the resource-less InvokeGuardrailChecks API and enforce thresholds.
1904 Detect-only: the API returns scores, never rewritten content. We map scores
1905 to a block decision via the configured thresholds. On a pass we return an
1906 empty ``BedrockGuardrailResponse`` (downstream masking treats it as a no-op).
1907 """
1908 start_time: Final = datetime.now(timezone.utc)
1910 checks_messages: Final = self._build_invoke_guardrail_checks_messages(
1911 source=source, messages=messages, response=response
1912 )
1913 if not checks_messages:
1914 # Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does.
1915 return BedrockGuardrailResponse()
1917 api_key: Final[str | None] = request_data.get("api_key") if request_data else None
1918 credentials, aws_region_name = await run_aws_signing(
1919 self._load_credentials, bearer_token=bedrock_bearer_token(api_key)
1920 )
1921 body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
1923 prepared_request: Final = await run_aws_signing(
1924 self._prepare_request,
1925 credentials=credentials,
1926 data=body,
1927 optional_params=self.optional_params,
1928 aws_region_name=aws_region_name,
1929 api_key=api_key,
1930 request_path=_BEDROCK_INVOKE_GUARDRAIL_CHECKS_PATH,
1931 )
1932 verbose_proxy_logger.debug("Bedrock InvokeGuardrailChecks request url: %s", prepared_request.url)
1934 event_type: Final = logging_event_type or (
1935 GuardrailEventHooks.pre_call if source == "INPUT" else GuardrailEventHooks.post_call
1936 )
1938 httpx_response: Final = await self._sign_and_post(
1939 prepared_request=prepared_request,
1940 request_data=request_data,
1941 event_type=event_type,
1942 start_time=start_time,
1943 )
1945 if httpx_response.status_code != 200:
1946 status_code, detail_message = self._parse_bedrock_guardrail_error_response(httpx_response)
1947 verbose_proxy_logger.error(
1948 "Bedrock InvokeGuardrailChecks: error response. Status %s: %s",
1949 httpx_response.status_code,
1950 detail_message,
1951 )
1952 self.add_standard_logging_guardrail_information_to_request_data(
1953 guardrail_provider=self.guardrail_provider,
1954 guardrail_json_response={"error": detail_message},
1955 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1956 guardrail_status="guardrail_failed_to_respond",
1957 start_time=start_time.timestamp(),
1958 end_time=datetime.now(timezone.utc).timestamp(),
1959 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1960 event_type=event_type,
1961 )
1962 raise HTTPException(status_code=status_code, detail=detail_message)
1964 try:
1965 json_response: Final = TypeAdapter(BedrockGuardrailChecksResponse).validate_python(httpx_response.json())
1966 except (ValidationError, ValueError) as e:
1967 verbose_proxy_logger.error("Bedrock InvokeGuardrailChecks: unparseable 200 response: %s", str(e))
1968 self.add_standard_logging_guardrail_information_to_request_data(
1969 guardrail_provider=self.guardrail_provider,
1970 guardrail_json_response={"error": str(e)},
1971 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1972 guardrail_status="guardrail_failed_to_respond",
1973 start_time=start_time.timestamp(),
1974 end_time=datetime.now(timezone.utc).timestamp(),
1975 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1976 event_type=event_type,
1977 )
1978 raise HTTPException(
1979 status_code=500,
1980 detail={"error": "Bedrock InvokeGuardrailChecks returned an unexpected response shape"},
1981 ) from e
1982 violations: Final = self._collect_invoke_checks_violations(json_response)
1984 # Log a copy with PII location offsets stripped: offsets + the (separately
1985 # logged) request messages would otherwise reconstruct the detected PII span.
1986 self.add_standard_logging_guardrail_information_to_request_data(
1987 guardrail_provider=self.guardrail_provider,
1988 guardrail_json_response=self._sanitize_invoke_checks_response_for_logging(json_response),
1989 request_data=request_data or {}, # mutable-ok: logging helper requires a dict
1990 guardrail_status=self._get_invoke_checks_status(bool(violations)),
1991 start_time=start_time.timestamp(),
1992 end_time=datetime.now(timezone.utc).timestamp(),
1993 duration=(datetime.now(timezone.utc) - start_time).total_seconds(),
1994 event_type=event_type,
1995 tracing_detail=self._build_invoke_checks_tracing_detail(violations) if violations else None,
1996 )
1998 if violations:
1999 raise self._get_block_exception_for_checks(violations, request_data=request_data)
2001 return BedrockGuardrailResponse()
2003 def _collect_invoke_checks_violations(
2004 self, response: BedrockGuardrailChecksResponse | None
2005 ) -> list[BedrockChecksViolation]:
2006 """Return the check results whose score meets/exceeds the configured threshold.
2008 Only checks present in the configured ``checks`` block are evaluated; a
2009 threshold of ``None`` makes that check detect-only (never contributes a
2010 violation). A truncated sensitiveInformation result counts as a violation
2011 (fail closed: omitted detections were never scored). Only the non-sensitive
2012 label (category/type) and the numeric score are kept -- never offsets or
2013 matched text.
2014 """
2015 results: Final[dict[str, Any]] = dict((response or {}).get("results") or {})
2016 # (results key, score field, label field, threshold). PII uses
2017 # confidenceScore/type; the other two use severityScore/category.
2018 check_specs: Final = [
2019 (
2020 "contentFilter",
2021 "severityScore",
2022 "category",
2023 self.content_filter_threshold,
2024 ),
2025 ("promptAttack", "severityScore", "category", self.prompt_attack_threshold),
2026 (
2027 "sensitiveInformation",
2028 "confidenceScore",
2029 "type",
2030 self.pii_confidence_threshold,
2031 ),
2032 ]
2034 configured_checks: Final = self.checks or {}
2035 violations: Final[list[BedrockChecksViolation]] = []
2036 for check_key, score_field, label_field, threshold in check_specs:
2037 if threshold is None or check_key not in configured_checks:
2038 continue
2039 check_result = results.get(check_key) or {}
2040 if check_key == "sensitiveInformation" and check_result.get("truncated"):
2041 violations.append({"check": check_key, "truncated": True})
2042 for entry in check_result.get("results") or []:
2043 score = entry.get(score_field)
2044 if isinstance(score, (int, float)) and float(score) >= threshold:
2045 violation: BedrockChecksViolation = (
2046 {"check": check_key, "category": entry.get("category"), "severityScore": float(score)}
2047 if score_field == "severityScore"
2048 else {"check": check_key, "type": entry.get("type"), "confidenceScore": float(score)}
2049 )
2050 violations.append(violation)
2051 return violations
2053 @staticmethod
2054 def _sanitize_invoke_checks_response_for_logging(
2055 response: BedrockGuardrailChecksResponse,
2056 ) -> dict[str, Any]:
2057 """Strip PII location offsets from a checks response before it is logged."""
2058 sanitized: Final[dict[str, Any]] = copy.deepcopy(dict(response))
2059 sensitive: Final = (sanitized.get("results") or {}).get("sensitiveInformation") or {}
2060 for entry in sensitive.get("results") or []:
2061 if isinstance(entry, dict):
2062 for key in _BEDROCK_CHECKS_PII_LOCATION_KEYS:
2063 entry.pop(key, None)
2064 return sanitized
2066 @staticmethod
2067 def _get_invoke_checks_status(over_threshold: bool) -> GuardrailStatus:
2068 return "guardrail_intervened" if over_threshold else "success"
2070 @staticmethod
2071 def _build_invoke_checks_tracing_detail(
2072 violations: list[BedrockChecksViolation],
2073 ) -> GuardrailTracingDetail:
2074 tracing_detail: Final[GuardrailTracingDetail] = {}
2075 categories: Final = [
2076 label
2077 for label in (v.get("category") or v.get("type") for v in violations)
2078 if isinstance(label, str) and label
2079 ]
2080 if categories:
2081 tracing_detail["violation_categories"] = categories
2082 tracing_detail["guardrail_action"] = "GUARDRAIL_INTERVENED" if violations else "NONE"
2083 return tracing_detail
2085 def _get_block_exception_for_checks(
2086 self, violations: list[BedrockChecksViolation], request_data: dict | None = None
2087 ) -> HTTPException | ModifyResponseException:
2088 """Build the block exception for an over-threshold InvokeGuardrailChecks result.
2090 Mirrors ``_get_http_exception_for_blocked_guardrail``'s return-type branching.
2091 The detail carries only non-sensitive labels + scores (no offsets / raw input).
2092 """
2093 if self.disable_exception_on_block is True:
2094 _request_data: Final = request_data or {}
2095 return ModifyResponseException(
2096 message="Violated guardrail policy",
2097 model=_request_data.get("model", "bedrock-guardrail"),
2098 request_data=_request_data,
2099 guardrail_name=self.guardrail_name,
2100 )
2101 return HTTPException(
2102 status_code=400,
2103 detail={
2104 "error": "Violated guardrail policy",
2105 "bedrock_guardrail_checks": violations,
2106 },
2107 )
2109 def _check_bedrock_response_for_exception(self, response) -> bool:
2110 """
2111 Return True if the Bedrock ApplyGuardrail response indicates an exception.
2113 Works with real httpx.Response objects and MagicMock responses used in tests.
2114 """
2115 payload = None
2117 try:
2118 json_method: Final = getattr(response, "json", None)
2119 if callable(json_method):
2120 payload = json_method()
2121 except Exception:
2122 payload = None
2124 if payload is None:
2125 try:
2126 raw: Final = getattr(response, "content", None)
2127 if isinstance(raw, (bytes, bytearray)):
2128 payload = json.loads(raw.decode("utf-8"))
2129 else:
2130 text: Final = getattr(response, "text", None)
2131 if isinstance(text, str):
2132 payload = json.loads(text)
2133 except Exception:
2134 # Can't parse -> assume no explicit Exception marker
2135 return False
2137 if not isinstance(payload, dict):
2138 return False
2140 return "Exception" in payload.get("Output", {}).get("__type", "")
2142 def _get_bedrock_guardrail_response_status(self, response: httpx.Response) -> GuardrailStatus:
2143 """
2144 Get the status of the bedrock guardrail response.
2146 Returns:
2147 "success": Content allowed through with no violations
2148 "guardrail_intervened": Content blocked due to policy violations
2149 "guardrail_failed_to_respond": Technical error or API failure
2150 """
2151 if response.status_code == 200:
2152 if self._check_bedrock_response_for_exception(response):
2153 return "guardrail_failed_to_respond"
2155 # Check if the guardrail would block content
2156 try:
2157 _json_response: Final = response.json()
2158 bedrock_guardrail_response: Final = BedrockGuardrailResponse(**_json_response)
2159 if self._should_raise_guardrail_blocked_exception(bedrock_guardrail_response):
2160 return "guardrail_intervened"
2161 except Exception:
2162 pass
2164 return "success"
2165 return "guardrail_failed_to_respond"
2167 def _parse_bedrock_guardrail_error_response(self, response: httpx.Response) -> tuple[int, str]:
2168 """
2169 Parse AWS Bedrock guardrail error response body to extract status code and message.
2171 AWS may return shapes like {"message": "..."} or {"error": {"message": "..."}}.
2172 Returns (status_code, message) for use in HTTPException.
2173 """
2174 status_code: Final = response.status_code
2175 message: Final = "Bedrock guardrail request failed"
2176 try:
2177 body: Final = response.json()
2178 except Exception:
2179 text: Final = getattr(response, "text", None) or ""
2180 if isinstance(text, str) and text.strip():
2181 return (status_code, text.strip())
2182 return (status_code, message)
2183 if isinstance(body, dict):
2184 if isinstance(body.get("message"), str):
2185 return (status_code, body["message"])
2186 err: Final = body.get("error")
2187 if isinstance(err, dict) and isinstance(err.get("message"), str):
2188 return (status_code, err["message"])
2189 if isinstance(err, str):
2190 return (status_code, err)
2191 return (status_code, message)
2193 def _build_tracing_detail(
2194 self, response: BedrockGuardrailResponse, aws_region_name: str | None
2195 ) -> GuardrailTracingDetail:
2196 """
2197 Build the tracing detail from the raw Bedrock response, before
2198 redaction, so downstream loggers (OTEL, Langfuse, ...) get the
2199 actual category names rather than the "[REDACTED]" sentinel that
2200 replaces customWords.match later. Bedrock's top-level ``action``
2201 field ("GUARDRAIL_INTERVENED" or "NONE") is also surfaced so the
2202 OTEL integration can expose it as a queryable span attribute
2203 without re-parsing the redacted guardrail_response blob.
2204 """
2205 violation_categories: Final = self._extract_violation_category_names(response)
2206 bedrock_action: Final = response.get("action")
2207 categories_detail: Final[GuardrailTracingDetail] = {"violation_categories": violation_categories}
2208 action_detail: Final[GuardrailTracingDetail] = {"guardrail_action": bedrock_action}
2209 tracing_detail: Final[GuardrailTracingDetail] = {
2210 **(categories_detail if violation_categories else _NO_TRACING_DETAIL),
2211 **(action_detail if isinstance(bedrock_action, str) else _NO_TRACING_DETAIL),
2212 **self._usage_tracing_detail(response.get("usage"), aws_region_name),
2213 }
2214 return tracing_detail
2216 @staticmethod
2217 def _usage_tracing_detail(
2218 usage: BedrockGuardrailUsage | None, aws_region_name: str | None
2219 ) -> GuardrailTracingDetail:
2220 if not isinstance(usage, dict):
2221 return _NO_TRACING_DETAIL
2222 usage_units: Final = { # mutable-ok: json.dumps'd into spend log metadata downstream
2223 key: value for key, value in usage.items() if isinstance(value, int)
2224 }
2225 if not usage_units:
2226 return _NO_TRACING_DETAIL
2227 cost_by_unit: Final = bedrock_guardrail_cost_by_unit(usage_units=usage_units, aws_region_name=aws_region_name)
2228 priced_detail: Final[GuardrailTracingDetail] = {"guardrail_cost_by_unit": cost_by_unit}
2229 usage_detail: Final[GuardrailTracingDetail] = {
2230 "guardrail_usage": usage_units,
2231 "guardrail_cost": guardrail_cost_total(cost_by_unit),
2232 **(priced_detail if cost_by_unit is not None else _NO_TRACING_DETAIL),
2233 }
2234 return usage_detail
2236 def _extract_violation_category_names(self, response: BedrockGuardrailResponse) -> list[str]:
2237 """
2238 Flatten the BLOCKED assessments into a list of human-readable category
2239 names suitable for queryable OTEL / standard-logging attributes.
2241 SECURITY: only emits the non-sensitive policy *label* (topic name,
2242 content-filter type, PII entity type, named-regex name). The raw
2243 ``match`` field is intentionally NOT used — it carries the user's
2244 original input that triggered the rule (e.g. a credit-card number
2245 that hit a regex, or the literal custom word). Surfacing it to
2246 telemetry would re-introduce the sensitive content the guardrail
2247 was supposed to keep out. Entries that only have a ``match`` (bare
2248 customWords, unnamed regexes) are therefore skipped — operators
2249 can still see the count in ``_extract_blocked_assessments`` which
2250 feeds the HTTP error detail.
2251 """
2252 names: Final[list[str]] = []
2253 for block in self._extract_blocked_assessments(response):
2254 for match in block.get("matches", []) or []:
2255 # Allow-list non-sensitive labels only. Never fall back to
2256 # `match.get("match")` — that's user-submitted content.
2257 label = match.get("name") or match.get("type")
2258 if isinstance(label, str) and label:
2259 names.append(label)
2260 return names
2262 def _extract_blocked_assessments(self, response: BedrockGuardrailResponse) -> list[dict]:
2263 """
2264 Walk the Bedrock guardrail response and emit a structured list of
2265 BLOCKED assessment entries describing exactly which policies fired.
2267 Mirrors the iteration in `_should_raise_guardrail_blocked_exception()`
2268 but produces a list of `{policy, matches}` dicts instead of a bool.
2269 Each `match` carries the originating subcategory, type, action, and
2270 matched term where available, so the client can render a precise
2271 explanation of the violation.
2272 """
2273 blocked: Final[list[dict]] = []
2274 assessments: Final = response.get("assessments", []) or []
2276 for assessment in assessments:
2277 # Topic policy
2278 topic_policy = assessment.get("topicPolicy")
2279 if topic_policy:
2280 topic_matches = [
2281 {
2282 "category": "topics",
2283 "name": t.get("name"),
2284 "type": t.get("type"),
2285 "action": t.get("action"),
2286 }
2287 for t in (topic_policy.get("topics") or [])
2288 if t.get("action") == "BLOCKED"
2289 ]
2290 if topic_matches:
2291 blocked.append({"policy": "topicPolicy", "matches": topic_matches})
2293 # Content policy
2294 content_policy = assessment.get("contentPolicy")
2295 if content_policy:
2296 content_matches = [
2297 {
2298 "category": "filters",
2299 "type": f.get("type"),
2300 "confidence": f.get("confidence"),
2301 "filterStrength": f.get("filterStrength"),
2302 "action": f.get("action"),
2303 }
2304 for f in (content_policy.get("filters") or [])
2305 if f.get("action") == "BLOCKED"
2306 ]
2307 if content_matches:
2308 blocked.append({"policy": "contentPolicy", "matches": content_matches})
2310 # Word policy
2311 word_policy = assessment.get("wordPolicy")
2312 if word_policy:
2313 word_matches: list[dict] = []
2314 for w in word_policy.get("customWords") or []:
2315 if w.get("action") == "BLOCKED":
2316 word_matches.append(
2317 {
2318 "category": "customWords",
2319 "match": w.get("match"),
2320 "action": w.get("action"),
2321 }
2322 )
2323 for mw in word_policy.get("managedWordLists") or []:
2324 if mw.get("action") == "BLOCKED":
2325 word_matches.append(
2326 {
2327 "category": "managedWordLists",
2328 "type": mw.get("type"),
2329 "match": mw.get("match"),
2330 "action": mw.get("action"),
2331 }
2332 )
2333 if word_matches:
2334 blocked.append({"policy": "wordPolicy", "matches": word_matches})
2336 # Sensitive information policy (PII)
2337 sensitive_info = assessment.get("sensitiveInformationPolicy")
2338 if sensitive_info:
2339 pii_matches: list[dict] = []
2340 for p in sensitive_info.get("piiEntities") or []:
2341 if p.get("action") == "BLOCKED":
2342 pii_matches.append(
2343 {
2344 "category": "piiEntities",
2345 "type": p.get("type"),
2346 "match": p.get("match"),
2347 "action": p.get("action"),
2348 }
2349 )
2350 for r in sensitive_info.get("regexes") or []:
2351 if r.get("action") == "BLOCKED":
2352 pii_matches.append(
2353 {
2354 "category": "regexes",
2355 "name": r.get("name"),
2356 "regex": r.get("regex"),
2357 "match": r.get("match"),
2358 "action": r.get("action"),
2359 }
2360 )
2361 if pii_matches:
2362 blocked.append(
2363 {
2364 "policy": "sensitiveInformationPolicy",
2365 "matches": pii_matches,
2366 }
2367 )
2369 # Contextual grounding policy
2370 contextual = assessment.get("contextualGroundingPolicy")
2371 if contextual:
2372 grounding_matches = [
2373 {
2374 "category": "filters",
2375 "type": f.get("type"),
2376 "threshold": f.get("threshold"),
2377 "score": f.get("score"),
2378 "action": f.get("action"),
2379 }
2380 for f in (contextual.get("filters") or [])
2381 if f.get("action") == "BLOCKED"
2382 ]
2383 if grounding_matches:
2384 blocked.append(
2385 {
2386 "policy": "contextualGroundingPolicy",
2387 "matches": grounding_matches,
2388 }
2389 )
2391 return blocked
2393 def _get_http_exception_for_blocked_guardrail(
2394 self, response: BedrockGuardrailResponse, request_data: dict | None = None
2395 ) -> HTTPException | ModifyResponseException:
2396 """
2397 Get the HTTP exception for a blocked guardrail.
2398 """
2399 bedrock_guardrail_output_text: str = ""
2400 outputs: Final[list[BedrockGuardrailOutput] | None] = response.get("outputs", []) or []
2401 if outputs:
2402 for output in outputs:
2403 if output.get("text"):
2404 bedrock_guardrail_output_text += output.get("text") or ""
2406 if self.disable_exception_on_block is True:
2407 _request_data: Final = request_data or {}
2408 return ModifyResponseException(
2409 message=bedrock_guardrail_output_text,
2410 model=_request_data.get("model", "bedrock-guardrail"),
2411 request_data=_request_data,
2412 guardrail_name=self.guardrail_name,
2413 )
2415 detail: Final[dict[str, object]] = {
2416 "error": "Violated guardrail policy",
2417 "bedrock_guardrail_response": bedrock_guardrail_output_text,
2418 }
2419 if self.guardrailIdentifier:
2420 detail["guardrailIdentifier"] = self.guardrailIdentifier
2421 if self.guardrailVersion:
2422 detail["guardrailVersion"] = self.guardrailVersion
2424 assessments: Final = self._extract_blocked_assessments(response)
2425 if assessments:
2426 detail["assessments"] = _redact_assessment_match_fields(assessments)
2428 return HTTPException(status_code=400, detail=detail)
2430 def _should_raise_guardrail_blocked_exception(self, response: BedrockGuardrailResponse) -> bool:
2431 """
2432 Only raise exception for "BLOCKED" actions, not for "ANONYMIZED" actions.
2434 If `self.mask_request_content` or `self.mask_response_content` is set to `True`,
2435 then use the output from the guardrail to mask the request or response content.
2437 However, even with masking enabled, content with action="BLOCKED" should still
2438 raise an exception, only content with action="ANONYMIZED" should be masked.
2439 """
2441 # if no intervention, return False
2442 if response.get("action") != "GUARDRAIL_INTERVENED":
2443 return False
2445 # Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED)
2446 # NOTE: Use `.get("k") or []` not `.get("k", [])` — Bedrock can return explicit
2447 # JSON null; dict.get("k", []) then yields None, and `for x in None` raises.
2448 assessments: Final = response.get("assessments") or []
2449 if not assessments:
2450 return False
2452 for assessment in assessments:
2453 # Check topic policy
2454 topic_policy = assessment.get("topicPolicy")
2455 if topic_policy:
2456 topics = topic_policy.get("topics") or []
2457 for topic in topics:
2458 if topic.get("action") == "BLOCKED":
2459 return True
2461 # Check content policy
2462 content_policy = assessment.get("contentPolicy")
2463 if content_policy:
2464 filters = content_policy.get("filters") or []
2465 for filter_item in filters:
2466 if filter_item.get("action") == "BLOCKED":
2467 return True
2469 # Check word policy
2470 word_policy = assessment.get("wordPolicy")
2471 if word_policy:
2472 custom_words = word_policy.get("customWords") or []
2473 for custom_word in custom_words:
2474 if custom_word.get("action") == "BLOCKED":
2475 return True
2476 managed_words = word_policy.get("managedWordLists") or []
2477 for managed_word in managed_words:
2478 if managed_word.get("action") == "BLOCKED":
2479 return True
2481 # Check sensitive information policy
2482 sensitive_info_policy = assessment.get("sensitiveInformationPolicy")
2483 if sensitive_info_policy:
2484 pii_entities = sensitive_info_policy.get("piiEntities") or []
2485 if pii_entities:
2486 for pii_entity in pii_entities:
2487 if pii_entity.get("action") == "BLOCKED":
2488 return True
2489 regexes = sensitive_info_policy.get("regexes") or []
2490 if regexes:
2491 for regex in regexes:
2492 if regex.get("action") == "BLOCKED":
2493 return True
2495 # Check contextual grounding policy
2496 contextual_grounding_policy = assessment.get("contextualGroundingPolicy")
2497 if contextual_grounding_policy:
2498 grounding_filters = contextual_grounding_policy.get("filters") or []
2499 for grounding_filter in grounding_filters:
2500 if grounding_filter.get("action") == "BLOCKED":
2501 return True
2503 # If we got here, intervention occurred but no BLOCKED actions found
2504 # This means all actions were ANONYMIZED or NONE, so don't raise exception
2505 return False
2507 async def async_pre_call_hook(
2508 self,
2509 user_api_key_dict: UserAPIKeyAuth,
2510 cache: DualCache,
2511 data: dict,
2512 call_type: CallTypesLiteral,
2513 ) -> Exception | str | dict | None:
2514 verbose_proxy_logger.debug("Inside Bedrock Pre-Call Hook for call_type: %s", call_type)
2516 from litellm.proxy.common_utils.callback_utils import (
2517 add_guardrail_to_applied_guardrails_header,
2518 )
2520 event_type: Final[GuardrailEventHooks] = GuardrailEventHooks.pre_call
2521 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
2522 return data
2524 new_messages: Final = self.get_guardrails_messages_for_call_type(
2525 call_type=cast(CallTypes, call_type),
2526 data=data,
2527 )
2529 # Handle None case
2530 if new_messages is None:
2531 verbose_proxy_logger.debug("No messages found for call_type, skipping guardrail")
2532 return data
2534 filter_result: Final = self._prepare_guardrail_messages_for_role(messages=new_messages)
2536 filtered_messages: Final = filter_result.payload_messages
2537 if not filtered_messages:
2538 verbose_proxy_logger.debug("No user-role messages available for guardrail payload")
2539 return data
2541 #########################################################
2542 ########## 1. Make the Bedrock API request ##########
2543 #########################################################
2544 # A block with disable_exception_on_block=True raises ModifyResponseException
2545 # from make_bedrock_api_request; that propagates to the endpoint handler,
2546 # which returns a 200 whose message is the guardrail's blockedInputMessaging.
2547 bedrock_guardrail_response: Final = await self.make_bedrock_api_request(
2548 source="INPUT",
2549 messages=filtered_messages,
2550 request_data=data,
2551 logging_event_type=GuardrailEventHooks.pre_call,
2552 )
2553 #########################################################
2555 #########################################################
2556 ########## 2. Update the messages with the guardrail response ##########
2557 #########################################################
2558 updated_subset: Final = self._update_messages_with_updated_bedrock_guardrail_response(
2559 messages=filtered_messages,
2560 bedrock_guardrail_response=bedrock_guardrail_response,
2561 )
2562 data["messages"] = self._merge_filtered_messages(
2563 original_messages=filter_result.original_messages or new_messages,
2564 updated_target_messages=updated_subset,
2565 target_indices=filter_result.target_indices,
2566 )
2568 #########################################################
2569 ########## 3. Add the guardrail to the applied guardrails header ##########
2570 #########################################################
2571 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
2572 return data
2574 async def async_moderation_hook(
2575 self,
2576 data: dict,
2577 user_api_key_dict: UserAPIKeyAuth,
2578 call_type: CallTypesLiteral,
2579 ):
2580 from litellm.proxy.common_utils.callback_utils import (
2581 add_guardrail_to_applied_guardrails_header,
2582 )
2584 event_type: GuardrailEventHooks = GuardrailEventHooks.during_call
2585 if call_type == CallTypes.call_mcp_tool.value:
2586 event_type = GuardrailEventHooks.during_mcp_call
2587 if self.should_run_guardrail(data=data, event_type=event_type) is not True:
2588 return
2590 new_messages: Final = self.get_guardrails_messages_for_call_type(
2591 call_type=cast(CallTypes, call_type),
2592 data=data,
2593 )
2595 if new_messages is None:
2596 verbose_proxy_logger.warning("Bedrock AI: not running guardrail. No messages in data")
2597 return
2599 filter_result: Final = self._prepare_guardrail_messages_for_role(messages=new_messages)
2600 filtered_messages: Final = filter_result.payload_messages
2601 if not filtered_messages:
2602 verbose_proxy_logger.debug("Bedrock AI: not running guardrail. No user-role messages")
2603 return
2605 #########################################################
2606 ########## 1. Make the Bedrock API request ##########
2607 #########################################################
2608 # A block with disable_exception_on_block=True raises ModifyResponseException
2609 # from make_bedrock_api_request. Because during_call runs in an asyncio.gather
2610 # alongside the LLM call (common_request_processing.py), swallowing the
2611 # exception here to set data["mock_response"] was ineffective: route_request
2612 # unpacked kwargs before this hook ran, and the LLM task's response was taken
2613 # unconditionally. Letting the exception propagate cancels the LLM task and
2614 # the endpoint handler returns the block response.
2615 bedrock_guardrail_response: Final = await self.make_bedrock_api_request(
2616 source="INPUT",
2617 messages=filtered_messages,
2618 request_data=data,
2619 logging_event_type=event_type,
2620 )
2621 #########################################################
2623 #########################################################
2624 ########## 2. Update the messages with the guardrail response ##########
2625 #########################################################
2626 updated_subset: Final = self._update_messages_with_updated_bedrock_guardrail_response(
2627 messages=filtered_messages,
2628 bedrock_guardrail_response=bedrock_guardrail_response,
2629 )
2630 data["messages"] = self._merge_filtered_messages(
2631 original_messages=filter_result.original_messages or new_messages,
2632 updated_target_messages=updated_subset,
2633 target_indices=filter_result.target_indices,
2634 )
2636 #########################################################
2637 ########## 3. Add the guardrail to the applied guardrails header ##########
2638 #########################################################
2639 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
2641 return data
2643 async def async_post_call_success_hook(
2644 self,
2645 data: dict,
2646 user_api_key_dict: UserAPIKeyAuth,
2647 response,
2648 ):
2649 from litellm.proxy.common_utils.callback_utils import (
2650 add_guardrail_to_applied_guardrails_header,
2651 )
2652 from litellm.types.guardrails import GuardrailEventHooks
2654 if self.should_run_guardrail(data=data, event_type=GuardrailEventHooks.post_call) is not True:
2655 return
2657 new_messages: Final[list[AllMessageValues] | None] = data.get("messages")
2658 if new_messages is None:
2659 verbose_proxy_logger.warning("Bedrock AI: not running guardrail. No messages in data")
2660 return
2662 # Check if the ModelResponse has text content in its choices
2663 # to avoid sending empty content to Bedrock (e.g., during tool calls)
2664 if isinstance(response, litellm.ModelResponse):
2665 has_text_content = False
2666 for choice in response.choices:
2667 if isinstance(choice, litellm.Choices):
2668 if choice.message.content and isinstance(choice.message.content, str):
2669 has_text_content = True
2670 break
2672 if not has_text_content:
2673 verbose_proxy_logger.warning("Bedrock AI: not running guardrail. No output text in response")
2674 return
2676 #########################################################
2677 ########## 1. Make Bedrock API requests ##########
2678 #########################################################
2679 # post_call is the response-validation hook by definition — only scan
2680 # OUTPUT. Input scanning belongs to pre_call / during_call hooks, which
2681 # users should configure if they want input validation. Running an
2682 # extra INPUT scan here produced a duplicate post-call entry in the
2683 # trace and made no semantic sense for a "post-call" event.
2684 # A block with disable_exception_on_block=True raises ModifyResponseException
2685 # from make_bedrock_api_request; that propagates to the endpoint handler,
2686 # which returns a 200 whose message is the guardrail's blockedInputMessaging.
2687 # Attach the LLM response to original_response so the synthetic block reply
2688 # reports the real token usage the upstream call consumed instead of zero.
2689 try:
2690 output_content_bedrock: Final = await self.make_bedrock_api_request(
2691 source="OUTPUT",
2692 response=response,
2693 messages=new_messages,
2694 request_data=data,
2695 logging_event_type=GuardrailEventHooks.post_call,
2696 )
2697 except ModifyResponseException as e:
2698 if e.original_response is None:
2699 e.original_response = response
2700 raise
2702 #########################################################
2703 ########## 2. Apply masking to response with output guardrail response ##########
2704 #########################################################
2705 if output_content_bedrock is not None:
2706 self._apply_masking_to_response(
2707 response=response,
2708 bedrock_guardrail_response=output_content_bedrock,
2709 )
2711 #########################################################
2712 ########## 3. Add the guardrail to the applied guardrails header ##########
2713 #########################################################
2714 add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=self.guardrail_name)
2716 ########### HELPER FUNCTIONS for bedrock guardrails ############################
2717 ##############################################################################
2718 ##############################################################################
2719 def _update_messages_with_updated_bedrock_guardrail_response(
2720 self,
2721 messages: list[AllMessageValues],
2722 bedrock_guardrail_response: BedrockGuardrailResponse,
2723 ) -> list[AllMessageValues]:
2724 """
2725 Use the output from the bedrock guardrail to mask sensitive content in messages.
2727 Args:
2728 messages: Original list of messages
2729 bedrock_guardrail_response: Response from Bedrock guardrail containing masked content
2731 Returns:
2732 List of messages with content masked according to guardrail response
2733 """
2734 # Get masked texts from guardrail response
2735 masked_texts: Final = self._extract_masked_texts_from_response(bedrock_guardrail_response)
2737 # If guardrail provided masked output, use it regardless of masking flags
2738 # because the guardrail has already determined this content needs anonymization
2739 if masked_texts:
2740 verbose_proxy_logger.debug("Bedrock guardrail provided masked output, applying to messages")
2741 return self._apply_masking_to_messages(messages=messages, masked_texts=masked_texts)
2743 # If masking is enabled but no masked texts available, still try to apply
2744 # (this maintains backward compatibility for edge cases)
2745 if self.mask_request_content or self.mask_response_content:
2746 verbose_proxy_logger.debug(
2747 "Masking enabled but no masked output from guardrail, returning original messages"
2748 )
2750 return messages
2752 async def async_post_call_streaming_iterator_hook(
2753 self,
2754 user_api_key_dict: UserAPIKeyAuth,
2755 response: Any,
2756 request_data: dict,
2757 ) -> AsyncGenerator[ModelResponseStream, None]:
2758 """
2759 Process streaming response chunks.
2761 Collect content from the stream and run the bedrock OUTPUT scan
2762 (post_call only validates the response).
2763 """
2764 if self._streams_incrementally():
2765 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
2766 UnifiedLLMGuardrails,
2767 )
2769 async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
2770 user_api_key_dict=user_api_key_dict,
2771 response=response,
2772 request_data=request_data,
2773 guardrail_to_apply=self,
2774 buffer_until_moderated_default=False,
2775 ):
2776 yield streamed_chunk
2777 return
2779 # Responses-API events are neither chat-completions chunks nor raw
2780 # Anthropic SSE, so the assembly below cannot scan them; the unified
2781 # guardrail's translation layer can, with buffering semantics kept.
2782 if _is_responses_api_route(user_api_key_dict.request_route):
2783 from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import (
2784 UnifiedLLMGuardrails,
2785 )
2787 async for translated_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook(
2788 user_api_key_dict=user_api_key_dict,
2789 response=response,
2790 request_data=request_data,
2791 guardrail_to_apply=self,
2792 buffer_until_moderated_default=True,
2793 ):
2794 yield translated_chunk
2795 return
2797 # Import here to avoid circular imports
2798 from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
2799 from litellm.main import stream_chunk_builder
2800 from litellm.types.utils import TextCompletionResponse
2802 # Collect all chunks to process them together
2803 started_at: Final = time.monotonic()
2804 all_chunks: Final[list[ModelResponseStream]] = []
2805 async for chunk in response:
2806 all_chunks.append(chunk)
2808 # /v1/messages arrives as SSE frames, which stream_chunk_builder cannot assemble
2809 raw_sse: Final = is_raw_sse_stream(all_chunks)
2810 assembled_model_response: ModelResponse | TextCompletionResponse | None = (
2811 assemble_anthropic_sse_stream(all_chunks, restore_identity=True)
2812 if raw_sse
2813 else stream_chunk_builder(chunks=all_chunks)
2814 )
2815 if isinstance(assembled_model_response, ModelResponse):
2816 pre_guardrail_text: Final = model_response_text(assembled_model_response)
2817 _pre_block_response: Final = assembled_model_response
2818 ####################################################################
2819 ########## 1. Make Bedrock Apply Guardrail API request ##########
2820 #
2821 # post_call only scans OUTPUT — input scanning belongs to
2822 # pre_call / during_call. Bedrock will raise if the response
2823 # violates the guardrail policy.
2824 ###################################################################
2825 # A block with disable_exception_on_block=True raises ModifyResponseException
2826 # from make_bedrock_api_request. Non-streaming paths let it propagate so
2827 # the endpoint handler turns it into a 200. Streaming can't do that: the
2828 # SSE response headers are already flushed, so a raise would be serialized
2829 # as an error frame by async_streaming_data_generator. Instead, replace
2830 # the assembled response with the synthetic block content in-place and
2831 # yield it as a normal stream, matching the shape a non-streaming block
2832 # produces.
2833 try:
2834 output_guardrail_response = await self.make_bedrock_api_request(
2835 source="OUTPUT",
2836 response=assembled_model_response,
2837 messages=request_data.get("messages"),
2838 request_data=request_data,
2839 logging_event_type=GuardrailEventHooks.post_call,
2840 )
2841 except HTTPException as block_exc:
2842 block_detail: Final = block_exc.detail
2843 # A policy block is the only 400 carrying a structured detail; a service failure
2844 # either details a plain string or reports a non-400 status. Re-raising a service
2845 # failure keeps its real status, but only while the headers are unflushed: past the
2846 # first keepalive ping the raise reaches nobody, so it has to travel as a frame too
2847 is_block: Final = raw_sse and block_exc.status_code == 400 and isinstance(block_detail, Mapping)
2848 headers_flushed: Final = keepalive_ping_has_fired(
2849 time.monotonic() - started_at, litellm.anthropic_sse_ping_interval_seconds
2850 )
2851 if not raw_sse or (not is_block and not headers_flushed):
2852 raise
2853 block_message, _ = serialize_http_exception_detail(block_detail)
2854 for error_frame in anthropic_sse_error_frames(
2855 block_message if is_block else f"{block_exc.status_code}: {block_message}"
2856 ):
2857 yield error_frame
2858 return
2859 except ModifyResponseException as e:
2860 if raw_sse:
2861 e.model = _pre_block_response.model or e.model
2862 if e.original_response is None:
2863 e.original_response = _pre_block_response
2864 for block_chunk in AnthropicMessagesHandler().build_block_sse_chunks(e, stream_started=False):
2865 yield block_chunk
2866 return
2867 # Preserve upstream usage from the LLM call we already
2868 # consumed. Non-streaming blocks carry it via
2869 # ModifyResponseException.original_response +
2870 # _blocked_response_usage; streaming has to do the copy
2871 # itself since the exception can't escape this generator.
2872 _original_usage: Final = getattr(assembled_model_response, "usage", None)
2873 assembled_model_response = ModelResponse(
2874 choices=[
2875 Choices(
2876 index=0,
2877 message=Message(role="assistant", content=e.message),
2878 finish_reason="content_filter",
2879 )
2880 ],
2881 model=e.model,
2882 )
2883 if _original_usage is not None:
2884 assembled_model_response.usage = _original_usage
2885 output_guardrail_response = None
2887 #########################################################################
2888 ########## 2. Apply masking to response with output guardrail response ##########
2889 #########################################################################
2890 if output_guardrail_response is not None:
2891 self._apply_masking_to_response(
2892 response=assembled_model_response,
2893 bedrock_guardrail_response=output_guardrail_response,
2894 )
2896 #########################################################################
2897 ########## 3. Return the (potentially masked) chunks ##########
2898 #########################################################################
2899 if raw_sse:
2900 for sse_chunk in (
2901 anthropic_sse_chunks_from_response(assembled_model_response)
2902 if model_response_text(assembled_model_response) != pre_guardrail_text
2903 else all_chunks
2904 ):
2905 yield sse_chunk
2906 return
2908 mock_response: Final = MockResponseIterator(model_response=assembled_model_response)
2910 # Return the reconstructed stream
2911 async for chunk in mock_response:
2912 yield chunk
2913 elif raw_sse:
2914 # Forwarding an unscannable stream would silently disable the guardrail, so fail closed.
2915 # A raise cannot reach the client once a keepalive ping has flushed the headers, so the
2916 # refusal travels as a frame, matching how a block is delivered above
2917 for error_frame in anthropic_sse_error_frames(
2918 f"{self.guardrail_name}: streamed response could not be assembled for scanning, blocking it"
2919 ):
2920 yield error_frame
2921 return
2922 else:
2923 for chunk in all_chunks:
2924 yield chunk
2926 def _extract_masked_texts_from_response(self, bedrock_guardrail_response: BedrockGuardrailResponse) -> list[str]:
2927 """
2928 Extract all masked text outputs from the guardrail response.
2930 Args:
2931 bedrock_guardrail_response: Response from Bedrock guardrail
2933 Returns:
2934 List of masked text strings
2935 """
2936 masked_output_text: Final[list[str]] = []
2937 masked_outputs: Final[list[BedrockGuardrailOutput] | None] = bedrock_guardrail_response.get("outputs", []) or []
2938 if not masked_outputs:
2939 verbose_proxy_logger.debug("No masked outputs found in guardrail response")
2940 return []
2942 for output in masked_outputs:
2943 text_content: str | None = output.get("text")
2944 if text_content is not None:
2945 masked_output_text.append(text_content)
2947 return masked_output_text
2949 def _apply_masking_to_messages(
2950 self, messages: list[AllMessageValues], masked_texts: list[str]
2951 ) -> list[AllMessageValues]:
2952 """
2953 Apply masked texts to message content using index tracking.
2955 Args:
2956 messages: Original messages
2957 masked_texts: List of masked text strings from guardrail
2959 Returns:
2960 Updated messages with masked content
2961 """
2962 updated_messages: Final = []
2963 masking_index = 0
2965 for message in messages:
2966 new_message = message.copy()
2967 content = new_message.get("content")
2969 # Skip messages with no content
2970 if content is None:
2971 updated_messages.append(new_message)
2972 continue
2974 # Handle string content
2975 if isinstance(content, str):
2976 if masking_index < len(masked_texts):
2977 new_message["content"] = masked_texts[masking_index]
2978 masking_index += 1
2979 # Handle list content
2980 elif isinstance(content, list):
2981 new_message["content"], masking_index = self._mask_content_list(
2982 content_list=content,
2983 masked_texts=masked_texts,
2984 masking_index=masking_index,
2985 )
2987 updated_messages.append(new_message)
2989 return updated_messages
2991 def _mask_content_list(
2992 self, content_list: Sequence[object], masked_texts: list[str], masking_index: int
2993 ) -> tuple[list[Any], int]:
2994 """
2995 Apply masking to a list of content items.
2997 Args:
2998 content_list: List of content items
2999 masked_texts: List of masked text strings
3000 starting_index: Starting index in the masked_texts list
3002 Returns:
3003 Updated content list with masked items
3004 """
3005 new_content: Final[list[dict[str, object] | str]] = []
3006 for item in content_list:
3007 if isinstance(item, dict) and "text" in item:
3008 new_item = item.copy()
3009 if masking_index < len(masked_texts):
3010 new_item["text"] = masked_texts[masking_index]
3011 masking_index += 1
3012 new_content.append(new_item)
3013 elif isinstance(item, str):
3014 if masking_index < len(masked_texts):
3015 item = masked_texts[masking_index]
3016 masking_index += 1
3017 if item is not None:
3018 new_content.append(item)
3020 return new_content, masking_index
3022 def _apply_masking_to_response(
3023 self,
3024 response: object,
3025 bedrock_guardrail_response: BedrockGuardrailResponse,
3026 ) -> None:
3027 """
3028 Apply masked content from bedrock guardrail to the response object.
3030 Args:
3031 response: The response object to modify
3032 bedrock_guardrail_response: Response from Bedrock guardrail containing masked content
3033 """
3034 # Get masked texts from guardrail response
3035 masked_texts: Final = self._extract_masked_texts_from_response(bedrock_guardrail_response)
3037 if not masked_texts:
3038 verbose_proxy_logger.debug("No masked outputs found, skipping response masking")
3039 return
3041 verbose_proxy_logger.debug("Applying masking to response with %d masked texts", len(masked_texts))
3043 # Apply masking to ModelResponse
3044 if isinstance(response, litellm.ModelResponse):
3045 self._apply_masking_to_model_response(response, masked_texts)
3046 else:
3047 verbose_proxy_logger.warning("Unsupported response type for masking: %s", type(response))
3049 def _apply_masking_to_model_response(self, response: litellm.ModelResponse, masked_texts: list[str]) -> None:
3050 """
3051 Apply masked texts to a ModelResponse object.
3053 Args:
3054 response: The ModelResponse object to modify in-place
3055 masked_texts: List of masked text strings from guardrail
3056 """
3057 masking_index = 0
3059 for choice in response.choices:
3060 if isinstance(choice, Choices):
3061 # For chat completions
3062 if choice.message.content and isinstance(choice.message.content, str):
3063 if masking_index < len(masked_texts):
3064 choice.message.content = masked_texts[masking_index]
3065 masking_index += 1
3066 verbose_proxy_logger.debug("Applied masking to choice message content")
3067 elif isinstance(choice, StreamingChoices):
3068 # For streaming responses, modify delta content
3069 if choice.delta.content and isinstance(choice.delta.content, str):
3070 if masking_index < len(masked_texts):
3071 choice.delta.content = masked_texts[masking_index]
3072 masking_index += 1
3073 verbose_proxy_logger.debug("Applied masking to choice delta content")
3074 elif isinstance(choice, TextChoices):
3075 # For text completions
3076 if choice.text and isinstance(choice.text, str):
3077 if masking_index < len(masked_texts):
3078 choice.text = masked_texts[masking_index]
3079 masking_index += 1
3080 verbose_proxy_logger.debug("Applied masking to choice text content")
3082 @staticmethod
3083 def _incremental_scan_cache() -> DualCache:
3084 """Resolve the cache used to remember which segments a session already scanned.
3086 Prefers the proxy's shared cache (``internal_usage_cache.dual_cache``), which is
3087 backed by Redis when the deployment configures it, so incremental state is shared
3088 across proxy instances. Falls back to a process-local ``DualCache`` singleton when
3089 the proxy is not running (e.g. unit tests), where sharing does not apply.
3090 """
3091 from litellm.integrations.custom_guardrail import dc as fallback_cache
3093 try:
3094 from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging
3095 except Exception: # noqa: BLE001 # proxy not importable outside the server; use local fallback
3096 return fallback_cache
3097 if _proxy_logging is not None:
3098 return _proxy_logging.internal_usage_cache.dual_cache
3099 return fallback_cache
3101 def _bedrock_response_has_masked_output(self, response: BedrockGuardrailResponse) -> bool:
3102 """Return True if the guardrail rewrote (masked/anonymized) any scanned text.
3104 Bedrock returns non-empty ``output``/``outputs`` text only when it changed the
3105 content; an ``action == "NONE"`` response leaves both empty.
3106 """
3107 for field in ("output", "outputs"):
3108 items = response.get(field) or []
3109 if any(isinstance(item, dict) and item.get("text") for item in items):
3110 return True
3111 return False
3113 async def _apply_incremental_request_scan(
3114 self,
3115 texts: list[str],
3116 inputs: "GenericGuardrailAPIInputs",
3117 request_data: dict,
3118 ) -> Optional["GenericGuardrailAPIInputs"]:
3119 """Scan only the text segments not already seen earlier in this session.
3121 Returns ``None`` when incremental scanning is inactive (feature off, no
3122 session id, masking enabled, or cache unavailable) or when the guardrail
3123 turns out to mask content, telling the caller to run the normal full scan.
3124 Otherwise scans only the new segments and skips the Bedrock call entirely
3125 when nothing is new. Incremental mode is for blocking/detection guardrails
3126 only: if the guardrail returns masked output it cannot be applied to the
3127 skipped context, so the scan falls back to the full path and no session
3128 state is recorded.
3129 """
3130 cache: Final = self._incremental_scan_cache()
3132 new_texts: Final = await self.filter_new_texts_for_session(
3133 texts=texts,
3134 request_data=request_data,
3135 cache=cache,
3136 )
3137 if new_texts is None:
3138 return None
3140 if not new_texts:
3141 verbose_proxy_logger.debug("Bedrock Guardrail: no new messages to scan for this session, skipping API call")
3142 return inputs
3144 bedrock_response: Final = await self.make_bedrock_api_request(
3145 source="INPUT",
3146 messages=[ChatCompletionUserMessage(role="user", content=text) for text in new_texts],
3147 request_data=request_data,
3148 logging_event_type=GuardrailEventHooks.pre_call,
3149 )
3151 if self._bedrock_response_has_masked_output(bedrock_response):
3152 verbose_proxy_logger.warning(
3153 "Bedrock Guardrail %s: guardrail returned masked/anonymized content; "
3154 "only_scan_new_messages cannot apply masking to skipped context, falling back to a full-context scan",
3155 self.guardrail_name,
3156 )
3157 return None
3159 await self.mark_texts_scanned(
3160 texts=texts,
3161 request_data=request_data,
3162 cache=cache,
3163 )
3164 return inputs
3166 async def apply_guardrail(
3167 self,
3168 inputs: "GenericGuardrailAPIInputs",
3169 request_data: dict,
3170 input_type: Literal["request", "response"],
3171 logging_obj: Optional["LiteLLMLoggingObj"] = None,
3172 ) -> "GenericGuardrailAPIInputs":
3173 """
3174 Apply Bedrock guardrail to a batch of texts for testing purposes.
3176 This method allows users to test Bedrock guardrails without making actual LLM calls.
3177 It creates mock messages to test the guardrail functionality.
3179 Args:
3180 inputs: Dictionary containing texts and optional images
3181 request_data: Request data dictionary for logging metadata
3182 input_type: Whether this is a "request" or "response"
3183 logging_obj: Optional logging object
3185 Returns:
3186 GenericGuardrailAPIInputs - processed_texts may be masked, images unchanged
3188 Raises:
3189 Exception: If content is blocked by Bedrock guardrail
3190 """
3191 # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None.
3192 # dict.get("texts", []) would return None if the key exists with a None value.
3193 texts: Final = inputs.get("texts") or []
3194 try:
3195 verbose_proxy_logger.debug("Bedrock Guardrail: Applying guardrail to %s text(s)", len(texts))
3197 if input_type == "request":
3198 incremental_result: Final = await self._apply_incremental_request_scan(
3199 texts=texts,
3200 inputs=inputs,
3201 request_data=request_data,
3202 )
3203 if incremental_result is not None:
3204 return incremental_result
3206 masked_texts = []
3208 selection: Final = self._select_messages_for_apply_guardrail(
3209 texts=texts,
3210 inputs=inputs,
3211 request_data=request_data,
3212 input_type=input_type,
3213 )
3214 if selection.skip_scan:
3215 return inputs
3216 filtered_messages: Final = selection.filtered_messages
3217 scanned_slice: Final = selection.scanned_slice
3218 scanned_role_subset: Final = selection.scanned_role_subset
3220 # Bedrock will throw an error if there is no text to process
3221 if filtered_messages:
3222 _log_hook = GuardrailEventHooks.pre_call if input_type == "request" else GuardrailEventHooks.post_call
3223 # Map the abstract input_type to the Bedrock source parameter.
3224 # "request" -> INPUT (scan user-supplied content)
3225 # "response" -> OUTPUT (scan model-generated content)
3226 # Bedrock guardrail policies are often configured differently
3227 # for Input vs Output (e.g. PII blocking only on Output), so
3228 # the source MUST match where the text originated.
3229 bedrock_source: Final[Literal["INPUT", "OUTPUT"]] = "OUTPUT" if input_type == "response" else "INPUT"
3230 if bedrock_source == "OUTPUT":
3231 # Build a synthetic ModelResponse whose choices carry the
3232 # text(s) to scan, so _create_bedrock_output_content_request
3233 # can produce the correct Bedrock OUTPUT payload.
3234 synthetic_response: Final = ModelResponse(
3235 choices=[
3236 Choices(
3237 index=_idx,
3238 message=Message(
3239 role="assistant",
3240 content=str(_msg.get("content") or ""),
3241 ),
3242 finish_reason="stop",
3243 )
3244 for _idx, _msg in enumerate(filtered_messages)
3245 ]
3246 )
3247 bedrock_response = await self.make_bedrock_api_request(
3248 source="OUTPUT",
3249 response=synthetic_response,
3250 messages=request_data.get("messages"),
3251 request_data=request_data,
3252 logging_event_type=_log_hook,
3253 )
3254 else:
3255 bedrock_response = await self.make_bedrock_api_request(
3256 source="INPUT",
3257 messages=filtered_messages,
3258 request_data=request_data,
3259 logging_event_type=_log_hook,
3260 )
3262 # Apply any masking that was applied by the guardrail
3263 output_list: Final = bedrock_response.get("output")
3264 if output_list:
3265 # If the guardrail returned modified content, use that
3266 for output_item in output_list:
3267 text_content = output_item.get("text")
3268 if text_content:
3269 masked_text = str(text_content)
3270 masked_texts.append(masked_text)
3271 else:
3272 outputs_list: Final = bedrock_response.get("outputs")
3273 if outputs_list:
3274 # Fallback to outputs field if output is not available
3275 for output_item in outputs_list:
3276 text_content = output_item.get("text")
3277 if text_content:
3278 masked_text = str(text_content)
3279 masked_texts.append(masked_text)
3281 # Reconcile masked output with the flat `texts` list (write back to
3282 # the scanned slice only, or skip if it can't be aligned).
3283 masked_texts = self._merge_masked_texts(
3284 masked_texts=masked_texts,
3285 texts=texts,
3286 scanned_slice=scanned_slice,
3287 scanned_role_subset=scanned_role_subset,
3288 )
3290 verbose_proxy_logger.debug("Bedrock Guardrail: Successfully applied guardrail")
3292 inputs["texts"] = masked_texts
3293 return inputs
3295 except (HTTPException, ModifyResponseException):
3296 # Let guardrail blocking exceptions propagate as-is so the proxy can
3297 # return the correct HTTP status (400 for HTTPException, 200 with the
3298 # block message for ModifyResponseException in disable_exception_on_block
3299 # mode). Without this, the generic except below wraps them into a plain
3300 # Exception, losing the semantics and preventing the proxy from
3301 # properly blocking the call.
3302 raise
3303 except Exception as e:
3304 verbose_proxy_logger.error("Bedrock Guardrail: Failed to apply guardrail: %s", str(e))
3305 raise Exception(f"Bedrock guardrail failed: {e}")