fix(bedrock): map guardrailConfig to InvokeModel guardrail headers

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,38 @@
 
 from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
 
+_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:
+        guardrail_config = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config)
+    except ValidationError as e:
+        raise BedrockError(
+            status_code=400,
+            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")
+    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 +424,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]

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):

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,145 @@ 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",
+        {},
+        {"trace": "enabled"},
+    ],
+)
+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"