From e6065b9f4f74ec82c09f93398bc74bdd03dcf083 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Tue, 20 Oct 2026 14:48:19 -0700 Subject: [PATCH] fix(bedrock)!: apply the deployment guardrailConfig on /bedrock pass-through routes (#47429) * fix(bedrock): apply the deployment guardrailConfig on the /bedrock pass-through routes The router-model pass-through dropped a Bedrock deployment's guardrailConfig, so converse, converse-stream, invoke and invoke-with-response-stream reached Bedrock unguarded while /v1/chat/completions on the same deployment was blocked. The signer now sends it the way Bedrock reads it, in the body for Converse and as the X-Amzn-Bedrock-Guardrail* headers for Invoke, replacing any guardrail the caller sent * test(bedrock): reject version-less guardrails in the scripted Bedrock and cover the 400 The mantle and fallback fixtures named a guardrail with no guardrailVersion, which real Bedrock refuses, yet the scripted Bedrock answered them as guarded successes. The fake now answers a version-less guardrail with a 400 ValidationException, the fixtures carry versions, and a new case checks a version-less deployment guardrail reaches Bedrock and the caller gets its rejection on all four actions --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../base_invoke_transformation.py | 4 +- .../bedrock/passthrough/transformation.py | 38 +- litellm/llms/bedrock/request_metadata.py | 11 +- litellm/passthrough/main.py | 3 +- ...drock_passthrough_guardrail_config_wire.py | 531 ++++++++++++++++++ ...test_bedrock_passthrough_transformation.py | 183 +++++- 6 files changed, 757 insertions(+), 13 deletions(-) create mode 100644 tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py 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 f026e4f9a99..87d4fc0b32e 100644 --- a/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py +++ b/litellm/llms/bedrock/chat/invoke_transformations/base_invoke_transformation.py @@ -60,7 +60,7 @@ def _response_header(headers: Mapping[str, str], name: str) -> str | None: return headers.get(name, None) -def _bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": +def bedrock_invoke_guardrail_headers(raw_guardrail_config: object) -> "dict[str, str]": try: guardrail_config: Final = _GUARDRAIL_CONFIG_VALIDATOR.validate_python(raw_guardrail_config) except ValidationError as e: @@ -440,7 +440,7 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM): guardrail_headers: Final = ( () if raw_guardrail_config is None - else tuple(_bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) + else tuple(bedrock_invoke_guardrail_headers(raw_guardrail_config).items()) ) owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params) return merge_bedrock_invoke_headers(headers, guardrail_headers, metadata_headers, owned_names) diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 9d895616336..edd30fff11b 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -14,6 +14,9 @@ from litellm.constants import AWS_ERROR_TYPE_HEADER from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + bedrock_invoke_guardrail_headers, +) from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream @@ -42,6 +45,9 @@ if TYPE_CHECKING: BEDROCK_STREAMING_ACTIONS: Final = frozenset({"invoke-with-response-stream", "converse-stream"}) _BEDROCK_CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) _BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) +_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset( + {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} +) def _bedrock_endpoint_action(endpoint: str) -> str: @@ -123,6 +129,31 @@ def _with_invoke_request_metadata( return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body +def _converse_guardrail_config(action: str, guardrail_config: object) -> object: + fields: Final = as_json_object(guardrail_config) + if action == "converse-stream" or fields is None: + return guardrail_config + return {key: value for key, value in fields.items() if key != "streamProcessingMode"} + + +def _with_deployment_guardrail( + action: str, + headers: Mapping[str, str], + body: Mapping[str, object], + guardrail_config: object, +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if guardrail_config is None: + return headers, body + if action in _BEDROCK_CONVERSE_ACTIONS: + return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)} + if action in _BEDROCK_INVOKE_ACTIONS: + guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items()) + return merge_bedrock_invoke_headers( + dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS + ), body + return headers, body + + def _is_cache_point(entry: object) -> bool: fields: Final = as_json_object(entry) return fields is not None and "cachePoint" in fields @@ -453,11 +484,14 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD body=_converse_body_to_send(body) if action in _BEDROCK_CONVERSE_ACTIONS else body, litellm_params=litellm_params, ) + guarded_headers, guarded_body = _with_deployment_guardrail( + action, signed_headers, signed_body, litellm_params.get("guardrailConfig") + ) return self._sign_request( service_name="bedrock", - headers=dict(signed_headers), + headers=dict(guarded_headers), optional_params=optional_params, - request_data=dict(signed_body), + request_data=dict(guarded_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 71a2dc78578..95a66387b28 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -256,12 +256,11 @@ def merge_bedrock_invoke_headers( """ Merge the ``X-Amzn-*`` headers the Invoke paths derive from params. - ``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is - the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's - headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry - proxy-authenticated identity into an AWS billing record that the caller must not be able to - write. Names are compared case-insensitively so a caller cannot leave a second spelling in - the dict and let the transport pick the winner. + ``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped + from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names + the caller must not be able to write, such as proxy-authenticated identity headed for an AWS + billing record or a deployment's guardrail. Names are compared case-insensitively so a caller + cannot leave a second spelling in the dict and let the transport pick the winner. """ if not caller_owned and not proxy_owned and not proxy_owned_names: return headers diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 744f9e1aa4c..4e94b1bb910 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -437,6 +437,7 @@ def llm_passthrough_route( params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, + guardrailConfig: object = None, **kwargs, ) -> ( httpx.Response @@ -547,7 +548,7 @@ def llm_passthrough_route( _request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None) headers, signed_json_body = provider_config.sign_request( headers=headers, - litellm_params=litellm_params_dict, + litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig}, request_data=_request_data, api_base=str(updated_url), model=model, diff --git a/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py new file mode 100644 index 00000000000..f995a7ac83f --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py @@ -0,0 +1,531 @@ +import base64 +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import boto3 +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BLOCKED_WORD: Final = "zzzforbidden" +_BLOCKED_TEXT: Final = "Blocked by guardrail" +_VERSIONLESS_GUARDRAIL_MESSAGE: Final = "guardrailIdentifier and guardrailVersion must be provided together." +_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream") +_DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +_CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +_MANTLE_KEY: Final = "scripted-mantle-key" +_CALLER_INVOKE_HEADERS: Final[Mapping[str, str]] = { + "x-pass-x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-pass-x-amzn-bedrock-guardrailversion": "9", + "x-pass-x-amzn-bedrock-trace": "DISABLED", +} + + +@dataclass(frozen=True, slots=True) +class Applied: + """The guardrail a request asked Bedrock for, read from where Bedrock reads it for that action.""" + + identifier: str + version: str | None + trace: str | None + stream_processing_mode: str | None + + +@dataclass(frozen=True, slots=True) +class Outcome: + """What the AWS SDK caller reads back.""" + + text: str + intervened: bool + traced_guardrail: str | None + + +def _action(request: Request) -> str: + return request.target.partition("?")[0].rsplit("/", 1)[-1] + + +def _applied(request: Request) -> Applied | None: + if _action(request).startswith("converse"): + config: Final = json.loads(request.body).get("guardrailConfig") + if config is None: + return None + return Applied( + config["guardrailIdentifier"], + config.get("guardrailVersion"), + config.get("trace"), + config.get("streamProcessingMode"), + ) + identifier: Final = request.headers.get("x-amzn-bedrock-guardrailidentifier") + if identifier is None: + return None + return Applied( + identifier, + request.headers.get("x-amzn-bedrock-guardrailversion"), + request.headers.get("x-amzn-bedrock-trace"), + None, + ) + + +def _identifier(request: Request) -> str | None: + applied: Final = _applied(request) + return None if applied is None else applied.identifier + + +def _assessment(applied: Applied) -> Mapping[str, JsonValue]: + return { + applied.identifier: { + "wordPolicy": {"customWords": [{"match": _BLOCKED_WORD, "action": "BLOCKED", "detected": True}]}, + "appliedGuardrailDetails": {"guardrailId": applied.identifier, "guardrailVersion": applied.version}, + } + } + + +def _bedrock_with_guardrails(request: Request) -> Reply: + """Bedrock's runtime as the caller sees it: the guardrail the request names blocks a prompt holding the blocked + word, answering with its message and, when the request enabled tracing, a trace naming that guardrail. A guardrail + named without a version is refused with a 400 ValidationException, as real Bedrock in us-east-1 refused it on + every one of these actions on 2026-10-19.""" + applied: Final = _applied(request) + if applied is not None and applied.version is None: + return Reply( + status=400, + body=json.dumps({"message": _VERSIONLESS_GUARDRAIL_MESSAGE}).encode(), + headers={"x-amzn-errortype": "ValidationException"}, + ) + blocked: Final = applied is not None and _BLOCKED_WORD.encode() in request.body + traced: Final = blocked and applied is not None and (applied.trace or "").lower() == "enabled" + text: Final = _BLOCKED_TEXT if blocked else "ok" + stop_reason: Final = "guardrail_intervened" if blocked else "end_turn" + usage: Final = {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6} + converse_trace: Final = {"trace": {"guardrail": {"inputAssessment": _assessment(applied)}}} if traced else {} + action: Final = _action(request) + if action == "converse": + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": stop_reason, + "usage": usage, + "metrics": {"latencyMs": 3}, + **converse_trace, + } + ).encode() + ) + if action == "converse-stream": + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": stop_reason}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + message: Final = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + **({"amazon-bedrock-guardrailAction": "INTERVENED"} if blocked else {}), + **({"amazon-bedrock-trace": {"guardrail": {"input": _assessment(applied)}}} if traced else {}), + } + if action == "invoke": + return Reply(body=json.dumps(message).encode()) + chunk: Final = {"bytes": base64.b64encode(json.dumps(message).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + + +def _bedrock_outage(request: Request) -> Reply: + return Reply( + status=500, + body=json.dumps({"message": "scripted outage"}).encode(), + headers={"x-amzn-errortype": "InternalServerException"}, + ) + + +def _traced_guardrail(assessment: object) -> str | None: + return None if not isinstance(assessment, dict) else next(iter(assessment)) + + +def _send(runtime: object, action: str, model: str, prompt: str, caller_guardrail: bool) -> Outcome: + messages: Final = [{"role": "user", "content": [{"text": prompt}]}] + invoke_body: Final = json.dumps( + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + ) + invoke_guardrail: Final = ( + {"guardrailIdentifier": "caller-guardrail", "guardrailVersion": "9", "trace": "DISABLED"} + if caller_guardrail + else {} + ) + if action == "converse": + response: Final = getattr(runtime, "converse")( + modelId=model, + messages=messages, + **({"guardrailConfig": dict(_CALLER_GUARDRAIL)} if caller_guardrail else {}), + ) + return Outcome( + text=response["output"]["message"]["content"][0]["text"], + intervened=response["stopReason"] == "guardrail_intervened", + traced_guardrail=_traced_guardrail(response.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "converse-stream": + events: Final = list( + getattr(runtime, "converse_stream")( + modelId=model, + messages=messages, + **( + {"guardrailConfig": {**_CALLER_GUARDRAIL, "streamProcessingMode": "async"}} + if caller_guardrail + else {} + ), + )["stream"] + ) + (stop,) = tuple(event["messageStop"]["stopReason"] for event in events if "messageStop" in event) + (metadata,) = tuple(event["metadata"] for event in events if "metadata" in event) + return Outcome( + text="".join( + event["contentBlockDelta"]["delta"]["text"] for event in events if "contentBlockDelta" in event + ), + intervened=stop == "guardrail_intervened", + traced_guardrail=_traced_guardrail(metadata.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "invoke": + message: Final = json.loads( + getattr(runtime, "invoke_model")(modelId=model, body=invoke_body, **invoke_guardrail)["body"].read() + ) + else: + (message,) = tuple( + json.loads(event["chunk"]["bytes"]) + for event in getattr(runtime, "invoke_model_with_response_stream")( + modelId=model, body=invoke_body, **invoke_guardrail + )["body"] + ) + return Outcome( + text=message["content"][0]["text"], + intervened=message.get("amazon-bedrock-guardrailAction") == "INTERVENED", + traced_guardrail=_traced_guardrail(message.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input")), + ) + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _deployment(wire: Wire, guardrail_config: Mapping[str, str] | None = None) -> Mapping[str, JsonValue]: + return { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + **({} if guardrail_config is None else {"guardrailConfig": dict(guardrail_config)}), + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + guarded_model: str + unguarded_model: str + balanced_model: str + balanced_wires: Mapping[str, Wire] + failing_model: str + failing_wire: Wire + fallback_wire: Wire + mantle_model: str + versionless_model: str + key: str + + def runtime( # noqa: ANN201 # boto3 clients are untyped + self, monkeypatch: pytest.MonkeyPatch, extra_headers: Mapping[str, str] = MappingProxyType({}) + ): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", self.key) + client: Final = boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}), + ) + client.meta.events.register( + "before-send.bedrock-runtime.*", + lambda request, **_: request.headers.update(extra_headers), + ) + return client + + def received(self, wire: Wire, marker: str) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + guarded_model: Final = f"bedrock-guarded-{suffix}" + unguarded_model: Final = f"bedrock-unguarded-{suffix}" + balanced_model: Final = f"bedrock-balanced-{suffix}" + failing_model: Final = f"bedrock-failing-{suffix}" + fallback_model: Final = f"bedrock-fallback-{suffix}" + mantle_model: Final = f"bedrock-mantle-{suffix}" + versionless_model: Final = f"bedrock-versionless-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-guardrail-config") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with ( + wire_server(_bedrock_with_guardrails) as wire, + wire_server(_bedrock_with_guardrails) as wire_a, + wire_server(_bedrock_with_guardrails) as wire_b, + wire_server(_bedrock_outage) as failing_wire, + wire_server(_bedrock_with_guardrails) as fallback_wire, + ): + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + { + "model_name": guarded_model, + "litellm_params": _deployment(wire, _DEPLOYMENT_GUARDRAIL), + }, + {"model_name": unguarded_model, "litellm_params": _deployment(wire)}, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_a, {"guardrailIdentifier": "guardrail-a", "guardrailVersion": "1"} + ), + }, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_b, {"guardrailIdentifier": "guardrail-b", "guardrailVersion": "2"} + ), + }, + { + "model_name": failing_model, + "litellm_params": _deployment( + failing_wire, + {"guardrailIdentifier": "failing-guardrail", "guardrailVersion": "5", "trace": "enabled"}, + ), + }, + { + "model_name": fallback_model, + "litellm_params": _deployment( + fallback_wire, + {"guardrailIdentifier": "fallback-guardrail", "guardrailVersion": "6", "trace": "enabled"}, + ), + }, + { + "model_name": versionless_model, + "litellm_params": _deployment( + wire, {"guardrailIdentifier": "versionless-guardrail", "trace": "enabled"} + ), + }, + { + "model_name": mantle_model, + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-sonnet-5", + "api_base": wire.url, + "api_key": _MANTLE_KEY, + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + "guardrailConfig": { + "guardrailIdentifier": "mantle-guardrail", + "guardrailVersion": "4", + "trace": "enabled", + }, + }, + }, + ], + "router_settings": { + **base_config.get("router_settings", {}), + "fallbacks": [{failing_model: [fallback_model]}], + }, + } + config_path: Final = directory / "bedrock-guardrail-config.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + yield Rig( + proxy=proxy, + wire=wire, + guarded_model=guarded_model, + unguarded_model=unguarded_model, + balanced_model=balanced_model, + balanced_wires={"guardrail-a": wire_a, "guardrail-b": wire_b}, + failing_model=failing_model, + failing_wire=failing_wire, + fallback_wire=fallback_wire, + mantle_model=mantle_model, + versionless_model=versionless_model, + key=scenario.key(key_alias=f"guardrail-key-{suffix}"), + ) + + +def _expected_on_the_wire(action: str, request: Request) -> Applied: + if action == "converse": + return Applied("deployment-guardrail", "3", "enabled", None) + if action == "converse-stream": + return Applied("deployment-guardrail", "3", "enabled", "sync") + assert {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} <= ( + _signed_header_names(request) + ), request.headers["authorization"] + return Applied("deployment-guardrail", "3", "ENABLED", None) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_passthrough_on_a_guardrailed_deployment_blocks_with_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + blocked: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + allowed: Final = _send(runtime, action, rig.guarded_model, f"{marker} say hello", caller_guardrail=False) + + assert blocked == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), blocked + assert allowed == Outcome("ok", intervened=False, traced_guardrail=None), allowed + received: Final = rig.received(rig.wire, marker) + assert [_action(request) for request in received] == [action, action], received + assert [_applied(request) for request in received] == [ + _expected_on_the_wire(action, request) for request in received + ], [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_caller_guardrail_is_replaced_by_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch, _CALLER_INVOKE_HEADERS) + + outcome: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == _expected_on_the_wire(action, request), (request.headers, request.body) + assert "caller-guardrail" not in json.dumps({**request.headers, "body": request.body.decode()}), request + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_mantle_deployment_passthrough_sends_its_guardrail_with_bearer_auth( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.mantle_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="mantle-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert request.headers["authorization"] == f"Bearer {_MANTLE_KEY}", request.headers + assert _applied(request) == Applied( + "mantle-guardrail", "4", "ENABLED" if action.startswith("invoke") else "enabled", None + ), (request.headers, request.body) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_guardrail_without_a_version_returns_the_bedrock_rejection( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + with pytest.raises(ClientError) as raised: + _send(runtime, action, rig.versionless_model, f"{marker} say hello", caller_guardrail=False) + + error: Final = raised.value.response.get("Error", {}) + assert (error.get("Code"), error.get("Message")) == ("ValidationException", _VERSIONLESS_GUARDRAIL_MESSAGE), error + received: Final = rig.received(rig.wire, marker) + assert {_applied(request) for request in received} == { + Applied("versionless-guardrail", None, "ENABLED" if action.startswith("invoke") else "enabled", None) + }, [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + if action.startswith("invoke"): + pytest.skip("BUG: the router-model invoke pass-through drops the caller's X-Amzn-Bedrock-Guardrail* headers") + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.unguarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail=None), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == Applied( + "caller-guardrail", "9", "disabled", "async" if action == "converse-stream" else None + ), request.body + + +def test_each_load_balanced_request_carries_the_guardrail_of_the_deployment_it_reached( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.balanced_model, f"{marker}-{attempt} say {_BLOCKED_WORD}", caller_guardrail=False) + for attempt, action in enumerate(_ACTIONS * 6) + ) + + assert all(outcome.intervened for outcome in outcomes), outcomes + reached: Final = {name: rig.received(wire, marker) for name, wire in rig.balanced_wires.items()} + assert sum(len(requests) for requests in reached.values()) == len(outcomes), reached + assert all(reached.values()), {name: len(requests) for name, requests in reached.items()} + assert {name: {_identifier(request) for request in requests} for name, requests in reached.items()} == { + "guardrail-a": {"guardrail-a"}, + "guardrail-b": {"guardrail-b"}, + }, reached + + +def test_a_fallback_sends_the_guardrail_of_the_deployment_it_fell_back_to( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.failing_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + for action in _ACTIONS + ) + + assert outcomes == (Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="fallback-guardrail"),) * 4, outcomes + failed: Final = rig.received(rig.failing_wire, marker) + served: Final = rig.received(rig.fallback_wire, marker) + assert [_action(request) for request in served] == list(_ACTIONS), served + assert {_identifier(request) for request in failed} == {"failing-guardrail"}, failed + assert {_identifier(request) for request in served} == {"fallback-guardrail"}, served diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index da5c4699925..44c0e82c006 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -5,8 +5,9 @@ import struct import sys import types from binascii import crc32 -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import Final from unittest.mock import patch @@ -922,7 +923,11 @@ INVOKE_BODY = { async def _sent_passthrough_request( - action: str, body: dict[str, object], identity: dict[str, object], **kwargs: object + action: str, + body: dict[str, object], + identity: dict[str, object], + deployment_params: Mapping[str, object] = MappingProxyType({}), + **kwargs: object, ) -> httpx.Request: requests: Final[list[httpx.Request]] = [] @@ -939,6 +944,7 @@ async def _sent_passthrough_request( "aws_access_key_id": "AKIAROUTERKEY", "aws_secret_access_key": "router-secret", "aws_region_name": "us-west-2", + **deployment_params, }, } ] @@ -1189,3 +1195,176 @@ async def test_passthrough_body_is_sent_as_is_when_params_are_dropped_and_forwar assert json.loads(sent.content) == caller_body assert REQUEST_METADATA_HEADER not in sent.headers + + +DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-amzn-bedrock-guardrailidentifier": "deployment-guardrail", + "x-amzn-bedrock-guardrailversion": "3", + "x-amzn-bedrock-trace": "ENABLED", +} +CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +CALLER_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "x-pass-X-Amzn-Bedrock-GuardrailVersion": "9", + "x-pass-X-Amzn-Bedrock-Trace": "DISABLED", +} + + +def _invoke_guardrail_headers(request: httpx.Request) -> dict[str, str]: + return { + name: value + for name, value in request.headers.items() + if name.lower().startswith("x-amzn-bedrock-guardrail") or name.lower() == "x-amzn-bedrock-trace" + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "sent_guardrail"), + [ + ("converse", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), + ("converse-stream", DEPLOYMENT_GUARDRAIL), + ], +) +async def test_converse_passthrough_sends_the_deployment_guardrail_in_the_body( + action: str, sent_guardrail: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + action, CONVERSE_BODY, {}, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": sent_guardrail} + assert _invoke_guardrail_headers(request) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_the_deployment_guardrail_as_signed_headers( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS) + + request: Final = await _sent_passthrough_request( + action, INVOKE_BODY, PROXY_IDENTITY, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert _invoke_guardrail_headers(request) == DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS + assert {*DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS, REQUEST_METADATA_HEADER} <= _signed_header_names(request) + assert json.loads(request.headers[REQUEST_METADATA_HEADER]) == EXPECTED_REQUEST_METADATA + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "caller_guardrail"), + [("converse", CALLER_GUARDRAIL), ("converse-stream", {**CALLER_GUARDRAIL, "streamProcessingMode": "async"})], +) +async def test_deployment_guardrail_replaces_the_caller_converse_guardrail( + action: str, caller_guardrail: Mapping[str, str], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + action, + {**CONVERSE_BODY, "guardrailConfig": caller_guardrail}, + {}, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_deployment_guardrail_replaces_every_invoke_guardrail_header_the_caller_sent( + action: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + {}, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS, + ) + + assert _invoke_guardrail_headers(request) == {"x-amzn-bedrock-guardrailidentifier": "deployment-guardrail"} + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +async def test_passthrough_without_a_deployment_guardrail_sends_the_caller_guardrail_as_is( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + caller_body: Final = {**CONVERSE_BODY, "guardrailConfig": CALLER_GUARDRAIL} + + converse: Final = await _sent_passthrough_request("converse", caller_body, {}) + invoke: Final = await _sent_passthrough_request( + "invoke", INVOKE_BODY, {}, request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS + ) + + assert json.loads(converse.content) == caller_body + assert _invoke_guardrail_headers(invoke) == { + "x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-amzn-bedrock-guardrailversion": "9", + "x-amzn-bedrock-trace": "DISABLED", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize("guardrail_config", ["deployment-guardrail", {"guardrailVersion": "3"}]) +async def test_invoke_passthrough_refuses_a_deployment_guardrail_it_cannot_send( + action: str, guardrail_config: object, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + with pytest.raises(BedrockError, match="guardrailConfig") as rejected: + await _sent_passthrough_request( + action, INVOKE_BODY, {}, deployment_params={"guardrailConfig": guardrail_config} + ) + + assert rejected.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_converse_passthrough_forwards_a_non_object_deployment_guardrail_unchanged( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + + request: Final = await _sent_passthrough_request( + "converse", CONVERSE_BODY, {}, deployment_params={"guardrailConfig": "deployment-guardrail"} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": "deployment-guardrail"} + + +@pytest.mark.asyncio +async def test_passthrough_leaves_actions_without_a_guardrail_field_as_sent(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + count_tokens_body: Final = {"input": {"converse": CONVERSE_BODY}} + + request: Final = await _sent_passthrough_request( + "count-tokens", count_tokens_body, {}, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == count_tokens_body + assert _invoke_guardrail_headers(request) == {}