mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-27 06:13:07 +00:00
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>
This commit is contained in:
yucheng-berri
GitHub
parent
6bc1b2b3e5
commit
e6065b9f4f
6 changed files with 757 additions and 13 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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) == {}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue