mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
feat(bedrock): apply bedrock_request_metadata_fields on the /bedrock pass-through routes (#47040)
* feat(bedrock): apply bedrock_request_metadata_fields on the /bedrock pass-through routes * fix(bedrock): move Invoke body requestMetadata into the signed header and match Bedrock's ASCII whitespace rule * fix(bedrock): keep Invoke body requestMetadata in the body and answer bad metadata as AWS does Invoke requests whose body already carries requestMetadata keep the resolved pairs in the body instead of the 8,500-character signed header. Malformed caller metadata now answers 400 with the ValidationException error type the AWS SDK expects, and drop_params drops bad pairs on both the pass-through routes and /v1/chat/completions. The caller-pair dedupe is linear, and the chat validator treats \s as ASCII whitespace, as Bedrock does * fix(bedrock): match Bedrock's 400s and header limit on pass-through request metadata Invoke metadata too long for the 8,500-character header now goes in the body, Invoke requestMetadata null answers 400 like Bedrock, Converse type errors carry SerializationException, violation messages name the failing key, and the reserved user_api_key_ prefix is matched case-insensitively after trimming whitespace * fix(bedrock): trim the Invoke metadata header to Bedrock's limit and drop bad requestMetadata under drop_params without the setting * fix(bedrock): leave /bedrock pass-through bodies as sent when request metadata forwarding is off Global drop_params no longer rewrites requestMetadata on the pass-through routes unless bedrock_request_metadata_fields is set. Chat keeps dropping bad pairs under drop_params. Adds a tests/integration wire test for the four pass-through actions and removes an unused parameter from the Invoke header helper. * test(bedrock): drive the pass-through metadata wire test with boto3's bearer key env var * fix(bedrock): trim the pass-through invoke metadata header and stop Nova invoke from checking metadata it drops * test(bedrock): cover pass-through metadata header trimming, malformed caller metadata, and chat drop_params on the wire * fix(bedrock): match requestMetadata keys and values against the whole string so a trailing newline cannot pass the 256-character limit
This commit is contained in:
Mateo Wang
GitHub
parent
caad4d6204
commit
82ca10cc2d
13 changed files with 1043 additions and 124 deletions
|
|
@ -66,7 +66,9 @@ from litellm.llms.bedrock.common_utils import (
|
|||
)
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
bedrock_request_metadata_is_owned,
|
||||
bedrock_request_metadata_is_rewritten,
|
||||
bedrock_request_metadata_violation,
|
||||
caller_request_metadata_violations_are_dropped,
|
||||
merge_bedrock_invoke_headers,
|
||||
resolve_bedrock_request_metadata,
|
||||
)
|
||||
|
|
@ -354,77 +356,15 @@ class AmazonConverseConfig(BaseConfig):
|
|||
and v is not None
|
||||
}
|
||||
|
||||
def _validate_request_metadata(self, metadata: dict) -> None:
|
||||
"""
|
||||
Validate requestMetadata according to AWS Bedrock Converse API constraints.
|
||||
|
||||
Constraints:
|
||||
- Maximum of 16 items
|
||||
- Keys: 1-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{1,256}
|
||||
- Values: 0-256 characters, pattern [a-zA-Z0-9\\s:_@$#=/+,-.]{0,256}
|
||||
"""
|
||||
import re
|
||||
|
||||
if not isinstance(metadata, dict):
|
||||
def _validate_request_metadata(self, metadata: object) -> None:
|
||||
violation: Final = bedrock_request_metadata_violation(metadata)
|
||||
if violation is not None:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="requestMetadata must be a dictionary",
|
||||
message=violation.message,
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
if len(metadata) > 16:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="requestMetadata can contain a maximum of 16 items",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
key_pattern: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$")
|
||||
value_pattern: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$")
|
||||
|
||||
for key, value in metadata.items():
|
||||
if not isinstance(key, str):
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="requestMetadata keys must be strings",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
if not isinstance(value, str):
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="requestMetadata values must be strings",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
if len(key) == 0 or len(key) > 256:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="requestMetadata key length must be 1-256 characters",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
if len(value) > 256:
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message="requestMetadata value length must be 0-256 characters",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
if not key_pattern.match(key):
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message=f"requestMetadata key '{key}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
if not value_pattern.match(value):
|
||||
raise litellm.exceptions.BadRequestError(
|
||||
message=f"requestMetadata value '{value}' contains invalid characters. Allowed: [a-zA-Z0-9\\s:_@$#=/+,.-]",
|
||||
model="bedrock",
|
||||
llm_provider="bedrock",
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _is_openai_gpt_reasoning_model(model: str) -> bool:
|
||||
return re.search(r"openai\.gpt-\d", model) is not None
|
||||
|
|
@ -1244,7 +1184,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
elif param == "context_management" and isinstance(value, (dict, list)):
|
||||
self._map_context_management_param(value, optional_params)
|
||||
if param == "requestMetadata":
|
||||
self._map_request_metadata_param(value, optional_params)
|
||||
self._map_request_metadata_param(value, optional_params, drop_params)
|
||||
if param == "service_tier" and isinstance(value, str):
|
||||
self._map_service_tier_param(value, optional_params)
|
||||
|
||||
|
|
@ -1277,9 +1217,10 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
return optional_params
|
||||
|
||||
def _map_request_metadata_param(self, value: object, optional_params: dict) -> None:
|
||||
def _map_request_metadata_param(self, value: object, optional_params: dict, drop_params: bool) -> None:
|
||||
if value is not None and isinstance(value, dict):
|
||||
self._validate_request_metadata(value)
|
||||
if not caller_request_metadata_violations_are_dropped(drop_params):
|
||||
self._validate_request_metadata(value)
|
||||
optional_params["requestMetadata"] = value
|
||||
|
||||
def _map_context_management_param(self, value: dict | list, optional_params: dict) -> None:
|
||||
|
|
@ -1737,7 +1678,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
|
||||
# Extract requestMetadata before processing other parameters
|
||||
request_metadata: Final = inference_params.pop("requestMetadata", None)
|
||||
if request_metadata is not None:
|
||||
if request_metadata is not None and not caller_request_metadata_violations_are_dropped(drop_params):
|
||||
self._validate_request_metadata(request_metadata)
|
||||
|
||||
output_config: OutputConfigBlock | None = inference_params.pop("outputConfig", None)
|
||||
|
|
@ -2164,7 +2105,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
request_metadata: Final = resolve_bedrock_request_metadata(
|
||||
litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata")
|
||||
)
|
||||
if bedrock_request_metadata_is_owned():
|
||||
if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True):
|
||||
_data.pop("requestMetadata", None)
|
||||
if request_metadata is not None:
|
||||
_data["requestMetadata"] = request_metadata
|
||||
|
|
@ -2244,7 +2185,7 @@ class AmazonConverseConfig(BaseConfig):
|
|||
request_metadata: Final = resolve_bedrock_request_metadata(
|
||||
litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata")
|
||||
)
|
||||
if bedrock_request_metadata_is_owned():
|
||||
if bedrock_request_metadata_is_rewritten(litellm_params.get("drop_params") is True):
|
||||
_data.pop("requestMetadata", None)
|
||||
if request_metadata is not None:
|
||||
_data["requestMetadata"] = request_metadata
|
||||
|
|
|
|||
|
|
@ -95,7 +95,12 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
|
|||
model_info: Mapping[str, object] | None = None,
|
||||
) -> dict:
|
||||
return AmazonConverseConfig.map_openai_params(
|
||||
self, non_default_params, optional_params, model, drop_params, model_info
|
||||
self,
|
||||
{key: value for key, value in non_default_params.items() if key != "requestMetadata"},
|
||||
optional_params,
|
||||
model,
|
||||
drop_params,
|
||||
model_info,
|
||||
)
|
||||
|
||||
def transform_request(
|
||||
|
|
@ -110,7 +115,9 @@ class AmazonInvokeNovaConfig(AmazonInvokeConfig, AmazonConverseConfig):
|
|||
self,
|
||||
model=model,
|
||||
messages=messages,
|
||||
optional_params=_without_tool_config_injection_points(optional_params),
|
||||
optional_params=_without_tool_config_injection_points(
|
||||
{key: value for key, value in optional_params.items() if key != "requestMetadata"}
|
||||
),
|
||||
litellm_params=litellm_params,
|
||||
headers=headers,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -1021,7 +1021,7 @@ class BedrockFilesConfig(BaseAWSLLM, BaseFilesConfig):
|
|||
|
||||
converse_config: Final = AmazonConverseConfig()
|
||||
mapped_params = converse_config.map_openai_params(
|
||||
non_default_params=optional_params,
|
||||
non_default_params={k: v for k, v in optional_params.items() if k != "requestMetadata"},
|
||||
optional_params={},
|
||||
model=model,
|
||||
drop_params=False,
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from httpx import Response
|
|||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
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.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage
|
||||
|
|
@ -17,6 +18,16 @@ from litellm.types.utils import ModelResponseStream
|
|||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo
|
||||
from ..request_metadata import (
|
||||
BEDROCK_REQUEST_METADATA_HEADER,
|
||||
BedrockRequestMetadataViolation,
|
||||
bedrock_request_metadata_headers,
|
||||
bedrock_request_metadata_is_owned,
|
||||
bedrock_request_metadata_violation,
|
||||
caller_request_metadata_violations_are_dropped,
|
||||
merge_bedrock_invoke_headers,
|
||||
resolve_bedrock_request_metadata,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from botocore.eventstream import EventStreamMessage
|
||||
|
|
@ -28,10 +39,87 @@ 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"})
|
||||
|
||||
|
||||
def _bedrock_endpoint_action(endpoint: str) -> str:
|
||||
return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1]
|
||||
|
||||
|
||||
def is_bedrock_streaming_endpoint(endpoint: str) -> bool:
|
||||
return endpoint.partition("?")[0].rstrip("/").rsplit("/", 1)[-1] in BEDROCK_STREAMING_ACTIONS
|
||||
return _bedrock_endpoint_action(endpoint) in BEDROCK_STREAMING_ACTIONS
|
||||
|
||||
|
||||
def _caller_request_metadata_violation(
|
||||
action: str, body: Mapping[str, object]
|
||||
) -> BedrockRequestMetadataViolation | None:
|
||||
if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body and body["requestMetadata"] is None:
|
||||
return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary")
|
||||
return bedrock_request_metadata_violation(body.get("requestMetadata"))
|
||||
|
||||
|
||||
def _aws_error_type(action: str, violation: BedrockRequestMetadataViolation) -> str:
|
||||
if action in _BEDROCK_CONVERSE_ACTIONS and violation.kind == "type":
|
||||
return "SerializationException"
|
||||
return "ValidationException"
|
||||
|
||||
|
||||
def _caller_request_metadata(action: str, body: Mapping[str, object], litellm_params: Mapping[str, object]) -> object:
|
||||
violation: Final = (
|
||||
None
|
||||
if caller_request_metadata_violations_are_dropped(litellm_params.get("drop_params") is True)
|
||||
else _caller_request_metadata_violation(action, body)
|
||||
)
|
||||
if violation is not None:
|
||||
raise BedrockError(
|
||||
status_code=400,
|
||||
message=violation.message,
|
||||
headers={AWS_ERROR_TYPE_HEADER: _aws_error_type(action, violation)},
|
||||
)
|
||||
return body.get("requestMetadata")
|
||||
|
||||
|
||||
def _with_body_request_metadata(
|
||||
action: str, body: Mapping[str, object], litellm_params: Mapping[str, object]
|
||||
) -> Mapping[str, object]:
|
||||
resolved: Final = resolve_bedrock_request_metadata(
|
||||
litellm_params, caller_metadata=_caller_request_metadata(action, body, litellm_params)
|
||||
)
|
||||
without_caller: Final = {key: value for key, value in body.items() if key != "requestMetadata"}
|
||||
return without_caller if resolved is None else {**without_caller, "requestMetadata": resolved}
|
||||
|
||||
|
||||
def _without_owned_request_metadata_header(headers: Mapping[str, str]) -> Mapping[str, str]:
|
||||
if not bedrock_request_metadata_is_owned():
|
||||
return headers
|
||||
return {name: value for name, value in headers.items() if name.lower() != BEDROCK_REQUEST_METADATA_HEADER.lower()}
|
||||
|
||||
|
||||
def _with_request_metadata(
|
||||
action: str | None,
|
||||
headers: Mapping[str, str],
|
||||
body: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[Mapping[str, str], Mapping[str, object]]:
|
||||
if action in _BEDROCK_CONVERSE_ACTIONS:
|
||||
return headers, _with_body_request_metadata(action, body, litellm_params)
|
||||
if action in _BEDROCK_INVOKE_ACTIONS and "requestMetadata" in body:
|
||||
return _without_owned_request_metadata_header(headers), _with_body_request_metadata(
|
||||
action, body, litellm_params
|
||||
)
|
||||
if action in _BEDROCK_INVOKE_ACTIONS:
|
||||
return _with_invoke_request_metadata(headers, body, litellm_params)
|
||||
return headers, body
|
||||
|
||||
|
||||
def _with_invoke_request_metadata(
|
||||
headers: Mapping[str, str],
|
||||
body: Mapping[str, object],
|
||||
litellm_params: Mapping[str, object],
|
||||
) -> tuple[Mapping[str, str], Mapping[str, object]]:
|
||||
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
|
||||
return merge_bedrock_invoke_headers(dict(headers), (), metadata_headers, owned_names), body
|
||||
|
||||
|
||||
_TEXT_ONLY_DELTA_FIELDS: Final = frozenset({"content", "role"})
|
||||
|
|
@ -284,18 +372,24 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD
|
|||
|
||||
def sign_request(
|
||||
self,
|
||||
headers: dict,
|
||||
litellm_params: dict,
|
||||
request_data: dict | None,
|
||||
headers: dict[str, str],
|
||||
litellm_params: dict[str, object],
|
||||
request_data: dict[str, object] | None,
|
||||
api_base: str,
|
||||
model: str | None = None,
|
||||
) -> tuple[dict, bytes | None]:
|
||||
optional_params: Final = litellm_params.copy()
|
||||
signed_headers, signed_body = _with_request_metadata(
|
||||
action=_bedrock_endpoint_action(api_base) if bedrock_request_metadata_is_owned() else None,
|
||||
headers=headers,
|
||||
body=request_data or {},
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
return self._sign_request(
|
||||
service_name="bedrock",
|
||||
headers=headers,
|
||||
headers=dict(signed_headers),
|
||||
optional_params=optional_params,
|
||||
request_data=request_data or {},
|
||||
request_data=dict(signed_body),
|
||||
api_base=api_base,
|
||||
model=model,
|
||||
api_key=self.get_bedrock_bearer_token(optional_params),
|
||||
|
|
|
|||
|
|
@ -8,17 +8,22 @@ sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the pr
|
|||
Two properties are load-bearing for that billing record and are asserted by the tests:
|
||||
proxy identity is resolved first so it can never be evicted by caller-supplied pairs, and the
|
||||
whole ``user_api_key_`` prefix is reserved so a caller cannot write a proxy-authoritative
|
||||
looking key. Values that break Bedrock's constraints are dropped rather than sanitised or
|
||||
rejected, because an operator flipping this setting on must not turn a working request into a
|
||||
400 and a silently rewritten attribution key is worse than an absent one.
|
||||
looking key. Identity and ``spend_logs_metadata`` values that break Bedrock's constraints are
|
||||
dropped rather than sanitised or rejected, because an operator flipping this setting on must not
|
||||
turn a working request into a 400 and a silently rewritten attribution key is worse than an
|
||||
absent one. For the same reason an Invoke header that would pass Bedrock's 8,500-character limit
|
||||
keeps only the leading pairs that fit. A caller's own malformed ``requestMetadata`` is rejected with
|
||||
a 400 unless ``drop_params`` is set, in which case its bad pairs are dropped the same way, with or
|
||||
without the setting.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
from collections.abc import Mapping, Sized
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
import litellm
|
||||
|
||||
|
|
@ -26,15 +31,27 @@ BEDROCK_REQUEST_METADATA_HEADER: Final = "X-Amzn-Bedrock-Request-Metadata"
|
|||
BEDROCK_REQUEST_METADATA_MAX_PAIRS: Final = 16
|
||||
BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX: Final = "user_api_key_"
|
||||
BEDROCK_REQUEST_METADATA_CLIENT_FIELD: Final = "spend_logs_metadata"
|
||||
BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS: Final = 8500
|
||||
|
||||
_METADATA_PARAM_NAMES: Final[tuple[str, ...]] = ("metadata", "litellm_metadata")
|
||||
_KEY_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{1,256}$")
|
||||
_VALUE_PATTERN: Final = re.compile(r"^[a-zA-Z0-9\s:_@$#=/+,.-]{0,256}$")
|
||||
_ALLOWED_CHARACTERS: Final = r"[a-zA-Z0-9\s:_@$#=/+,.-]"
|
||||
_KEY_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{1,256}}", re.ASCII)
|
||||
_VALUE_PATTERN: Final = re.compile(rf"{_ALLOWED_CHARACTERS}{{0,256}}", re.ASCII)
|
||||
_OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),))
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BedrockRequestMetadataViolation:
|
||||
kind: Literal["type", "constraint"]
|
||||
message: str
|
||||
|
||||
|
||||
def _is_forwardable(key: str, value: str) -> bool:
|
||||
return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None
|
||||
return _KEY_PATTERN.fullmatch(key) is not None and _VALUE_PATTERN.fullmatch(value) is not None
|
||||
|
||||
|
||||
def _is_reserved(key: str) -> bool:
|
||||
return key.strip().lower().startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX)
|
||||
|
||||
|
||||
def _text_pairs(source: object) -> tuple[tuple[str, str], ...]:
|
||||
|
|
@ -99,13 +116,16 @@ def _client_pairs(
|
|||
candidates: Final = tuple(
|
||||
(key, value)
|
||||
for key, value in (*_text_pairs(caller_metadata), *spend_logs_pairs)
|
||||
if not key.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX) and _is_forwardable(key, value)
|
||||
if not _is_reserved(key) and _is_forwardable(key, value)
|
||||
)
|
||||
return tuple(
|
||||
pair
|
||||
for index, pair in enumerate(candidates)
|
||||
if pair[0] not in tuple(earlier for earlier, _ in candidates[:index])
|
||||
)[:budget]
|
||||
first_index: Final = {key: index for index, (key, _) in reversed(tuple(enumerate(candidates)))}
|
||||
return tuple(pair for index, pair in enumerate(candidates) if first_index[pair[0]] == index)[:budget]
|
||||
|
||||
|
||||
def _forwardable_caller_pairs(caller_metadata: object) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(pair for pair in _text_pairs(caller_metadata) if _is_forwardable(*pair))[
|
||||
:BEDROCK_REQUEST_METADATA_MAX_PAIRS
|
||||
]
|
||||
|
||||
|
||||
def resolve_bedrock_request_metadata(
|
||||
|
|
@ -113,17 +133,18 @@ def resolve_bedrock_request_metadata(
|
|||
caller_metadata: object = None,
|
||||
) -> dict[str, str] | None:
|
||||
"""
|
||||
Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when the feature is
|
||||
off or nothing survives Bedrock's constraints. The result is a plain dict because it is
|
||||
written straight onto the Converse body, which Bedrock types as ``dict[str, str]``.
|
||||
Resolve the ``requestMetadata`` pairs to send to Bedrock, or ``None`` when nothing survives
|
||||
Bedrock's constraints. The result is a plain dict because it is written straight onto the
|
||||
Converse body, which Bedrock types as ``dict[str, str]``.
|
||||
|
||||
``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. It has already
|
||||
been validated (and rejected with a 400) by the Converse transformation, so it is only
|
||||
filtered here for the reserved identity prefix and the remaining slot budget.
|
||||
``caller_metadata`` is any ``requestMetadata`` the caller passed explicitly. Callers reject a
|
||||
malformed one with a 400 unless ``drop_params`` is set, so here its pairs are only filtered for
|
||||
Bedrock's rules, the reserved identity prefix, and the remaining slot budget. With the feature
|
||||
off only Bedrock's rules apply, since no identity is written for a caller to forge.
|
||||
"""
|
||||
allowed_fields: Final = _allowed_fields()
|
||||
if not allowed_fields:
|
||||
return None
|
||||
return dict(_forwardable_caller_pairs(caller_metadata)) or None
|
||||
sources: Final = _metadata_sources(litellm_params)
|
||||
identity: Final = _identity_pairs(sources, allowed_fields)
|
||||
client: Final = _client_pairs(
|
||||
|
|
@ -136,6 +157,36 @@ def resolve_bedrock_request_metadata(
|
|||
return resolved or None
|
||||
|
||||
|
||||
def _pair_violation(key: str, value: str) -> BedrockRequestMetadataViolation | None:
|
||||
if _KEY_PATTERN.fullmatch(key) is None:
|
||||
return BedrockRequestMetadataViolation(
|
||||
"constraint", f"requestMetadata key {key!r} must be 1-256 characters of {_ALLOWED_CHARACTERS}"
|
||||
)
|
||||
if _VALUE_PATTERN.fullmatch(value) is None:
|
||||
return BedrockRequestMetadataViolation(
|
||||
"constraint", f"requestMetadata value for key {key!r} must be 0-256 characters of {_ALLOWED_CHARACTERS}"
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def bedrock_request_metadata_violation(caller_metadata: object) -> BedrockRequestMetadataViolation | None:
|
||||
"""Why a caller's ``requestMetadata`` breaks Bedrock's rules, in the order Bedrock checks them: a
|
||||
wrong JSON type before the pair count and character rules. ``None`` means absent or valid."""
|
||||
text_pairs: Final = _text_pairs(caller_metadata)
|
||||
entry_count: Final = len(caller_metadata) if isinstance(caller_metadata, Sized) else 0
|
||||
if caller_metadata is None:
|
||||
return None
|
||||
if not isinstance(caller_metadata, Mapping):
|
||||
return BedrockRequestMetadataViolation("type", "requestMetadata must be a dictionary")
|
||||
if len(text_pairs) != entry_count:
|
||||
return BedrockRequestMetadataViolation("type", "requestMetadata keys and values must be strings")
|
||||
if len(text_pairs) > BEDROCK_REQUEST_METADATA_MAX_PAIRS:
|
||||
return BedrockRequestMetadataViolation(
|
||||
"constraint", f"requestMetadata can contain a maximum of {BEDROCK_REQUEST_METADATA_MAX_PAIRS} items"
|
||||
)
|
||||
return next((violation for key, value in text_pairs if (violation := _pair_violation(key, value))), None)
|
||||
|
||||
|
||||
def bedrock_request_metadata_is_owned() -> bool:
|
||||
"""
|
||||
Whether the proxy OWNS the request-metadata field and header name for this request.
|
||||
|
|
@ -149,25 +200,51 @@ def bedrock_request_metadata_is_owned() -> bool:
|
|||
return bool(_allowed_fields())
|
||||
|
||||
|
||||
def caller_request_metadata_violations_are_dropped(drop_params: bool) -> bool:
|
||||
return drop_params or litellm.drop_params is True
|
||||
|
||||
|
||||
def bedrock_request_metadata_is_rewritten(drop_params: bool) -> bool:
|
||||
return bedrock_request_metadata_is_owned() or caller_request_metadata_violations_are_dropped(drop_params)
|
||||
|
||||
|
||||
def _header_value(pairs: Mapping[str, str]) -> str:
|
||||
return json.dumps(pairs, separators=(",", ":"))
|
||||
|
||||
|
||||
def _fits_header(pairs: Mapping[str, str]) -> bool:
|
||||
return len(_header_value(pairs)) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS
|
||||
|
||||
|
||||
def _longest_fitting_prefix(resolved: Mapping[str, str]) -> dict[str, str]:
|
||||
pairs: Final = tuple(resolved.items())
|
||||
return next(
|
||||
(dict(pairs[:count]) for count in range(len(pairs), 0, -1) if _fits_header(dict(pairs[:count]))),
|
||||
{},
|
||||
)
|
||||
|
||||
|
||||
def bedrock_request_metadata_headers(
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
) -> tuple[frozenset[str], tuple[tuple[str, str], ...]]:
|
||||
"""
|
||||
The signed ``X-Amzn-Bedrock-Request-Metadata`` header for the Invoke paths, which have no
|
||||
body field for request metadata.
|
||||
The signed ``X-Amzn-Bedrock-Request-Metadata`` header for an Invoke request whose body carries
|
||||
no ``requestMetadata`` of its own.
|
||||
|
||||
Returns the header names the proxy OWNS and, separately, the pairs to send. Ownership is
|
||||
reported whenever forwarding is enabled, including when nothing resolves, because a caller
|
||||
can suppress the resolver (omit the allow-listed fields, or send values that all fail
|
||||
Bedrock's rules) and an owned-but-empty result must still evict the caller's header rather
|
||||
than fall back to it.
|
||||
than fall back to it. Only the leading pairs that fit Bedrock's header limit are sent, so
|
||||
identity always survives and long ``spend_logs_metadata`` pairs drop off the end.
|
||||
"""
|
||||
if not bedrock_request_metadata_is_owned():
|
||||
return frozenset(), ()
|
||||
resolved: Final = resolve_bedrock_request_metadata(litellm_params)
|
||||
if resolved is None:
|
||||
fitting: Final = {} if resolved is None else _longest_fitting_prefix(resolved)
|
||||
if not fitting:
|
||||
return _OWNED_HEADER_NAMES, ()
|
||||
return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),)
|
||||
return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, _header_value(fitting)),)
|
||||
|
||||
|
||||
def merge_bedrock_invoke_headers(
|
||||
|
|
|
|||
|
|
@ -1087,6 +1087,16 @@ def _guardrail_refusal_status(error: Exception) -> int | None:
|
|||
return None
|
||||
|
||||
|
||||
def _aws_error_type(error: Exception, handled: ProxyException) -> str | None:
|
||||
"""The AWS error type to answer the handled error with: ValidationException for a guardrail block, else the type
|
||||
the error carried, as LiteLLM's own Bedrock request checks set it. None when the shared handler changed the
|
||||
guardrail block's status."""
|
||||
refusal_status: Final = _guardrail_refusal_status(error)
|
||||
if refusal_status is not None:
|
||||
return "ValidationException" if str(refusal_status) == handled.code else None
|
||||
return next((value for name, value in handled.headers.items() if name.lower() == AWS_ERROR_TYPE_HEADER), None)
|
||||
|
||||
|
||||
def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None:
|
||||
"""Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error
|
||||
type."""
|
||||
|
|
@ -1111,9 +1121,9 @@ async def bedrock_passthrough_error_response(
|
|||
user_api_key_dict: UserAPIKeyAuth,
|
||||
proxy_logging_obj: ProxyLoggingType,
|
||||
) -> Response:
|
||||
"""Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error
|
||||
Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the
|
||||
shared handler raises it."""
|
||||
"""Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block, an error
|
||||
Bedrock returned, or an error carrying an AWS error type the way Bedrock does, so an AWS SDK reads its code and
|
||||
message. Anything else is raised as the shared handler raises it."""
|
||||
try:
|
||||
await processor.handle_llm_api_exception(
|
||||
e=error,
|
||||
|
|
@ -1126,9 +1136,9 @@ async def bedrock_passthrough_error_response(
|
|||
raise
|
||||
return response
|
||||
except ProxyException as handled:
|
||||
refusal_status: Final = _guardrail_refusal_status(error)
|
||||
if refusal_status is not None and str(refusal_status) == handled.code:
|
||||
raise AwsErrorTypeProxyException(handled, "ValidationException") from handled
|
||||
aws_error_type: Final = _aws_error_type(error, handled)
|
||||
if aws_error_type is not None:
|
||||
raise AwsErrorTypeProxyException(handled, aws_error_type) from handled
|
||||
raise
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,316 @@
|
|||
import base64
|
||||
import json
|
||||
import uuid
|
||||
from urllib.parse import unquote
|
||||
from collections.abc import Iterator, Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
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() + 120)
|
||||
|
||||
_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0"
|
||||
_EVENT_STREAM: Final = "application/vnd.amazon.eventstream"
|
||||
_METADATA_HEADER: Final = "x-amzn-bedrock-request-metadata"
|
||||
_HEADER_MAX_CHARS: Final = 8500
|
||||
_LONG_SPEND_LOGS_METADATA: Final = tuple((f"{index:02d}".ljust(256, "k"), "\n" * 256) for index in range(12))
|
||||
_CONVERSE_REPLY: Final[Mapping[str, JsonValue]] = {
|
||||
"output": {"message": {"role": "assistant", "content": [{"text": "ok"}]}},
|
||||
"stopReason": "end_turn",
|
||||
"usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6},
|
||||
"metrics": {"latencyMs": 3},
|
||||
}
|
||||
_INVOKE_REPLY: Final[Mapping[str, JsonValue]] = {
|
||||
"id": "msg_bedrock",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude",
|
||||
"content": [{"type": "text", "text": "ok"}],
|
||||
"stop_reason": "end_turn",
|
||||
"stop_sequence": None,
|
||||
"usage": {"input_tokens": 5, "output_tokens": 1},
|
||||
}
|
||||
_CONVERSE_STREAM_FRAMES: Final = (
|
||||
_aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"),
|
||||
_aws_event_frame("messageStop", {"stopReason": "end_turn"}, "sc", "u"),
|
||||
_aws_event_frame("metadata", {"usage": {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6}}, "sc", "u"),
|
||||
)
|
||||
|
||||
|
||||
def _bedrock_peer(request: Request) -> Reply:
|
||||
if request.target.endswith("/converse-stream"):
|
||||
return Reply(chunks=_CONVERSE_STREAM_FRAMES, content_type=_EVENT_STREAM)
|
||||
if request.target.endswith("/invoke-with-response-stream"):
|
||||
chunk: Final = {"bytes": base64.b64encode(json.dumps(_INVOKE_REPLY).encode()).decode()}
|
||||
return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM)
|
||||
if request.target.endswith("/converse"):
|
||||
return Reply(body=json.dumps(_CONVERSE_REPLY).encode())
|
||||
return Reply(body=json.dumps(_INVOKE_REPLY).encode())
|
||||
|
||||
|
||||
def _messages(marker: str) -> Sequence[JsonValue]:
|
||||
return [{"role": "user", "content": [{"text": marker}]}]
|
||||
|
||||
|
||||
def _invoke_body(marker: str) -> Mapping[str, JsonValue]:
|
||||
return {
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 16,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
}
|
||||
|
||||
|
||||
def _signed_header_names(request: Request) -> frozenset[str]:
|
||||
return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";"))
|
||||
|
||||
|
||||
def _header_value(pairs: tuple[tuple[str, str], ...]) -> str:
|
||||
return json.dumps(dict(pairs), separators=(",", ":"))
|
||||
|
||||
|
||||
def _send_with_caller_metadata(
|
||||
runtime: object, action: str, model: str, marker: str, caller_metadata: Mapping[str, JsonValue]
|
||||
) -> None:
|
||||
if action == "invoke":
|
||||
getattr(runtime, "invoke_model")(
|
||||
modelId=model, body=json.dumps({**_invoke_body(marker), "requestMetadata": caller_metadata})
|
||||
)
|
||||
return
|
||||
getattr(runtime, action.replace("-", "_"))(
|
||||
modelId=model, messages=_messages(marker), requestMetadata=caller_metadata
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Rig:
|
||||
proxy: Gateway
|
||||
wire: Wire
|
||||
strict_model: str
|
||||
dropping_model: str
|
||||
chat_model: str
|
||||
key: str
|
||||
identity: Mapping[str, str]
|
||||
long_key: str
|
||||
long_identity: tuple[tuple[str, str], ...]
|
||||
|
||||
def runtime(self, key: str, monkeypatch: pytest.MonkeyPatch): # noqa: ANN201 # boto3 clients are untyped
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key)
|
||||
return 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},
|
||||
parameter_validation=False,
|
||||
),
|
||||
)
|
||||
|
||||
def received(self, marker: str) -> tuple[Request, ...]:
|
||||
return tuple(request for request in self.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]
|
||||
strict_model: Final = f"bedrock-metadata-{suffix}"
|
||||
dropping_model: Final = f"bedrock-metadata-dropping-{suffix}"
|
||||
chat_model: Final = f"bedrock-metadata-chat-{suffix}"
|
||||
directory: Final = tmp_path_factory.mktemp("bedrock-request-metadata")
|
||||
base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
with wire_server(_bedrock_peer) as wire:
|
||||
deployment: Final = {
|
||||
"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,
|
||||
}
|
||||
config: Final = {
|
||||
**base_config,
|
||||
"model_list": [
|
||||
*base_config.get("model_list", []),
|
||||
{"model_name": strict_model, "litellm_params": deployment},
|
||||
{"model_name": dropping_model, "litellm_params": {**deployment, "drop_params": True}},
|
||||
{
|
||||
"model_name": chat_model,
|
||||
"litellm_params": {**deployment, "model": f"bedrock/converse/{_MODEL_ID}", "drop_params": True},
|
||||
},
|
||||
],
|
||||
"litellm_settings": {
|
||||
**base_config.get("litellm_settings", {}),
|
||||
"drop_params": False,
|
||||
"bedrock_request_metadata_fields": [
|
||||
"user_api_key_alias",
|
||||
"user_api_key_team_alias",
|
||||
"spend_logs_metadata",
|
||||
],
|
||||
},
|
||||
}
|
||||
config_path: Final = directory / "bedrock-request-metadata.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,
|
||||
):
|
||||
team_alias: Final = f"platform-{suffix}"
|
||||
team_id: Final = scenario.team(team_alias=team_alias)
|
||||
key_alias: Final = f"prod-key-{suffix}"
|
||||
long_key_alias: Final = f"long-key-{suffix}"
|
||||
yield Rig(
|
||||
proxy=proxy,
|
||||
wire=wire,
|
||||
strict_model=strict_model,
|
||||
dropping_model=dropping_model,
|
||||
chat_model=chat_model,
|
||||
key=scenario.key(
|
||||
key_alias=key_alias,
|
||||
team_id=team_id,
|
||||
metadata={"spend_logs_metadata": {"agent": "triage-bot"}},
|
||||
),
|
||||
identity={
|
||||
"user_api_key_alias": key_alias,
|
||||
"user_api_key_team_alias": team_alias,
|
||||
"agent": "triage-bot",
|
||||
},
|
||||
long_key=scenario.key(
|
||||
key_alias=long_key_alias,
|
||||
team_id=team_id,
|
||||
metadata={"spend_logs_metadata": dict(_LONG_SPEND_LOGS_METADATA)},
|
||||
),
|
||||
long_identity=(
|
||||
("user_api_key_alias", long_key_alias),
|
||||
("user_api_key_team_alias", team_alias),
|
||||
*_LONG_SPEND_LOGS_METADATA,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_boto3_passthrough_calls_carry_the_key_identity_in_signed_request_metadata(
|
||||
rig: Rig, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
marker: Final = f"who-is-calling-{uuid.uuid4().hex}"
|
||||
forged: Final = {"user_api_key_team_alias": "forged", "ticket": "T-42", "bad!": "x"}
|
||||
runtime: Final = rig.runtime(rig.key, monkeypatch)
|
||||
model: Final = rig.dropping_model
|
||||
|
||||
converse: Final = runtime.converse(modelId=model, messages=_messages(marker), requestMetadata=forged)
|
||||
converse_stream: Final = [
|
||||
next(iter(event))
|
||||
for event in runtime.converse_stream(modelId=model, messages=_messages(marker), requestMetadata=forged)[
|
||||
"stream"
|
||||
]
|
||||
]
|
||||
invoke: Final = json.loads(
|
||||
runtime.invoke_model(modelId=model, body=json.dumps(_invoke_body(marker)))["body"].read()
|
||||
)
|
||||
invoke_stream: Final = [
|
||||
json.loads(event["chunk"]["bytes"])
|
||||
for event in runtime.invoke_model_with_response_stream(modelId=model, body=json.dumps(_invoke_body(marker)))[
|
||||
"body"
|
||||
]
|
||||
]
|
||||
assert converse["output"] == _CONVERSE_REPLY["output"], converse
|
||||
assert converse_stream == ["messageStart", "messageStop", "metadata"], converse_stream
|
||||
assert invoke == _INVOKE_REPLY, invoke
|
||||
assert invoke_stream == [_INVOKE_REPLY], invoke_stream
|
||||
|
||||
received: Final = rig.received(marker)
|
||||
assert [request.target for request in received] == [
|
||||
f"/model/{_MODEL_ID}/converse",
|
||||
f"/model/{_MODEL_ID}/converse-stream",
|
||||
f"/model/{_MODEL_ID}/invoke",
|
||||
f"/model/{_MODEL_ID}/invoke-with-response-stream",
|
||||
], received
|
||||
converse_requests: Final = received[:2]
|
||||
invoke_requests: Final = received[2:]
|
||||
assert [json.loads(request.body) for request in converse_requests] == [
|
||||
{"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}},
|
||||
{"messages": _messages(marker), "requestMetadata": {**rig.identity, "ticket": "T-42"}},
|
||||
], [request.body for request in converse_requests]
|
||||
assert [_METADATA_HEADER in request.headers for request in converse_requests] == [False, False]
|
||||
assert [json.loads(request.body) for request in invoke_requests] == [_invoke_body(marker), _invoke_body(marker)]
|
||||
assert [json.loads(request.headers[_METADATA_HEADER]) for request in invoke_requests] == [
|
||||
rig.identity,
|
||||
rig.identity,
|
||||
], [request.headers for request in invoke_requests]
|
||||
assert all(_METADATA_HEADER in _signed_header_names(request) for request in invoke_requests), [
|
||||
request.headers["authorization"] for request in invoke_requests
|
||||
]
|
||||
|
||||
|
||||
def test_invoke_metadata_too_long_for_the_header_signs_the_leading_pairs_that_fit(
|
||||
rig: Rig, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
marker: Final = f"long-key-{uuid.uuid4().hex}"
|
||||
runtime: Final = rig.runtime(rig.long_key, monkeypatch)
|
||||
|
||||
invoke: Final = json.loads(
|
||||
runtime.invoke_model(modelId=rig.strict_model, body=json.dumps(_invoke_body(marker)))["body"].read()
|
||||
)
|
||||
|
||||
assert invoke == _INVOKE_REPLY, invoke
|
||||
(request,) = rig.received(marker)
|
||||
assert json.loads(request.body) == _invoke_body(marker)
|
||||
header: Final = request.headers[_METADATA_HEADER]
|
||||
sent_pairs: Final = len(json.loads(header))
|
||||
assert tuple(json.loads(header).items()) == rig.long_identity[:sent_pairs], header
|
||||
assert 2 <= sent_pairs < len(rig.long_identity), header
|
||||
assert len(header) <= _HEADER_MAX_CHARS < len(_header_value(rig.long_identity[: sent_pairs + 1]))
|
||||
assert _METADATA_HEADER in _signed_header_names(request), request.headers["authorization"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("action", "caller_metadata", "error_code"),
|
||||
[
|
||||
pytest.param("converse", {"ticket": 42}, "SerializationException", id="converse-number-value"),
|
||||
pytest.param("converse-stream", {"bad!": "x"}, "ValidationException", id="converse-stream-bad-key"),
|
||||
pytest.param("invoke", {"bad!": "x"}, "ValidationException", id="invoke-bad-key"),
|
||||
],
|
||||
)
|
||||
def test_malformed_caller_request_metadata_reaches_boto3_as_the_bedrock_error_without_an_upstream_call(
|
||||
rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str, caller_metadata: Mapping[str, JsonValue], error_code: str
|
||||
) -> None:
|
||||
marker: Final = f"malformed-{uuid.uuid4().hex}"
|
||||
runtime: Final = rig.runtime(rig.key, monkeypatch)
|
||||
|
||||
with pytest.raises(ClientError) as raised:
|
||||
_send_with_caller_metadata(runtime, action, rig.strict_model, marker, caller_metadata)
|
||||
|
||||
assert raised.value.response["Error"]["Code"] == error_code, raised.value.response
|
||||
assert raised.value.response["ResponseMetadata"]["HTTPStatusCode"] == 400, raised.value.response
|
||||
assert rig.received(marker) == ()
|
||||
|
||||
|
||||
def test_chat_on_a_dropping_converse_deployment_sends_only_the_valid_caller_pairs(rig: Rig) -> None:
|
||||
marker: Final = f"chat-{uuid.uuid4().hex}"
|
||||
|
||||
response: Final = rig.proxy.client.post(
|
||||
"/chat/completions",
|
||||
json={
|
||||
"model": rig.chat_model,
|
||||
"messages": [{"role": "user", "content": marker}],
|
||||
"requestMetadata": {"ticket": "T-42", "bad!": "x"},
|
||||
},
|
||||
headers={"Authorization": f"Bearer {rig.key}"},
|
||||
timeout=case_timeout_seconds(),
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["choices"][0]["message"]["content"] == "ok", response.text
|
||||
(request,) = rig.received(marker)
|
||||
assert unquote(request.target) == f"/model/{_MODEL_ID}/converse", request.target
|
||||
assert json.loads(request.body)["requestMetadata"] == {**rig.identity, "ticket": "T-42"}, request.body
|
||||
|
|
@ -115,3 +115,22 @@ def test_tool_config_injection_point_is_neither_placed_nor_credited(local_model_
|
|||
assert [tool["toolSpec"]["name"] for tool in request["toolConfig"]["tools"]] == ["f"]
|
||||
assert "cachePoint" not in json.dumps(request)
|
||||
assert GATEWAY_INJECTED_CACHE_METADATA_KEY not in metadata
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["b c", "bad!"], ids=["nbsp", "outside-bedrock-charset"])
|
||||
def test_caller_request_metadata_is_ignored_because_invoke_model_has_no_field_for_it(local_model_cost_map, value):
|
||||
caller_metadata = {"ticket": value}
|
||||
optional_params = AmazonInvokeNovaConfig().map_openai_params(
|
||||
non_default_params={"max_tokens": 8, "requestMetadata": caller_metadata},
|
||||
optional_params={},
|
||||
model=MODEL,
|
||||
drop_params=False,
|
||||
)
|
||||
|
||||
request = _transform_request(
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
optional_params={**optional_params, "requestMetadata": caller_metadata},
|
||||
)
|
||||
|
||||
assert "requestMetadata" not in request
|
||||
assert request["inferenceConfig"] == {"maxTokens": 8}
|
||||
|
|
|
|||
|
|
@ -4199,7 +4199,7 @@ def test_request_metadata_key_constraints():
|
|||
# Test empty key
|
||||
invalid_metadata = {"": "value"}
|
||||
|
||||
with pytest.raises(Exception, match=r"(?i)key length|empty"):
|
||||
with pytest.raises(Exception, match=r"key '' must be 1-256 characters"):
|
||||
config._transform_request(
|
||||
model="anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
messages=messages,
|
||||
|
|
@ -4241,6 +4241,18 @@ def test_request_metadata_value_constraints():
|
|||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata", [{"ticket": "a\u00a0b"}, {"a\u2003b": "ticket"}])
|
||||
def test_request_metadata_rejects_whitespace_outside_ascii(metadata):
|
||||
with pytest.raises(litellm.exceptions.BadRequestError, match=r"must be [01]-256 characters of"):
|
||||
AmazonConverseConfig()._transform_request(
|
||||
model="anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
messages=[{"role": "user", "content": "Hello!"}],
|
||||
optional_params={"requestMetadata": metadata},
|
||||
litellm_params={},
|
||||
headers={},
|
||||
)
|
||||
|
||||
|
||||
def test_request_metadata_character_pattern():
|
||||
"""Test character pattern validation for keys and values."""
|
||||
config = AmazonConverseConfig()
|
||||
|
|
|
|||
|
|
@ -1760,6 +1760,23 @@ class TestBedrockBatchNonChatEndpointRecords:
|
|||
assert "metadata" not in model_input
|
||||
assert "requestMetadata" not in model_input
|
||||
|
||||
def test_nova_record_request_metadata_stays_out_of_the_bedrock_body(self):
|
||||
model_input = self._transform(
|
||||
{
|
||||
"custom_id": "4d",
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": self.NOVA_MODEL,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"requestMetadata": {"ticket": "b\u00a0c"},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
assert "requestMetadata" not in model_input
|
||||
assert model_input["messages"] == [{"role": "user", "content": [{"text": "hi"}]}]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"body",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@ import litellm
|
|||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.passthrough.transformation import (
|
||||
BedrockPassthroughConfig,
|
||||
is_bedrock_streaming_endpoint,
|
||||
|
|
@ -898,3 +899,293 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint():
|
|||
collector.add(_event_frame("messageStart", {"role": "assistant"}))
|
||||
|
||||
assert collector.build_logged_response(_converse_stream_logging_obj()) is None
|
||||
|
||||
|
||||
REQUEST_METADATA_HEADER = "x-amzn-bedrock-request-metadata"
|
||||
REQUEST_METADATA_FIELDS = ["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"]
|
||||
PROXY_IDENTITY = {
|
||||
"user_api_key_alias": "prod-key",
|
||||
"user_api_key_team_alias": "platform",
|
||||
"spend_logs_metadata": {"agent": "triage-bot"},
|
||||
}
|
||||
EXPECTED_REQUEST_METADATA = {
|
||||
"user_api_key_alias": "prod-key",
|
||||
"user_api_key_team_alias": "platform",
|
||||
"agent": "triage-bot",
|
||||
}
|
||||
CONVERSE_BODY = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]}
|
||||
INVOKE_BODY = {
|
||||
"anthropic_version": "bedrock-2023-05-31",
|
||||
"max_tokens": 8,
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
}
|
||||
|
||||
|
||||
async def _sent_passthrough_request(
|
||||
action: str, body: dict[str, object], identity: dict[str, object], **kwargs: object
|
||||
) -> httpx.Request:
|
||||
requests: Final[list[httpx.Request]] = []
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
requests.append(request)
|
||||
return httpx.Response(200, json={"ok": True}, request=request)
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "bedrock-pt",
|
||||
"litellm_params": {
|
||||
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
||||
"aws_access_key_id": "AKIAROUTERKEY",
|
||||
"aws_secret_access_key": "router-secret",
|
||||
"aws_region_name": "us-west-2",
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
await router.allm_passthrough_route(
|
||||
model="bedrock-pt",
|
||||
method="POST",
|
||||
endpoint=f"/model/bedrock-pt/{action}",
|
||||
data=body,
|
||||
client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)),
|
||||
litellm_metadata=identity,
|
||||
**kwargs,
|
||||
)
|
||||
return requests[0]
|
||||
|
||||
|
||||
def _signed_header_names(request: httpx.Request) -> frozenset[str]:
|
||||
return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";"))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("action", "query_params"),
|
||||
[("converse", None), ("converse-stream", None), ("converse/", None), ("converse", {"trace": "1"})],
|
||||
)
|
||||
async def test_converse_passthrough_signs_proxy_identity_into_request_metadata(
|
||||
action: str, query_params: dict[str, str] | None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}}
|
||||
|
||||
request: Final = await _sent_passthrough_request(
|
||||
action, caller_body, PROXY_IDENTITY, request_query_params=query_params
|
||||
)
|
||||
|
||||
assert json.loads(request.content) == {
|
||||
**CONVERSE_BODY,
|
||||
"requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"},
|
||||
}
|
||||
assert caller_body["requestMetadata"] == {"user_api_key_team_alias": "forged", "ticket": "T-42"}
|
||||
assert REQUEST_METADATA_HEADER not in request.headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
|
||||
async def test_invoke_passthrough_sends_proxy_identity_as_a_signed_header(
|
||||
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,
|
||||
request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'},
|
||||
)
|
||||
|
||||
assert json.loads(request.headers[REQUEST_METADATA_HEADER]) == EXPECTED_REQUEST_METADATA
|
||||
assert REQUEST_METADATA_HEADER in _signed_header_names(request)
|
||||
assert json.loads(request.content) == INVOKE_BODY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
|
||||
async def test_invoke_passthrough_keeps_caller_body_request_metadata_in_the_body(
|
||||
action: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
caller_body: Final = {**INVOKE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged", "ticket": "T-42"}}
|
||||
|
||||
request: Final = await _sent_passthrough_request(
|
||||
action,
|
||||
caller_body,
|
||||
PROXY_IDENTITY,
|
||||
request_headers={f"x-pass-{REQUEST_METADATA_HEADER.upper()}": '{"user_api_key_team_alias":"forged"}'},
|
||||
)
|
||||
|
||||
assert json.loads(request.content) == {
|
||||
**INVOKE_BODY,
|
||||
"requestMetadata": {**EXPECTED_REQUEST_METADATA, "ticket": "T-42"},
|
||||
}
|
||||
assert REQUEST_METADATA_HEADER not in request.headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("setting", [None, []])
|
||||
async def test_passthrough_requests_are_untouched_when_forwarding_is_off(
|
||||
setting: list[str] | None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", setting)
|
||||
caller_body: Final = {**CONVERSE_BODY, "requestMetadata": {"user_api_key_team_alias": "forged"}}
|
||||
|
||||
converse: Final = await _sent_passthrough_request("converse", caller_body, PROXY_IDENTITY)
|
||||
invoke: Final = await _sent_passthrough_request("invoke", INVOKE_BODY, PROXY_IDENTITY)
|
||||
|
||||
assert json.loads(converse.content) == caller_body
|
||||
assert REQUEST_METADATA_HEADER not in invoke.headers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)])
|
||||
async def test_passthrough_drops_reserved_caller_keys_when_nothing_resolves(
|
||||
action: str, body: dict[str, object], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged"}}
|
||||
|
||||
request: Final = await _sent_passthrough_request(action, caller_body, {})
|
||||
|
||||
assert json.loads(request.content) == body
|
||||
assert REQUEST_METADATA_HEADER not in request.headers
|
||||
|
||||
|
||||
MALFORMED_REQUEST_METADATA: Final = (
|
||||
pytest.param(
|
||||
{"user_api_key_team_alias": "forged", "bad!": "x"}, "key 'bad!' must be", "ValidationException", id="bad-key"
|
||||
),
|
||||
pytest.param({"ticket": "a\u00a0b"}, "value for key 'ticket' must be", "ValidationException", id="bad-value"),
|
||||
pytest.param("ticket=T-42", "must be a dictionary", "SerializationException", id="not-a-dictionary"),
|
||||
pytest.param({f"key_{index}": "v" for index in range(17)}, "maximum of 16 items", "ValidationException", id="17"),
|
||||
pytest.param({"ticket": 42}, "must be strings", "SerializationException", id="non-string-value"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["converse", "converse-stream"])
|
||||
@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA)
|
||||
async def test_malformed_converse_request_metadata_is_rejected_with_the_error_type_bedrock_returns(
|
||||
action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
|
||||
with pytest.raises(BedrockError, match=violation) as rejected:
|
||||
await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY)
|
||||
|
||||
assert rejected.value.status_code == 400
|
||||
assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == converse_error_type
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
|
||||
@pytest.mark.parametrize(("caller_metadata", "violation", "converse_error_type"), MALFORMED_REQUEST_METADATA)
|
||||
async def test_malformed_invoke_request_metadata_is_rejected_as_a_bedrock_validation_error(
|
||||
action: str, caller_metadata: object, violation: str, converse_error_type: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
|
||||
with pytest.raises(BedrockError, match=violation) as rejected:
|
||||
await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": caller_metadata}, PROXY_IDENTITY)
|
||||
|
||||
assert rejected.value.status_code == 400
|
||||
assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
|
||||
async def test_invoke_null_request_metadata_is_rejected_like_bedrock_does(
|
||||
action: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
|
||||
with pytest.raises(BedrockError, match="must be a dictionary") as rejected:
|
||||
await _sent_passthrough_request(action, {**INVOKE_BODY, "requestMetadata": None}, PROXY_IDENTITY)
|
||||
|
||||
assert rejected.value.status_code == 400
|
||||
assert httpx.Headers(rejected.value.headers).get("x-amzn-errortype") == "ValidationException"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["converse", "converse-stream"])
|
||||
async def test_converse_null_request_metadata_is_treated_as_absent_like_bedrock_does(
|
||||
action: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
|
||||
request: Final = await _sent_passthrough_request(action, {**CONVERSE_BODY, "requestMetadata": None}, PROXY_IDENTITY)
|
||||
|
||||
assert json.loads(request.content) == {**CONVERSE_BODY, "requestMetadata": EXPECTED_REQUEST_METADATA}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"])
|
||||
async def test_invoke_metadata_too_long_for_the_header_sends_the_leading_pairs_that_fit(
|
||||
action: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)}
|
||||
resolved: Final = tuple(
|
||||
{"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform", **long_pairs}.items()
|
||||
)
|
||||
|
||||
request: Final = await _sent_passthrough_request(
|
||||
action,
|
||||
INVOKE_BODY,
|
||||
{**PROXY_IDENTITY, "spend_logs_metadata": long_pairs},
|
||||
request_headers={f"x-pass-{REQUEST_METADATA_HEADER}": '{"user_api_key_team_alias":"forged"}'},
|
||||
)
|
||||
|
||||
header: Final = request.headers[REQUEST_METADATA_HEADER]
|
||||
sent: Final = json.loads(header)
|
||||
one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":"))
|
||||
assert sent == dict(resolved[: len(sent)])
|
||||
assert len(header) <= 8500 < len(one_more_pair)
|
||||
assert json.loads(request.content) == INVOKE_BODY
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("action", "body"), [("converse", CONVERSE_BODY), ("invoke", INVOKE_BODY)])
|
||||
@pytest.mark.parametrize("drop_globally", [True, False])
|
||||
async def test_malformed_caller_request_metadata_is_dropped_when_params_are_dropped(
|
||||
action: str, body: dict[str, object], drop_globally: bool, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", REQUEST_METADATA_FIELDS)
|
||||
monkeypatch.setattr(litellm, "drop_params", drop_globally)
|
||||
caller_body: Final = {**body, "requestMetadata": {"user_api_key_team_alias": "forged", "bad!": "x"}}
|
||||
|
||||
dropped: Final = await _sent_passthrough_request(
|
||||
action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True})
|
||||
)
|
||||
|
||||
assert json.loads(dropped.content) == {**body, "requestMetadata": EXPECTED_REQUEST_METADATA}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("action", "body", "caller_metadata"),
|
||||
[
|
||||
("converse", CONVERSE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}),
|
||||
("invoke", INVOKE_BODY, {"user_api_key_team_alias": "caller-team", "bad!": "x"}),
|
||||
("invoke", INVOKE_BODY, None),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("drop_globally", [True, False])
|
||||
async def test_passthrough_body_is_sent_as_is_when_params_are_dropped_and_forwarding_is_off(
|
||||
action: str,
|
||||
body: dict[str, object],
|
||||
caller_metadata: dict[str, str] | None,
|
||||
drop_globally: bool,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None)
|
||||
monkeypatch.setattr(litellm, "drop_params", drop_globally)
|
||||
caller_body: Final = {**body, "requestMetadata": caller_metadata}
|
||||
|
||||
sent: Final = await _sent_passthrough_request(
|
||||
action, caller_body, PROXY_IDENTITY, **({} if drop_globally else {"drop_params": True})
|
||||
)
|
||||
|
||||
assert json.loads(sent.content) == caller_body
|
||||
assert REQUEST_METADATA_HEADER not in sent.headers
|
||||
|
|
|
|||
|
|
@ -1,8 +1,10 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Callable, Mapping
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
|
|
@ -18,6 +20,7 @@ from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_tran
|
|||
)
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
BEDROCK_REQUEST_METADATA_HEADER,
|
||||
BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS,
|
||||
BEDROCK_REQUEST_METADATA_MAX_PAIRS,
|
||||
resolve_bedrock_request_metadata,
|
||||
)
|
||||
|
|
@ -128,6 +131,20 @@ def test_identity_survives_a_caller_filling_every_slot(
|
|||
)
|
||||
|
||||
|
||||
def test_a_key_sent_twice_keeps_the_caller_value_and_uses_one_slot(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
|
||||
caller: Final = {
|
||||
f"client_{index:02d}": "caller" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY) - 1)
|
||||
}
|
||||
spend_logs: Final = {**{key: "spend" for key in caller}, "extra": "spend"}
|
||||
|
||||
resolved: Final = resolve_bedrock_request_metadata(
|
||||
litellm_params("metadata", spend_logs_metadata=spend_logs, **IDENTITY), caller_metadata=caller
|
||||
)
|
||||
|
||||
assert resolved == {**IDENTITY, **caller, "extra": "spend"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field_order",
|
||||
[
|
||||
|
|
@ -158,7 +175,14 @@ def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(
|
|||
@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"])
|
||||
@pytest.mark.parametrize(
|
||||
"forged_key",
|
||||
["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"],
|
||||
[
|
||||
"user_api_key_team_alias",
|
||||
"user_api_key_org_alias",
|
||||
"user_api_key_hash",
|
||||
"USER_API_KEY_TEAM_ALIAS",
|
||||
"User_Api_Key_Hash",
|
||||
" user_api_key_user_id",
|
||||
],
|
||||
)
|
||||
def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(
|
||||
forged_key,
|
||||
|
|
@ -181,14 +205,21 @@ def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(
|
|||
assert "attacker-controlled" not in resolved.values()
|
||||
|
||||
|
||||
def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(monkeypatch: pytest.MonkeyPatch):
|
||||
@pytest.mark.parametrize(
|
||||
"team_alias",
|
||||
["O'Brien's team", "platform\u00a0team", "platform\u3000team", "x" * 256 + "\n"],
|
||||
ids=["apostrophe", "no-break-space", "ideographic-space", "257-chars-ending-in-newline"],
|
||||
)
|
||||
def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds(
|
||||
team_alias: str, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""A team alias with an apostrophe must not turn a working request into a 400 the moment
|
||||
an operator flips the setting on."""
|
||||
an operator flips the setting on. Bedrock also rejects non-ASCII whitespace (checked live 2026-10-16)."""
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
|
||||
params = litellm_params(
|
||||
"metadata",
|
||||
user_api_key_alias="prod-key",
|
||||
user_api_key_team_alias="O'Brien's team",
|
||||
user_api_key_team_alias=team_alias,
|
||||
user_api_key_user_email="x" * 300,
|
||||
)
|
||||
|
||||
|
|
@ -198,13 +229,54 @@ def test_identity_violating_the_character_class_is_dropped_and_the_request_succe
|
|||
assert body["messages"]
|
||||
|
||||
|
||||
def test_caller_supplied_violation_still_raises_bad_request(monkeypatch: pytest.MonkeyPatch):
|
||||
@pytest.mark.parametrize(
|
||||
"caller_metadata",
|
||||
[{"team": "O'Brien's team"}, {"k" * 256 + "\n": "v"}, {"team": "v" * 256 + "\n"}],
|
||||
ids=["bad-character", "257-char-key-ending-in-newline", "257-char-value-ending-in-newline"],
|
||||
)
|
||||
def test_caller_supplied_violation_still_raises_bad_request(
|
||||
caller_metadata: Mapping[str, str], monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
|
||||
|
||||
with pytest.raises(litellm.exceptions.BadRequestError):
|
||||
converse_body(
|
||||
litellm_params("metadata", **IDENTITY),
|
||||
{"requestMetadata": {"team": "O'Brien's team"}},
|
||||
converse_body(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_metadata})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
|
||||
@pytest.mark.parametrize("drop_globally", [True, False])
|
||||
@pytest.mark.parametrize(("fields", "forwarded_identity"), [(ALL_FIELDS, IDENTITY), (None, {})])
|
||||
def test_caller_violations_are_dropped_like_the_passthrough_routes_when_params_are_dropped(
|
||||
driver: Callable[[Mapping[str, object], Mapping[str, object]], Mapping[str, object]],
|
||||
drop_globally: bool,
|
||||
fields: list[str] | None,
|
||||
forwarded_identity: Mapping[str, str],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields)
|
||||
monkeypatch.setattr(litellm, "drop_params", drop_globally)
|
||||
params: Final = {**litellm_params("metadata", **IDENTITY), **({} if drop_globally else {"drop_params": True})}
|
||||
|
||||
body: Final = driver(params, {"requestMetadata": {"team": "O'Brien's team", "ticket": "T-42"}})
|
||||
|
||||
assert body["requestMetadata"] == {**forwarded_identity, "ticket": "T-42"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fields", [ALL_FIELDS, None])
|
||||
def test_mapped_caller_violations_are_dropped_when_params_are_dropped(
|
||||
fields: list[str] | None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", fields)
|
||||
caller: Final = {"team": "O'Brien's team", "ticket": "T-42"}
|
||||
|
||||
mapped: Final = AmazonConverseConfig().map_openai_params(
|
||||
non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=True
|
||||
)
|
||||
|
||||
assert mapped["requestMetadata"] == caller
|
||||
with pytest.raises(litellm.exceptions.BadRequestError):
|
||||
AmazonConverseConfig().map_openai_params(
|
||||
non_default_params={"requestMetadata": caller}, optional_params={}, model=MODEL, drop_params=False
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -320,6 +392,27 @@ def metadata_header_values(headers):
|
|||
return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
|
||||
def test_the_metadata_header_keeps_the_leading_pairs_that_fit_bedrocks_header_limit(
|
||||
driver: Callable[[Mapping[str, str], Mapping[str, object]], Mapping[str, str]],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", ALL_FIELDS)
|
||||
long_pairs: Final = {f"{index:02d}".ljust(256, "k"): "\n" * 256 for index in range(12)}
|
||||
params: Final = litellm_params("metadata", spend_logs_metadata=long_pairs, **IDENTITY)
|
||||
resolved: Final = tuple((resolve_bedrock_request_metadata(params) or {}).items())
|
||||
|
||||
headers: Final = driver({}, params)
|
||||
(header,) = (value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower())
|
||||
sent: Final = TypeAdapter(dict[str, str]).validate_json(header)
|
||||
one_more_pair: Final = json.dumps(dict(resolved[: len(sent) + 1]), separators=(",", ":"))
|
||||
|
||||
assert len(header) <= BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS
|
||||
assert tuple(sent.items()) == resolved[: len(sent)]
|
||||
assert IDENTITY.items() <= sent.items()
|
||||
assert len(one_more_pair) > BEDROCK_REQUEST_METADATA_HEADER_MAX_CHARS
|
||||
|
||||
|
||||
def test_converse_still_sets_the_bearer_authorization_header(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Converse owns the metadata header now, and that must not disturb the api_key path its
|
||||
validate_environment existed for. Closing the forgery hole cannot break authentication."""
|
||||
|
|
|
|||
|
|
@ -33,6 +33,7 @@ from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException,
|
|||
from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
from litellm.passthrough.main import raise_for_provider_status
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe
|
||||
from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
|
|
@ -2137,6 +2138,47 @@ class TestBedrockPassthroughAwsErrors:
|
|||
assert exc_info.value.aws_error_type == "ValidationException"
|
||||
assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_request_check_is_raised_with_the_aws_error_type_it_carries(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
error: Final = BedrockError(
|
||||
status_code=400,
|
||||
message="requestMetadata can contain a maximum of 16 items",
|
||||
headers={"x-amzn-errortype": "ValidationException"},
|
||||
)
|
||||
|
||||
with pytest.raises(AwsErrorTypeProxyException) as exc_info:
|
||||
await self._answer(error)
|
||||
|
||||
assert exc_info.value.aws_error_type == "ValidationException"
|
||||
assert "requestMetadata can contain a maximum of 16 items" in exc_info.value.error.message
|
||||
assert exc_info.value.error.code == "400"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_error_without_an_aws_error_type_keeps_the_proxy_shape(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await self._answer(BedrockError(status_code=400, message="bad request"))
|
||||
|
||||
assert exc_info.value.code == "400"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_error_with_a_non_ascii_header_value_keeps_the_proxy_shape(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
error: Final = HTTPException(status_code=403, detail="model not allowed", headers={"x-model-group": "modèle"})
|
||||
|
||||
with pytest.raises(ProxyException) as exc_info:
|
||||
await self._answer(error)
|
||||
|
||||
assert (exc_info.value.code, exc_info.value.headers.get("x-model-group")) == ("403", "modèle")
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("error", "message", "code"),
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue