From 3a66dcc5a40bd2f042e086926b33951d0b97c2d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:08:21 +0000 Subject: [PATCH 1/2] fix(bedrock): map guardrailConfig to InvokeModel guardrail headers The InvokeModel API takes the guardrail identifier, version and trace as X-Amzn-Bedrock-* request headers, unlike Converse which takes them in the request body. The invoke transformer never set these headers, so guardrailConfig was silently dropped (or leaked into the request body) and Bedrock guardrails never ran on invoke-route models. Pop guardrailConfig in AmazonInvokeConfig.validate_environment, validate it, and set the headers before SigV4 signing; explicitly passed headers keep winning over guardrailConfig so existing workarounds are unaffected --- .../base_invoke_transformation.py | 37 ++++- litellm/types/llms/bedrock.py | 2 +- .../test_base_invoke_transformation.py | 141 ++++++++++++++++++ 3 files changed, 178 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index bbe16e26713..b22d49a3713 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union, cast, get_args import httpx +from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger @@ -24,6 +25,7 @@ HTTPHandler, _get_httpx_client, ) +from litellm.types.llms.bedrock import GuardrailConfigBlock from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import ModelResponse, Usage from litellm.utils import CustomStreamWrapper @@ -37,6 +39,30 @@ from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM +_GUARDRAIL_CONFIG_VALIDATOR: "TypeAdapter[GuardrailConfigBlock]" = TypeAdapter(GuardrailConfigBlock) + + +def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": + try: + guardrail_config = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config) + except ValidationError as e: + raise BedrockError( + status_code=400, + message=( + "Invalid guardrailConfig={}. Expected format: {{'guardrailIdentifier': str, " + "'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}}. Error: {}".format( + raw_guardrail_config, e + ) + ), + ) + trace = guardrail_config.get("trace") + candidate_headers = { + "X-Amzn-Bedrock-GuardrailIdentifier": guardrail_config.get("guardrailIdentifier"), + "X-Amzn-Bedrock-GuardrailVersion": guardrail_config.get("guardrailVersion"), + "X-Amzn-Bedrock-Trace": trace.upper() if trace is not None else None, + } + return {name: value for name, value in candidate_headers.items() if value is not None} + class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): def __init__(self, **kwargs): @@ -390,7 +416,16 @@ def validate_environment( api_key: Optional[str] = None, api_base: Optional[str] = None, ) -> dict: - return headers + raw_guardrail_config = optional_params.pop("guardrailConfig", None) + if raw_guardrail_config is None: + return headers + existing_header_names = frozenset(name.lower() for name in headers) + guardrail_headers = { + name: value + for name, value in _bedrock_invoke_guardrail_headers(raw_guardrail_config).items() + if name.lower() not in existing_header_names + } + return {**headers, **guardrail_headers} def get_error_class( self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] diff --git a/litellm/types/llms/bedrock.py b/litellm/types/llms/bedrock.py index a1dc08e5f29..bdf6b8fefed 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -325,7 +325,7 @@ class ToolConfigBlock(TypedDict, total=False): class GuardrailConfigBlock(TypedDict, total=False): guardrailIdentifier: str guardrailVersion: str - trace: Literal["enabled", "disabled"] + trace: Literal["enabled", "disabled", "enabled_full"] class InferenceConfig(TypedDict, total=False): diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index aff89f02ff2..4b39b13d3ac 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -14,6 +14,7 @@ from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( AmazonInvokeConfig, ) +from litellm.llms.bedrock.common_utils import BedrockError @pytest.mark.parametrize( @@ -39,3 +40,143 @@ def test_transform_request_drops_stream_chunk_size(config, model): ) assert "stream_chunk_size" not in json.dumps(request_body) + + +def test_validate_environment_maps_guardrail_config_to_invoke_headers(): + """The InvokeModel API takes the guardrail identifier/version/trace as + X-Amzn-Bedrock-* request headers, unlike Converse which takes them in the + body. https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_InvokeModel.html""" + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "DRAFT", + "trace": "enabled", + }, + "max_tokens": 10, + } + + headers = AmazonInvokeConfig().validate_environment( + headers={}, + model="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params=optional_params, + litellm_params={}, + ) + + assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "ff6ujrregl1q" + assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT" + assert headers["X-Amzn-Bedrock-Trace"] == "ENABLED" + assert "guardrailConfig" not in optional_params + + +def test_validate_environment_without_guardrail_config_leaves_headers_untouched(): + headers = AmazonInvokeConfig().validate_environment( + headers={"foo": "bar"}, + model="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params={"max_tokens": 10}, + litellm_params={}, + ) + + assert headers == {"foo": "bar"} + + +def test_validate_environment_skips_absent_guardrail_fields(): + headers = AmazonInvokeConfig().validate_environment( + headers={}, + model="amazon.titan-text-express-v1", + messages=[{"role": "user", "content": "hi"}], + optional_params={"guardrailConfig": {"guardrailIdentifier": "gr-id", "guardrailVersion": "1"}}, + litellm_params={}, + ) + + assert headers == { + "X-Amzn-Bedrock-GuardrailIdentifier": "gr-id", + "X-Amzn-Bedrock-GuardrailVersion": "1", + } + + +def test_validate_environment_does_not_clobber_explicit_guardrail_headers(): + """Users worked around the missing guardrailConfig support by passing the + AWS headers directly; an explicit header must keep winning over + guardrailConfig regardless of casing.""" + headers = AmazonInvokeConfig().validate_environment( + headers={"x-amzn-bedrock-guardrailidentifier": "explicit-id"}, + model="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "guardrailConfig": {"guardrailIdentifier": "config-id", "guardrailVersion": "2"}, + }, + litellm_params={}, + ) + + assert headers["x-amzn-bedrock-guardrailidentifier"] == "explicit-id" + assert "X-Amzn-Bedrock-GuardrailIdentifier" not in headers + assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "2" + + +@pytest.mark.parametrize( + "bad_guardrail_config", + [ + {"guardrailIdentifier": "gr-id", "trace": "verbose"}, + {"guardrailIdentifier": ["gr-id"]}, + "gr-id", + ], +) +def test_validate_environment_rejects_malformed_guardrail_config(bad_guardrail_config): + with pytest.raises(BedrockError) as excinfo: + AmazonInvokeConfig().validate_environment( + headers={}, + model="anthropic.claude-3-sonnet-20240229-v1:0", + messages=[{"role": "user", "content": "hi"}], + optional_params={"guardrailConfig": bad_guardrail_config}, + litellm_params={}, + ) + + assert excinfo.value.status_code == 400 + assert "guardrailConfig" in str(excinfo.value) + + +@pytest.mark.parametrize( + "model", + [ + "anthropic.claude-3-sonnet-20240229-v1:0", + "amazon.titan-text-express-v1", + "mistral.mistral-7b-instruct-v0:2", + "meta.llama3-8b-instruct-v1:0", + ], +) +def test_guardrail_config_flows_to_headers_not_request_body(model): + """Mirrors the handler flow (validate_environment then transform_request): + guardrailConfig must end up in the signed headers and never leak into the + request body, where Bedrock rejects it as an extra input.""" + config = AmazonInvokeConfig() + optional_params = { + "guardrailConfig": { + "guardrailIdentifier": "ff6ujrregl1q", + "guardrailVersion": "DRAFT", + "trace": "disabled", + }, + "max_tokens": 10, + } + messages = [{"role": "user", "content": "hi"}] + + headers = config.validate_environment( + headers={}, + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + ) + request_body = config.transform_request( + model=model, + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers=headers, + ) + + assert "guardrailConfig" not in json.dumps(request_body) + assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "ff6ujrregl1q" + assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT" + assert headers["X-Amzn-Bedrock-Trace"] == "DISABLED" From 6e9a46fd2267e7ba66b51ae5a5b46c62cbdc9826 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 2 Jul 2026 19:24:11 +0000 Subject: [PATCH 2/2] fix(bedrock): reject guardrailConfig missing guardrailIdentifier A guardrailConfig without guardrailIdentifier (e.g. an empty dict) would validate, produce no guardrail headers, and let the request proceed with guardrails silently not applied; that silent skip is the exact failure mode this fix exists to remove, so fail fast with a 400 instead --- .../base_invoke_transformation.py | 18 +++++++++++++----- .../test_base_invoke_transformation.py | 2 ++ 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py index b22d49a3713..dd7cf12604d 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -41,6 +41,10 @@ _GUARDRAIL_CONFIG_VALIDATOR: "TypeAdapter[GuardrailConfigBlock]" = TypeAdapter(GuardrailConfigBlock) +_GUARDRAIL_CONFIG_EXPECTED_FORMAT = ( + "{'guardrailIdentifier': str, 'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}" +) + def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": try: @@ -48,11 +52,15 @@ def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str except ValidationError as e: raise BedrockError( status_code=400, - message=( - "Invalid guardrailConfig={}. Expected format: {{'guardrailIdentifier': str, " - "'guardrailVersion': str, 'trace': 'enabled'|'disabled'|'enabled_full'}}. Error: {}".format( - raw_guardrail_config, e - ) + message="Invalid guardrailConfig={}. Expected format: {}. Error: {}".format( + raw_guardrail_config, _GUARDRAIL_CONFIG_EXPECTED_FORMAT, e + ), + ) + if "guardrailIdentifier" not in guardrail_config: + raise BedrockError( + status_code=400, + message="guardrailConfig={} is missing 'guardrailIdentifier'. Expected format: {}".format( + raw_guardrail_config, _GUARDRAIL_CONFIG_EXPECTED_FORMAT ), ) trace = guardrail_config.get("trace") diff --git a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index 4b39b13d3ac..5fefae7e411 100644 --- a/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -121,6 +121,8 @@ def test_validate_environment_does_not_clobber_explicit_guardrail_headers(): {"guardrailIdentifier": "gr-id", "trace": "verbose"}, {"guardrailIdentifier": ["gr-id"]}, "gr-id", + {}, + {"trace": "enabled"}, ], ) def test_validate_environment_rejects_malformed_guardrail_config(bad_guardrail_config):