fix(bedrock guardrails): gate plain-message grounding behind contextual_grounding_from_messages

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/_lazy_openapi_snapshot.json

@@ -11503,6 +11503,12 @@
               "description": "Enable content moderation to check for harmful content (harassment, hate speech, etc.).",
               "title": "Content Moderation Check"
             },
+            "contextual_grounding_from_messages": {
+              "default": false,
+              "description": "ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.",
+              "title": "Contextual Grounding From Messages",
+              "type": "boolean"
+            },
             "credentials": {
               "anyOf": [
                 {

litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py

@@ -244,6 +244,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         prompt_attack_threshold: float | None = 0.5,
         pii_confidence_threshold: float | None = 0.5,
         chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
+        contextual_grounding_from_messages: bool = False,
         streaming_buffer_until_moderated: bool | None = None,
         streaming_sampling_rate: int | None = None,
         streaming_end_of_stream_only: bool | None = None,
@@ -265,6 +266,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         self.guardrailVersion = guardrailVersion
         self.guardrail_provider = "bedrock"
         self.chunk_budget_chars = chunk_budget_chars
+        self.contextual_grounding_from_messages = contextual_grounding_from_messages
         self.experimental_use_latest_role_message_only = bool(kwargs.get("experimental_use_latest_role_message_only"))
 
         # Resource-less, detect-only InvokeGuardrailChecks mode. Present `checks`
@@ -459,8 +461,8 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         """
         Flatten a message into text blocks, preserving any contextual-grounding
         qualifier carried by the content-block ``type`` (grounding_source / query).
-        Untagged text keeps ``qualifier=None`` so the payload is unchanged for
-        callers that do not use grounding.
+        Untagged text keeps ``qualifier=None``; the OUTPUT scan decides whether to
+        derive grounding qualifiers from it.
         """
         content: Final = message.get("content")
         if content is None:
@@ -494,8 +496,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
         contextual-grounding check to grade the response against. ``query`` is accepted
         from any role (it is the user's question).
 
-        A request with no tagged blocks falls back to the plain messages: system /
-        developer text is the grounding source and the latest user message is the query.
+        With ``contextual_grounding_from_messages`` on, a request with no tagged blocks
+        falls back to the plain messages: system / developer text is the grounding
+        source and the latest user message is the query.
         """
         grounding: Final[list[QualifiedTextBlock]] = []
         for message in messages or []:
@@ -507,13 +510,13 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
                     and role in _GROUNDING_SOURCE_TRUSTED_ROLES
                 ):
                     grounding.append(block)
-        return grounding or self._derive_grounding_blocks_from_plain_messages(messages)
+        if grounding or not self.contextual_grounding_from_messages:
+            return grounding
+        return self._derive_grounding_blocks_from_plain_messages(messages)
 
     def _derive_grounding_blocks_from_plain_messages(
         self, messages: list[AllMessageValues] | None
     ) -> list[QualifiedTextBlock]:
-        """Bedrock scores grounding only when source, query and response are all present,
-        and rejects a source without a query, so return nothing unless both exist."""
         if not messages:
             return []
         latest_user_index: Final = self._find_latest_message_index(messages, target_role="user")

litellm/proxy/guardrails/guardrail_initializers.py

@@ -23,6 +23,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail):
         prompt_attack_threshold=litellm_params.prompt_attack_threshold,
         pii_confidence_threshold=litellm_params.pii_confidence_threshold,
         chunk_budget_chars=litellm_params.chunk_budget_chars,
+        contextual_grounding_from_messages=litellm_params.contextual_grounding_from_messages,
         default_on=litellm_params.default_on,
         disable_exception_on_block=litellm_params.disable_exception_on_block,
         mask_request_content=litellm_params.mask_request_content,

litellm/types/guardrails.py

@@ -552,6 +552,16 @@ class BedrockGuardrailConfigModel(BaseModel):
         "still rejects is bisected automatically, so this value only trades round trips against "
         "batch size and cannot fail a request on its own.",
     )
+    contextual_grounding_from_messages: bool = Field(
+        default=False,
+        description="ApplyGuardrail: when True, post-call scans of a request with no grounding_source / "
+        "query content parts send the system and developer messages as the grounding source and "
+        "the latest user message as the query, so the guardrail's contextual grounding policy can "
+        "score the response. Bedrock bills contextual grounding units for these scans and rejects "
+        "queries, sources and responses over its contextual grounding length limits, so leave this "
+        "off for guardrails without a contextual grounding policy. Default False: plain messages "
+        "are never sent as grounding context.",
+    )
 
 
 class BedrockGuardrailStreamingParams(BaseModel):

tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py

@@ -2375,8 +2375,12 @@ _GROUNDING_QUERY_TEXT = "What is the capital of Japan?"
 _GROUNDING_RESPONSE_TEXT = "The capital of Japan is Tokyo."
 
 
-def _grounding_guardrail() -> BedrockGuardrail:
-    return BedrockGuardrail(guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT")
+def _grounding_guardrail(from_messages: bool = False) -> BedrockGuardrail:
+    return BedrockGuardrail(
+        guardrailIdentifier="test-guardrail",
+        guardrailVersion="DRAFT",
+        contextual_grounding_from_messages=from_messages,
+    )
 
 
 def _grounding_messages() -> list:
@@ -2418,9 +2422,11 @@ def _input_request(messages: list) -> dict:
     return _grounding_guardrail().convert_to_bedrock_format(source="INPUT", messages=messages)
 
 
-def _output_request(messages: list, response=None) -> dict:
+def _output_request(messages: list, response=None, from_messages: bool = False) -> dict:
     """Arrange a guardrail and act: build the Bedrock OUTPUT payload."""
-    return _grounding_guardrail().convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages)
+    return _grounding_guardrail(from_messages).convert_to_bedrock_format(
+        source="OUTPUT", response=response, messages=messages
+    )
 
 
 def test_grounding_input_strips_grounding_and_query_qualifiers():
@@ -2475,8 +2481,9 @@ def test_grounding_output_keeps_legacy_payload_without_tags():
 
 
 def test_grounding_output_derives_source_and_query_from_plain_messages():
-    """Untagged string system + user messages become grounding_source + query, so a
-    guardrail with a grounding threshold actually grades the response."""
+    """With contextual_grounding_from_messages on, untagged string system + user
+    messages become grounding_source + query, so a guardrail with a grounding
+    threshold grades the response."""
     messages = [
         {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
         {"role": "user", "content": _GROUNDING_QUERY_TEXT},
@@ -2486,6 +2493,19 @@ def test_grounding_output_derives_source_and_query_from_plain_messages():
         "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK],
     }
 
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
+
+    assert actual_request == expected_request
+
+
+def test_grounding_output_plain_messages_stay_legacy_when_flag_is_off():
+    """Default config: plain system + user text is never sent as grounding context."""
+    messages = [
+        {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
+        {"role": "user", "content": _GROUNDING_QUERY_TEXT},
+    ]
+    expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]}
+
     actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT))
 
     assert actual_request == expected_request
@@ -2513,7 +2533,7 @@ def test_grounding_output_derived_query_is_latest_user_turn_only():
         ],
     }
 
-    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT))
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
 
     assert actual_request == expected_request
 
@@ -2550,7 +2570,7 @@ def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(mes
     request that cannot supply both from trusted roles keeps the untagged payload."""
     expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]}
 
-    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT))
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
 
     assert actual_request == expected_request
 
