From 26e2e208ad347542f7ec4a5c34a2c4a02a5c9145 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 21:18:08 +0000 Subject: [PATCH 1/4] fix(bedrock guardrails): derive contextual grounding source and query from plain messages Bedrock only runs a contextualGroundingPolicy when the ApplyGuardrail payload carries grounding_source and query qualifiers. Callers sending ordinary system and user messages never got those, so a configured grounding threshold was silently skipped on /v1/chat/completions and /guardrails/apply_guardrail. When no explicit grounding_source or query tags are present, system and developer text is sent as grounding_source and the latest user message as query. The apply_guardrail response branch now forwards the request messages, which it previously dropped. Resolves LIT-4224 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/bedrock_guardrails.py | 32 +++- .../test_bedrock_guardrails.py | 154 ++++++++++++++++++ 2 files changed, 185 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 6a7ac4361b9..fe254803d04 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -493,6 +493,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): result carrying externally-influenced content can supply fake evidence for the 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. """ grounding: Final[list[QualifiedTextBlock]] = [] for message in messages or []: @@ -504,7 +507,33 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): and role in _GROUNDING_SOURCE_TRUSTED_ROLES ): grounding.append(block) - return grounding + return grounding or 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") + if latest_user_index is None: + return [] + sources: Final = tuple( + QualifiedTextBlock(text=block.text, qualifier="grounding_source") + for message in messages + if message.get("role") in _GROUNDING_SOURCE_TRUSTED_ROLES + for block in self.get_content_items_for_message(message=message) or [] + if block.text + ) + queries: Final = tuple( + QualifiedTextBlock(text=block.text, qualifier="query") + for block in self.get_content_items_for_message(message=messages[latest_user_index]) or [] + if block.text + ) + if not sources or not queries: + return [] + return [*sources, *queries] def supports_scan_only_tool_results(self) -> bool: return self.experimental_use_latest_role_message_only is not True @@ -3210,6 +3239,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): bedrock_response = await self.make_bedrock_api_request( source="OUTPUT", response=synthetic_response, + messages=request_data.get("messages"), request_data=request_data, logging_event_type=_log_hook, ) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index c0762edec92..a9a208e1807 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2474,6 +2474,122 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): assert actual_request == expected_request +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.""" + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "OUTPUT", + "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], + } + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +def test_grounding_output_derived_query_is_latest_user_turn_only(): + """In a multi-turn chat only the latest user message is the query; earlier user + turns and assistant turns are not sent as query or source. Developer messages + count as source alongside system.""" + developer_text = "Answer in one sentence." + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "developer", "content": [{"type": "text", "text": developer_text}]}, + {"role": "user", "content": "Hi"}, + {"role": "assistant", "content": "Hello, how can I help?"}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "OUTPUT", + "content": [ + _GROUNDING_SOURCE_BLOCK, + {"text": {"text": developer_text, "qualifiers": ["grounding_source"]}}, + _QUERY_BLOCK, + _GUARD_BLOCK, + ], + } + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +@pytest.mark.parametrize( + "messages", + [ + pytest.param([{"role": "system", "content": _GROUNDING_SOURCE_TEXT}], id="system-without-user"), + pytest.param( + [ + {"role": "tool", "content": _GROUNDING_SOURCE_TEXT, "tool_call_id": "c1"}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ], + id="tool-result-is-not-a-source", + ), + pytest.param( + [ + {"role": "system", "content": ""}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ], + id="empty-system-prompt", + ), + pytest.param( + [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": [{"type": "image_url", "image_url": {"url": "https://x.test/a.png"}}]}, + ], + id="image-only-user-turn", + ), + ], +) +def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(messages): + """Bedrock rejects a grounding_source without a query (and vice versa), so a + 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)) + + assert actual_request == expected_request + + +def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): + """A caller that tags blocks keeps full control: untagged system text is not + added as a second source and the untagged user text is not a second query.""" + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + *_grounding_messages(), + {"role": "user", "content": "Please be brief."}, + ] + expected_request = { + "source": "OUTPUT", + "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], + } + + actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT)) + + assert actual_request == expected_request + + +def test_grounding_input_ignores_plain_message_derivation(): + """Derivation is OUTPUT-only: an INPUT scan of plain system + user text stays an + untagged payload, so input policies keep scanning every block.""" + messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "INPUT", + "content": [{"text": {"text": _GROUNDING_SOURCE_TEXT}}, {"text": {"text": _GROUNDING_QUERY_TEXT}}], + } + + actual_request = _input_request(messages) + + assert actual_request == expected_request + + def test_grounding_output_combines_multiple_sources(): """Every grounding_source block is emitted; Bedrock combines them into one corpus.""" uk_source_text = "London is the capital of UK." @@ -2597,6 +2713,44 @@ async def test_grounding_output_blocked_raises_400(): assert exc_info.value.status_code == 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() + request_messages = [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ] + expected_request = { + "source": "OUTPUT", + "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK], + } + + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + with ( + patch.object(guardrail.async_handler, "post", new_callable=AsyncMock) as mock_post, + patch.object(guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1")), + patch.object(guardrail, "_prepare_request", return_value=MagicMock()) as mock_prepare, + ): + mock_post.return_value = _passing_bedrock_httpx_response(_GROUNDING_RESPONSE_TEXT) + + await guardrail.apply_guardrail( + inputs={"texts": [_GROUNDING_RESPONSE_TEXT]}, + request_data={"messages": request_messages}, + input_type="response", + ) + + assert mock_prepare.call_count == 1 + assert json.loads(json.dumps(mock_prepare.call_args.kwargs["data"])) == expected_request + + ############################################################################### # LIT-4186: disable_exception_on_block regression tests # From 442af3aab69f39d490829714025ffc2e49e7acc3 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:04:42 +0000 Subject: [PATCH 2/4] 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 | 6 +++ .../guardrail_hooks/bedrock_guardrails.py | 17 ++++--- .../guardrails/guardrail_initializers.py | 1 + litellm/types/guardrails.py | 10 ++++ .../test_bedrock_guardrails.py | 48 +++++++++++++------ .../proxy/guardrails/test_init_guardrails.py | 46 ++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 +++ 7 files changed, 112 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 7a110eff080..fb2f014e3d8 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/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": [ { diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index fe254803d04..2c407d91a48 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/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") diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 16369abbfb0..7858adeb55d 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/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, diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 69cb88bfa2f..ab2f9edea1d 100644 --- a/litellm/types/guardrails.py +++ b/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): diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index a9a208e1807..d14ff34e0f4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -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}, diff --git a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py b/tests/test_litellm/proxy/guardrails/test_init_guardrails.py index ceb084b4a4d..8377db57b6e 100644 --- a/tests/test_litellm/proxy/guardrails/test_init_guardrails.py +++ b/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 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0b0e3e18215..d76b256f752 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/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 From 2e11f7bc7b8eef524418dcf00305b2d80a556f71 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:09:59 +0000 Subject: [PATCH 3/4] test(bedrock guardrails): shorten grounding test docstrings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_hooks/test_bedrock_guardrails.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index d14ff34e0f4..81bbaf86e46 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2481,9 +2481,7 @@ def test_grounding_output_keeps_legacy_payload_without_tags(): def test_grounding_output_derives_source_and_query_from_plain_messages(): - """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.""" + """Flag on: untagged system + user text is sent as grounding_source + query.""" messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, {"role": "user", "content": _GROUNDING_QUERY_TEXT}, @@ -2512,9 +2510,7 @@ def test_grounding_output_plain_messages_stay_legacy_when_flag_is_off(): def test_grounding_output_derived_query_is_latest_user_turn_only(): - """In a multi-turn chat only the latest user message is the query; earlier user - turns and assistant turns are not sent as query or source. Developer messages - count as source alongside system.""" + """Only the latest user turn is the query; system and developer turns are the source.""" developer_text = "Answer in one sentence." messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, @@ -2566,8 +2562,7 @@ def test_grounding_output_derived_query_is_latest_user_turn_only(): ], ) def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(messages): - """Bedrock rejects a grounding_source without a query (and vice versa), so a - request that cannot supply both from trusted roles keeps the untagged payload.""" + """Bedrock rejects a source without a query and vice versa, so send neither.""" expected_request = {"source": "OUTPUT", "content": [{"text": {"text": _GROUNDING_RESPONSE_TEXT}}]} actual_request = _output_request(messages, _model_response(_GROUNDING_RESPONSE_TEXT), from_messages=True) @@ -2576,8 +2571,7 @@ def test_grounding_output_stays_legacy_when_plain_source_or_query_is_missing(mes def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): - """A caller that tags blocks keeps full control: untagged system text is not - added as a second source and the untagged user text is not a second query.""" + """Tagged blocks win: untagged text around them is not added as source or query.""" messages = [ {"role": "system", "content": "You are a helpful assistant."}, *_grounding_messages(), @@ -2594,8 +2588,7 @@ def test_grounding_output_explicit_tags_take_precedence_over_plain_messages(): def test_grounding_input_ignores_plain_message_derivation(): - """Derivation is OUTPUT-only: an INPUT scan of plain system + user text stays an - untagged payload, so input policies keep scanning every block.""" + """INPUT scans never derive grounding qualifiers from plain messages.""" messages = [ {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, {"role": "user", "content": _GROUNDING_QUERY_TEXT}, From d8d7210b5734c521a1d0fd0f651b620785492515 Mon Sep 17 00:00:00 2001 From: yucheng Date: Mon, 14 Sep 2026 22:28:58 +0000 Subject: [PATCH 4/4] test(bedrock guardrails): cover tagged messages on the apply_guardrail response path with the flag off Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../test_bedrock_guardrails.py | 28 +++++++++++++++---- 1 file changed, 22 insertions(+), 6 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 81bbaf86e46..d173c5f5c70 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -2729,12 +2729,28 @@ async def test_grounding_output_blocked_raises_400(): @pytest.mark.asyncio -async def test_apply_guardrail_response_forwards_request_messages_for_grounding(): - guardrail = _grounding_guardrail(from_messages=True) - request_messages = [ - {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, - {"role": "user", "content": _GROUNDING_QUERY_TEXT}, - ] +@pytest.mark.parametrize( + "from_messages, request_messages", + [ + ( + True, + [ + {"role": "system", "content": _GROUNDING_SOURCE_TEXT}, + {"role": "user", "content": _GROUNDING_QUERY_TEXT}, + ], + ), + ( + False, + [ + {"role": "system", "content": [{"type": "grounding_source", "text": _GROUNDING_SOURCE_TEXT}]}, + {"role": "user", "content": [{"type": "query", "text": _GROUNDING_QUERY_TEXT}]}, + ], + ), + ], + ids=["plain-messages-flag-on", "tagged-messages-flag-off"], +) +async def test_apply_guardrail_response_forwards_request_messages_for_grounding(from_messages, request_messages): + guardrail = _grounding_guardrail(from_messages=from_messages) expected_request = { "source": "OUTPUT", "content": [_GROUNDING_SOURCE_BLOCK, _QUERY_BLOCK, _GUARD_BLOCK],