@@ -2568,7 +2588,7 @@ def test_grounding_output_explicit_tags_take_precedence_over_plain_messages():
         "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK],
     }
 
-    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT))
+    actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True)
 
     assert actual_request == expected_request
 
@@ -2585,7 +2605,9 @@ def test_grounding_input_ignores_plain_message_derivation():
         "content": [{"text": {"text": _GROUNDING_SOURCE_TEXT}}, {"text": {"text": _GROUNDING_QUERY_TEXT}}],
     }
 
-    actual_request = _input_request(messages)
+    actual_request = _grounding_guardrail(from_messages=True).convert_to_bedrock_format(
+        source="INPUT", messages=messages
+    )
 
     assert actual_request == expected_request
 
@@ -2715,11 +2737,7 @@ async def test_grounding_output_blocked_raises_400():
 
 @pytest.mark.asyncio
 async def test_apply_guardrail_response_forwards_request_messages_for_grounding():
-    """/guardrails/apply_guardrail with input_type=response: the request messages
-    stored in request_data must reach the OUTPUT payload as grounding_source + query
-    around the guarded text. Before LIT-4224 the response branch dropped them, so
-    Bedrock never ran its contextual-grounding policy on this route."""
-    guardrail = _grounding_guardrail()
+    guardrail = _grounding_guardrail(from_messages=True)
     request_messages = [
         {"role": "system", "content": _GROUNDING_SOURCE_TEXT},
         {"role": "user", "content": _GROUNDING_QUERY_TEXT},

tests/test_litellm/proxy/guardrails/test_init_guardrails.py

@@ -71,6 +71,52 @@ def test_initialize_bedrock_forwards_chunk_budget_chars():
     assert initialized[-1].chunk_budget_chars == 60_000
 
 
+def test_initialize_bedrock_forwards_contextual_grounding_from_messages():
+    """`contextual_grounding_from_messages: true` in config.yaml must make the post-call
+    payload carry the plain system prompt and user turn as grounding_source and query."""
+    import litellm
+    from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail
+    from litellm.types.utils import Choices, Message, ModelResponse
+
+    test_guardrail = {
+        "guardrail_name": "test_bedrock_grounding_from_messages",
+        "litellm_params": {
+            "guardrail": SupportedGuardrailIntegrations.BEDROCK.value,
+            "mode": "post_call",
+            "guardrailIdentifier": "test-guardrail",
+            "guardrailVersion": "DRAFT",
+            "contextual_grounding_from_messages": True,
+        },
+    }
+    messages = [
+        {"role": "system", "content": "Returns are accepted for 30 days."},
+        {"role": "user", "content": "How long is the return window?"},
+    ]
+    response = ModelResponse(
+        choices=[Choices(index=0, message=Message(role="assistant", content="30 days."), finish_reason="stop")]
+    )
+    expected_request = {
+        "source": "OUTPUT",
+        "content": [
+            {"text": {"text": "Returns are accepted for 30 days.", "qualifiers": ["grounding_source"]}},
+            {"text": {"text": "How long is the return window?", "qualifiers": ["query"]}},
+            {"text": {"text": "30 days.", "qualifiers": ["guard_content"]}},
+        ],
+    }
+
+    guardrail_handler = InMemoryGuardrailHandler()
+    guardrail_handler.initialize_guardrail(guardrail=test_guardrail)
+
+    initialized = [
+        callback
+        for callback in litellm.callbacks
+        if isinstance(callback, BedrockGuardrail) and callback.guardrail_name == "test_bedrock_grounding_from_messages"
+    ]
+    assert initialized, "bedrock guardrail was not registered as a callback"
+    actual_request = initialized[-1].convert_to_bedrock_format(source="OUTPUT", response=response, messages=messages)
+    assert json.loads(json.dumps(actual_request)) == expected_request
+
+
 def test_initialize_guardrail_preserves_guardrail_info():
     """
     Regression (LIT-2529): initialize_guardrail must carry guardrail_info into the

ui/litellm-dashboard/src/lib/http/schema.d.ts

@@ -30875,6 +30875,12 @@ export interface components {
              * @description Enable content moderation to check for harmful content (harassment, hate speech, etc.).
              */
             content_moderation_check?: boolean | null;
+            /**
+             * Contextual Grounding From Messages
+             * @description ApplyGuardrail: when True, post-call scans of a request with no grounding_source / query content parts send the system and developer messages as the grounding source and the latest user message as the query, so the guardrail's contextual grounding policy can score the response. Bedrock bills contextual grounding units for these scans and rejects queries, sources and responses over its contextual grounding length limits, so leave this off for guardrails without a contextual grounding policy. Default False: plain messages are never sent as grounding context.
+             * @default false
+             */
+            contextual_grounding_from_messages: boolean;
             /**
              * Credentials
              * @description Path to Google Cloud credentials JSON file or JSON string