mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
feat(bedrock): forward LiteLLM identity and metadata into Bedrock requestMetadata (#36861)
Adds an opt-in operator allow-list, litellm_settings::bedrock_request_metadata_fields, that forwards LiteLLM key, team and end-user identity plus client spend_logs_metadata into Bedrock request metadata so Bedrock spend can be grouped in AWS Cost Explorer. Covers all three Bedrock surfaces: the Converse body requestMetadata field, and a signed X-Amzn-Bedrock-Request-Metadata header on Invoke chat completions and on Invoke /v1/messages, where the header is the only viable leg. The resolver reads both metadata variable names, reserves the whole user_api_key_ prefix against caller-supplied keys, caps the client slot budget explicitly at 16 minus the reserved count, and drops rather than rejects auto-injected values that violate Bedrock constraints. Caller-supplied requestMetadata keeps its existing 400 semantics. The request-metadata field and header are proxy-owned whenever forwarding is enabled. A caller-supplied value, reachable through the generic extra_headers passthrough, is dropped unconditionally and compared case-insensitively, and is replaced only by the proxy's own value, so identity in the AWS billing record cannot be forged. Absence of a resolved value still means absence on the wire rather than a fallback to the caller's. The guardrail headers keep their existing no-displace behaviour.
This commit is contained in:
Yassin Kortam
GitHub
parent
1139012b45
commit
ee08b63657
7 changed files with 673 additions and 13 deletions
|
|
@ -26,6 +26,7 @@ def _dev_env_hot_reload_enabled() -> bool:
|
|||
if os.getenv("LITELLM_MODE", "DEV") == "DEV":
|
||||
_dotenv.load_dotenv(override=_dev_env_hot_reload_enabled())
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import (
|
||||
Any,
|
||||
Callable,
|
||||
|
|
@ -217,6 +218,9 @@ add_user_information_to_llm_headers: Optional[bool] = (
|
|||
overwrite_user_with_key_hash: bool = (
|
||||
False # force the outgoing `user` param to the hashed api key, so providers see a stable, tamper-proof id
|
||||
)
|
||||
bedrock_request_metadata_fields: Optional[Sequence[str]] = (
|
||||
None # allow-list of `user_api_key_*` fields (+ `spend_logs_metadata`) sent as Bedrock `requestMetadata`
|
||||
)
|
||||
store_audit_logs = False # Enterprise feature, allow users to see audit logs
|
||||
skip_system_message_in_guardrail: bool = False
|
||||
skip_tool_message_in_guardrail: bool = False
|
||||
|
|
|
|||
|
|
@ -39,6 +39,12 @@ from litellm.llms.anthropic.chat.transformation import (
|
|||
AnthropicConfig,
|
||||
)
|
||||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
bedrock_request_metadata_is_owned,
|
||||
merge_bedrock_invoke_headers,
|
||||
resolve_bedrock_request_metadata,
|
||||
)
|
||||
from litellm.types.llms.bedrock import *
|
||||
from litellm.types.llms.openai import (
|
||||
AllMessageValues,
|
||||
|
|
@ -1652,6 +1658,13 @@ class AmazonConverseConfig(BaseConfig):
|
|||
user_continue_message=litellm_params.pop("user_continue_message", None),
|
||||
)
|
||||
|
||||
request_metadata: Final = resolve_bedrock_request_metadata(
|
||||
litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata")
|
||||
)
|
||||
if bedrock_request_metadata_is_owned():
|
||||
_data.pop("requestMetadata", None)
|
||||
if request_metadata is not None:
|
||||
_data["requestMetadata"] = request_metadata
|
||||
data: Final[RequestObject] = {"messages": bedrock_messages, **_data}
|
||||
|
||||
return data
|
||||
|
|
@ -1705,6 +1718,13 @@ class AmazonConverseConfig(BaseConfig):
|
|||
user_continue_message=litellm_params.pop("user_continue_message", None),
|
||||
)
|
||||
|
||||
request_metadata: Final = resolve_bedrock_request_metadata(
|
||||
litellm_params=litellm_params, caller_metadata=_data.get("requestMetadata")
|
||||
)
|
||||
if bedrock_request_metadata_is_owned():
|
||||
_data.pop("requestMetadata", None)
|
||||
if request_metadata is not None:
|
||||
_data["requestMetadata"] = request_metadata
|
||||
data: Final[RequestObject] = {"messages": bedrock_messages, **_data}
|
||||
|
||||
return data
|
||||
|
|
@ -2258,7 +2278,8 @@ class AmazonConverseConfig(BaseConfig):
|
|||
) -> dict:
|
||||
if api_key:
|
||||
headers["Authorization"] = f"Bearer {api_key}"
|
||||
return headers
|
||||
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
|
||||
return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names)
|
||||
|
||||
def should_fake_stream(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -13,6 +13,10 @@ import httpx
|
|||
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
merge_bedrock_invoke_headers,
|
||||
)
|
||||
from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig
|
||||
from litellm.passthrough.utils import CommonUtils
|
||||
from litellm.types.llms.openai import AllMessageValues
|
||||
|
|
@ -169,9 +173,12 @@ class AmazonBedrockOpenAIConfig(OpenAIGPTConfig, BaseAWSLLM):
|
|||
"""
|
||||
Validate the environment and return headers.
|
||||
|
||||
For Bedrock, we don't need Bearer token auth since we use AWS SigV4.
|
||||
For Bedrock, we don't need Bearer token auth since we use AWS SigV4. This path signs the
|
||||
same ``/model/{id}/invoke`` endpoint as ``AmazonInvokeConfig``, so it owns the request
|
||||
metadata header on the same terms rather than letting a caller supply it.
|
||||
"""
|
||||
return headers
|
||||
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
|
||||
return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names)
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BedrockError:
|
||||
"""Return the appropriate error class for Bedrock."""
|
||||
|
|
|
|||
|
|
@ -20,6 +20,10 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
|
||||
from litellm.llms.bedrock.chat.invoke_handler import make_call, make_sync_call
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
merge_bedrock_invoke_headers,
|
||||
)
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
AsyncHTTPHandler,
|
||||
HTTPHandler,
|
||||
|
|
@ -417,15 +421,13 @@ class AmazonInvokeConfig(BaseConfig, BaseAWSLLM):
|
|||
api_base: str | None = None,
|
||||
) -> dict:
|
||||
raw_guardrail_config: Final = optional_params.pop("guardrailConfig", None)
|
||||
if raw_guardrail_config is None:
|
||||
return headers
|
||||
existing_header_names: Final = frozenset(name.lower() for name in headers)
|
||||
guardrail_headers: Final = {
|
||||
name: value
|
||||
for name, value in _bedrock_invoke_guardrail_headers(raw_guardrail_config).items()
|
||||
if name.lower() not in existing_header_names
|
||||
}
|
||||
return {**headers, **guardrail_headers}
|
||||
guardrail_headers: Final = (
|
||||
()
|
||||
if raw_guardrail_config is None
|
||||
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)
|
||||
|
||||
def get_error_class(self, error_message: str, status_code: int, headers: dict | httpx.Headers) -> BaseLLMException:
|
||||
return BedrockError(status_code=status_code, message=error_message)
|
||||
|
|
|
|||
|
|
@ -37,6 +37,10 @@ from litellm.llms.bedrock.common_utils import (
|
|||
normalize_tool_input_schema_types_for_bedrock_invoke,
|
||||
pop_bedrock_invoke_output_config_format,
|
||||
)
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
bedrock_request_metadata_headers,
|
||||
merge_bedrock_invoke_headers,
|
||||
)
|
||||
from litellm.types.llms.anthropic import (
|
||||
ANTHROPIC_BETA_HEADER_VALUES,
|
||||
ANTHROPIC_TOOL_SEARCH_BETA_HEADER,
|
||||
|
|
@ -89,7 +93,8 @@ class AmazonAnthropicClaudeMessagesConfig(
|
|||
api_key: str | None = None,
|
||||
api_base: str | None = None,
|
||||
) -> tuple[dict, str | None]:
|
||||
return headers, api_base
|
||||
owned_names, metadata_headers = bedrock_request_metadata_headers(litellm_params)
|
||||
return merge_bedrock_invoke_headers(headers, (), metadata_headers, owned_names), api_base
|
||||
|
||||
def sign_request(
|
||||
self,
|
||||
|
|
|
|||
199
litellm/llms/bedrock/request_metadata.py
Normal file
199
litellm/llms/bedrock/request_metadata.py
Normal file
|
|
@ -0,0 +1,199 @@
|
|||
"""
|
||||
Resolve AWS Bedrock ``requestMetadata`` from LiteLLM proxy identity and caller metadata.
|
||||
|
||||
Bedrock attaches request metadata to CloudTrail records and to the dimension AWS Cost
|
||||
Explorer groups on, so everything here is opt-in: nothing is forwarded unless the operator
|
||||
sets ``litellm.bedrock_request_metadata_fields`` (``litellm_settings`` on the proxy).
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import litellm
|
||||
|
||||
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"
|
||||
|
||||
_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}$")
|
||||
_OWNED_HEADER_NAMES: Final[frozenset[str]] = frozenset((BEDROCK_REQUEST_METADATA_HEADER.lower(),))
|
||||
|
||||
|
||||
def _is_forwardable(key: str, value: str) -> bool:
|
||||
return _KEY_PATTERN.match(key) is not None and _VALUE_PATTERN.match(value) is not None
|
||||
|
||||
|
||||
def _text_pairs(source: object) -> tuple[tuple[str, str], ...]:
|
||||
if not isinstance(source, Mapping):
|
||||
return ()
|
||||
return tuple((key, value) for key, value in source.items() if isinstance(key, str) and isinstance(value, str))
|
||||
|
||||
|
||||
def _allowed_fields() -> tuple[str, ...]:
|
||||
"""
|
||||
The operator allow-list, deduplicated so a field repeated in config cannot consume a second
|
||||
reserved slot and shrink the client budget for nothing. First occurrence wins, which keeps
|
||||
the operator's declared precedence intact.
|
||||
"""
|
||||
configured: Final[object] = litellm.bedrock_request_metadata_fields
|
||||
if not isinstance(configured, (list, tuple)):
|
||||
return ()
|
||||
fields: Final = tuple(str(field) for field in configured)
|
||||
return tuple(field for index, field in enumerate(fields) if field not in fields[:index])
|
||||
|
||||
|
||||
def _metadata_sources(litellm_params: Mapping[str, object] | None) -> tuple[Mapping[str, object], ...]:
|
||||
"""``metadata`` on /v1/chat/completions, ``litellm_metadata`` on the LITELLM_METADATA_ROUTES."""
|
||||
if litellm_params is None:
|
||||
return ()
|
||||
return tuple(
|
||||
source
|
||||
for name in _METADATA_PARAM_NAMES
|
||||
for source in (litellm_params.get(name),)
|
||||
if isinstance(source, Mapping)
|
||||
)
|
||||
|
||||
|
||||
def _identity_pairs(
|
||||
sources: tuple[Mapping[str, object], ...],
|
||||
allowed_fields: tuple[str, ...],
|
||||
) -> tuple[tuple[str, str], ...]:
|
||||
return tuple(
|
||||
(field, value)
|
||||
for field in allowed_fields
|
||||
if field.startswith(BEDROCK_REQUEST_METADATA_IDENTITY_PREFIX)
|
||||
for value in (_first_text(sources, field),)
|
||||
if value is not None and _is_forwardable(field, value)
|
||||
)[:BEDROCK_REQUEST_METADATA_MAX_PAIRS]
|
||||
|
||||
|
||||
def _first_text(sources: tuple[Mapping[str, object], ...], field: str) -> str | None:
|
||||
return next((value for source in sources if isinstance(value := source.get(field), str)), None)
|
||||
|
||||
|
||||
def _client_pairs(
|
||||
sources: tuple[Mapping[str, object], ...],
|
||||
allowed_fields: tuple[str, ...],
|
||||
caller_metadata: object,
|
||||
budget: int,
|
||||
) -> tuple[tuple[str, str], ...]:
|
||||
spend_logs_pairs: Final = (
|
||||
tuple(pair for source in sources for pair in _text_pairs(source.get(BEDROCK_REQUEST_METADATA_CLIENT_FIELD)))
|
||||
if BEDROCK_REQUEST_METADATA_CLIENT_FIELD in allowed_fields
|
||||
else ()
|
||||
)
|
||||
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)
|
||||
)
|
||||
return tuple(
|
||||
pair
|
||||
for index, pair in enumerate(candidates)
|
||||
if pair[0] not in tuple(earlier for earlier, _ in candidates[:index])
|
||||
)[:budget]
|
||||
|
||||
|
||||
def resolve_bedrock_request_metadata(
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
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]``.
|
||||
|
||||
``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.
|
||||
"""
|
||||
allowed_fields: Final = _allowed_fields()
|
||||
if not allowed_fields:
|
||||
return None
|
||||
sources: Final = _metadata_sources(litellm_params)
|
||||
identity: Final = _identity_pairs(sources, allowed_fields)
|
||||
client: Final = _client_pairs(
|
||||
sources=sources,
|
||||
allowed_fields=allowed_fields,
|
||||
caller_metadata=caller_metadata,
|
||||
budget=BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(identity),
|
||||
)
|
||||
resolved: Final = {key: value for key, value in (*identity, *client)}
|
||||
return resolved or None
|
||||
|
||||
|
||||
def bedrock_request_metadata_is_owned() -> bool:
|
||||
"""
|
||||
Whether the proxy OWNS the request-metadata field and header name for this request.
|
||||
|
||||
Ownership follows the operator's opt-in alone, never whether anything resolved, because a
|
||||
caller can suppress the resolver by omitting the allow-listed fields or by sending values
|
||||
that all fail Bedrock's rules. Owned-but-empty has to mean "absent on the wire" rather than
|
||||
"fall back to whatever the caller supplied", or the reserved-prefix guarantee is bypassable
|
||||
by anyone who can make the resolver produce nothing.
|
||||
"""
|
||||
return bool(_allowed_fields())
|
||||
|
||||
|
||||
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.
|
||||
|
||||
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.
|
||||
"""
|
||||
if not bedrock_request_metadata_is_owned():
|
||||
return frozenset(), ()
|
||||
resolved: Final = resolve_bedrock_request_metadata(litellm_params)
|
||||
if resolved is None:
|
||||
return _OWNED_HEADER_NAMES, ()
|
||||
return _OWNED_HEADER_NAMES, ((BEDROCK_REQUEST_METADATA_HEADER, json.dumps(resolved, separators=(",", ":"))),)
|
||||
|
||||
|
||||
def merge_bedrock_invoke_headers(
|
||||
headers: dict[str, str],
|
||||
caller_owned: tuple[tuple[str, str], ...],
|
||||
proxy_owned: tuple[tuple[str, str], ...],
|
||||
proxy_owned_names: frozenset[str],
|
||||
) -> dict[str, str]:
|
||||
"""
|
||||
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.
|
||||
"""
|
||||
if not caller_owned and not proxy_owned and not proxy_owned_names:
|
||||
return headers
|
||||
existing_names: Final = frozenset(name.lower() for name in headers)
|
||||
return {
|
||||
name: value
|
||||
for name, value in (
|
||||
*((n, v) for n, v in headers.items() if n.lower() not in proxy_owned_names),
|
||||
*((n, v) for n, v in caller_owned if n.lower() not in existing_names),
|
||||
*proxy_owned,
|
||||
)
|
||||
}
|
||||
422
tests/test_litellm/llms/bedrock/test_request_metadata.py
Normal file
422
tests/test_litellm/llms/bedrock/test_request_metadata.py
Normal file
|
|
@ -0,0 +1,422 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.path.insert(0, os.path.abspath("../../../.."))
|
||||
|
||||
import litellm
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.chat.converse_transformation import AmazonConverseConfig
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.amazon_openai_transformation import (
|
||||
AmazonBedrockOpenAIConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
AmazonInvokeConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.messages.invoke_transformations.anthropic_claude3_transformation import (
|
||||
AmazonAnthropicClaudeMessagesConfig,
|
||||
)
|
||||
from litellm.llms.bedrock.request_metadata import (
|
||||
BEDROCK_REQUEST_METADATA_HEADER,
|
||||
BEDROCK_REQUEST_METADATA_MAX_PAIRS,
|
||||
resolve_bedrock_request_metadata,
|
||||
)
|
||||
|
||||
MODEL = "anthropic.claude-3-5-sonnet-20240620-v1:0"
|
||||
MESSAGES = [{"role": "user", "content": "hi"}]
|
||||
ALL_FIELDS = [
|
||||
"user_api_key_alias",
|
||||
"user_api_key_team_alias",
|
||||
"user_api_key_user_email",
|
||||
"spend_logs_metadata",
|
||||
]
|
||||
IDENTITY = {"user_api_key_alias": "prod-key", "user_api_key_team_alias": "platform"}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_setting():
|
||||
previous = litellm.bedrock_request_metadata_fields
|
||||
yield
|
||||
litellm.bedrock_request_metadata_fields = previous
|
||||
|
||||
|
||||
def litellm_params(metadata_key, **metadata):
|
||||
return {metadata_key: dict(metadata)}
|
||||
|
||||
|
||||
def converse_body(litellm_params_value, optional_params=None):
|
||||
return AmazonConverseConfig()._transform_request(
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params=dict(optional_params or {}),
|
||||
litellm_params=dict(litellm_params_value),
|
||||
)
|
||||
|
||||
|
||||
def converse_body_async(litellm_params_value, optional_params=None):
|
||||
"""The proxy serves completions through the async transform, so every rule asserted against
|
||||
the sync body has to be asserted against this one too or half the product is untested."""
|
||||
return asyncio.run(
|
||||
AmazonConverseConfig()._async_transform_request(
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params=dict(optional_params or {}),
|
||||
litellm_params=dict(litellm_params_value),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
CONVERSE_DRIVERS = [converse_body, converse_body_async]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("setting", [None, []])
|
||||
def test_feature_off_by_default_leaves_body_and_headers_untouched(setting):
|
||||
litellm.bedrock_request_metadata_fields = setting
|
||||
params = litellm_params("metadata", spend_logs_metadata={"team": "x"}, **IDENTITY)
|
||||
|
||||
assert "requestMetadata" not in converse_body(params)
|
||||
assert BEDROCK_REQUEST_METADATA_HEADER not in AmazonInvokeConfig().validate_environment(
|
||||
headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=dict(params)
|
||||
)
|
||||
messages_headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
|
||||
headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=dict(params)
|
||||
)
|
||||
assert BEDROCK_REQUEST_METADATA_HEADER not in messages_headers
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
def test_resolver_reads_both_metadata_variable_names(metadata_key):
|
||||
"""`/v1/chat/completions` populates `metadata`; the LITELLM_METADATA_ROUTES populate
|
||||
`litellm_metadata`. Reading only one silently forwards nothing on the other route."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
params = litellm_params(metadata_key, spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY)
|
||||
|
||||
assert converse_body(params)["requestMetadata"] == {**IDENTITY, "cost_center": "cc-1"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"])
|
||||
def test_invoke_messages_header_reads_both_metadata_variable_names(metadata_key):
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
params = litellm_params(metadata_key, **IDENTITY)
|
||||
|
||||
headers, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
|
||||
headers={}, model=MODEL, messages=MESSAGES, optional_params={}, litellm_params=params
|
||||
)
|
||||
|
||||
assert json.loads(headers[BEDROCK_REQUEST_METADATA_HEADER]) == IDENTITY
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reverse_client_keys", [False, True])
|
||||
@pytest.mark.parametrize("field_order", [ALL_FIELDS, list(reversed(ALL_FIELDS))])
|
||||
@pytest.mark.parametrize("client_source", ["spend_logs_metadata", "requestMetadata"])
|
||||
def test_identity_survives_a_caller_filling_every_slot(reverse_client_keys, field_order, client_source):
|
||||
"""A caller sending 16 keys of its own must not evict the identity the feature exists to
|
||||
produce. Driven over every input ordering so the invariant is not an accident of one."""
|
||||
litellm.bedrock_request_metadata_fields = field_order
|
||||
client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS)]
|
||||
client_pairs = {key: "v" for key in (reversed(client_keys) if reverse_client_keys else client_keys)}
|
||||
if client_source == "spend_logs_metadata":
|
||||
params, optional_params = litellm_params("metadata", spend_logs_metadata=client_pairs, **IDENTITY), {}
|
||||
else:
|
||||
params, optional_params = litellm_params("metadata", **IDENTITY), {"requestMetadata": client_pairs}
|
||||
|
||||
resolved = converse_body(params, optional_params)["requestMetadata"]
|
||||
|
||||
assert len(resolved) == BEDROCK_REQUEST_METADATA_MAX_PAIRS
|
||||
for key, value in IDENTITY.items():
|
||||
assert resolved[key] == value
|
||||
assert len([key for key in resolved if key.startswith("client_")]) == (
|
||||
BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field_order",
|
||||
[
|
||||
["user_api_key_alias", "user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata"],
|
||||
["user_api_key_alias", "user_api_key_team_alias", "user_api_key_alias", "spend_logs_metadata"],
|
||||
["user_api_key_alias", "user_api_key_team_alias", "spend_logs_metadata", "user_api_key_team_alias"],
|
||||
],
|
||||
)
|
||||
def test_a_field_repeated_in_the_allow_list_does_not_consume_a_client_slot(field_order):
|
||||
"""An operator repeating a field in YAML must not inflate the reserved count and shrink the
|
||||
client budget. Asserts the client keys that should have fitted actually reach the wire, since
|
||||
asserting only that identity survives passes with or without the deduplication."""
|
||||
litellm.bedrock_request_metadata_fields = field_order
|
||||
client_keys = [f"client_{index:02d}" for index in range(BEDROCK_REQUEST_METADATA_MAX_PAIRS - 1)]
|
||||
params = litellm_params("metadata", spend_logs_metadata={key: "v" for key in client_keys}, **IDENTITY)
|
||||
|
||||
resolved = converse_body(params)["requestMetadata"]
|
||||
|
||||
expected_client_slots = BEDROCK_REQUEST_METADATA_MAX_PAIRS - len(IDENTITY)
|
||||
assert resolved == {**IDENTITY, **{key: "v" for key in client_keys[:expected_client_slots]}}
|
||||
assert len(resolved) == BEDROCK_REQUEST_METADATA_MAX_PAIRS
|
||||
assert client_keys[expected_client_slots - 1] in resolved
|
||||
|
||||
|
||||
@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"],
|
||||
)
|
||||
def test_caller_cannot_forge_or_shadow_a_reserved_identity_key(forged_key, client_source):
|
||||
"""`user_api_key_org_alias` and `user_api_key_hash` are names the proxy does not set here,
|
||||
so an exact-key reservation would let the forged value through under a name that reads as
|
||||
proxy-authoritative in the AWS billing record."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
forged = {forged_key: "attacker-controlled"}
|
||||
if client_source == "spend_logs_metadata":
|
||||
params, optional_params = litellm_params("metadata", spend_logs_metadata=forged, **IDENTITY), {}
|
||||
else:
|
||||
params, optional_params = litellm_params("metadata", **IDENTITY), {"requestMetadata": forged}
|
||||
|
||||
resolved = converse_body(params, optional_params)["requestMetadata"]
|
||||
|
||||
assert resolved == IDENTITY
|
||||
assert "attacker-controlled" not in resolved.values()
|
||||
|
||||
|
||||
def test_identity_violating_the_character_class_is_dropped_and_the_request_succeeds():
|
||||
"""A team alias with an apostrophe must not turn a working request into a 400 the moment
|
||||
an operator flips the setting on."""
|
||||
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_user_email="x" * 300,
|
||||
)
|
||||
|
||||
body = converse_body(params)
|
||||
|
||||
assert body["requestMetadata"] == {"user_api_key_alias": "prod-key"}
|
||||
assert body["messages"]
|
||||
|
||||
|
||||
def test_caller_supplied_violation_still_raises_bad_request():
|
||||
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"}},
|
||||
)
|
||||
|
||||
|
||||
def test_non_string_and_absent_identity_values_are_dropped():
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS + ["user_api_key_spend"]
|
||||
params = litellm_params("metadata", user_api_key_alias="prod-key", user_api_key_spend=1.25)
|
||||
|
||||
assert converse_body(params)["requestMetadata"] == {"user_api_key_alias": "prod-key"}
|
||||
|
||||
|
||||
def test_email_is_separately_opt_in():
|
||||
"""PII crossing into CloudTrail only when the operator names the field."""
|
||||
identity_with_email = {**IDENTITY, "user_api_key_user_email": "owner@example.com"}
|
||||
litellm.bedrock_request_metadata_fields = ["user_api_key_alias", "user_api_key_team_alias"]
|
||||
assert (
|
||||
"user_api_key_user_email"
|
||||
not in converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"]
|
||||
)
|
||||
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
assert converse_body(litellm_params("metadata", **identity_with_email))["requestMetadata"] == identity_with_email
|
||||
|
||||
|
||||
def test_resolver_returns_none_when_nothing_survives():
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
assert resolve_bedrock_request_metadata(litellm_params=None) is None
|
||||
assert resolve_bedrock_request_metadata(litellm_params={"metadata": {"unrelated": "x"}}) is None
|
||||
|
||||
|
||||
def test_invoke_header_is_json_encoded_and_signed():
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
params = litellm_params("metadata", spend_logs_metadata={"cost_center": "cc-1"}, **IDENTITY)
|
||||
|
||||
headers = AmazonInvokeConfig().validate_environment(
|
||||
headers={"anthropic-version": "bedrock-2023-05-31"},
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params={},
|
||||
litellm_params=params,
|
||||
)
|
||||
|
||||
assert json.loads(headers[BEDROCK_REQUEST_METADATA_HEADER]) == {**IDENTITY, "cost_center": "cc-1"}
|
||||
signed = BaseAWSLLM()._filter_headers_for_aws_signature(headers)
|
||||
assert BEDROCK_REQUEST_METADATA_HEADER in signed
|
||||
assert "anthropic-version" not in signed
|
||||
|
||||
|
||||
def test_a_caller_supplied_guardrail_header_still_wins():
|
||||
"""The no-displace rule is deliberate for the guardrail headers and must survive the
|
||||
request-metadata header becoming proxy-owned."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
|
||||
headers = AmazonInvokeConfig().validate_environment(
|
||||
headers={"X-Amzn-Bedrock-GuardrailIdentifier": "caller-set"},
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params={"guardrailConfig": {"guardrailIdentifier": "gid", "guardrailVersion": "DRAFT"}},
|
||||
litellm_params=litellm_params("metadata", **IDENTITY),
|
||||
)
|
||||
|
||||
assert headers["X-Amzn-Bedrock-GuardrailIdentifier"] == "caller-set"
|
||||
assert headers["X-Amzn-Bedrock-GuardrailVersion"] == "DRAFT"
|
||||
|
||||
|
||||
FORGED = '{"user_api_key_alias":"FORGED-KEY","user_api_key_team_alias":"FORGED-TEAM"}'
|
||||
|
||||
|
||||
def invoke_headers(caller_headers, params, optional_params=None):
|
||||
return AmazonInvokeConfig().validate_environment(
|
||||
headers=dict(caller_headers),
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params=dict(optional_params or {}),
|
||||
litellm_params=dict(params),
|
||||
)
|
||||
|
||||
|
||||
def messages_headers(caller_headers, params):
|
||||
resolved, _ = AmazonAnthropicClaudeMessagesConfig().validate_anthropic_messages_environment(
|
||||
headers=dict(caller_headers),
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params={},
|
||||
litellm_params=dict(params),
|
||||
)
|
||||
return resolved
|
||||
|
||||
|
||||
def openai_invoke_headers(caller_headers, params):
|
||||
return AmazonBedrockOpenAIConfig().validate_environment(
|
||||
headers=dict(caller_headers),
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params={},
|
||||
litellm_params=dict(params),
|
||||
)
|
||||
|
||||
|
||||
def converse_headers(caller_headers, params):
|
||||
return AmazonConverseConfig().validate_environment(
|
||||
headers=dict(caller_headers),
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params={},
|
||||
litellm_params=dict(params),
|
||||
)
|
||||
|
||||
|
||||
HEADER_DRIVERS = [invoke_headers, messages_headers, openai_invoke_headers, converse_headers]
|
||||
|
||||
|
||||
def metadata_header_values(headers):
|
||||
return [value for name, value in headers.items() if name.lower() == BEDROCK_REQUEST_METADATA_HEADER.lower()]
|
||||
|
||||
|
||||
def test_converse_still_sets_the_bearer_authorization_header():
|
||||
"""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."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
|
||||
headers = AmazonConverseConfig().validate_environment(
|
||||
headers={},
|
||||
model=MODEL,
|
||||
messages=MESSAGES,
|
||||
optional_params={},
|
||||
litellm_params=dict(litellm_params("metadata", **IDENTITY)),
|
||||
api_key="sk-converse-bearer",
|
||||
)
|
||||
|
||||
assert headers["Authorization"] == "Bearer sk-converse-bearer"
|
||||
assert metadata_header_values(headers) == [json.dumps(IDENTITY, separators=(",", ":"))]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
|
||||
@pytest.mark.parametrize(
|
||||
"caller_header_name",
|
||||
[BEDROCK_REQUEST_METADATA_HEADER, BEDROCK_REQUEST_METADATA_HEADER.lower(), "x-AMZN-bedrock-Request-METADATA"],
|
||||
)
|
||||
def test_a_caller_cannot_forge_the_request_metadata_header(driver, caller_header_name):
|
||||
"""`extra_headers` puts caller-supplied names into the same dict the proxy merges into, so a
|
||||
deferring merge would sign the caller's forged identity into the AWS billing record. Every
|
||||
spelling must lose, or a second variant is left for the transport to choose between."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
|
||||
headers = driver({caller_header_name: FORGED}, litellm_params("metadata", **IDENTITY))
|
||||
|
||||
values = metadata_header_values(headers)
|
||||
assert values == [json.dumps(IDENTITY, separators=(",", ":"))]
|
||||
assert "FORGED" not in json.dumps(headers)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
|
||||
def test_a_caller_cannot_forge_the_header_when_the_resolver_yields_nothing(driver):
|
||||
"""Forwarding enabled but nothing resolvable, which a caller can arrange by supplying values
|
||||
that all fail Bedrock's rules. Owned-but-empty must mean no header on the wire, never a
|
||||
fallback to the caller's."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
unresolvable = litellm_params("metadata", user_api_key_alias="O'Brien's key", user_api_key_team_alias="x" * 300)
|
||||
|
||||
headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, unresolvable)
|
||||
|
||||
assert metadata_header_values(headers) == []
|
||||
assert "FORGED" not in json.dumps(headers)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
|
||||
@pytest.mark.parametrize(
|
||||
"forged_key",
|
||||
["user_api_key_team_alias", "user_api_key_org_alias", "user_api_key_hash"],
|
||||
)
|
||||
def test_a_caller_cannot_keep_reserved_body_keys_when_the_resolver_yields_nothing(forged_key, driver):
|
||||
"""The Converse body has the same fail-open shape as the header: with forwarding on and
|
||||
nothing resolvable, leaving the caller's `requestMetadata` in place would keep their
|
||||
reserved-prefix keys on the wire. Owned-but-empty must remove the field outright."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
|
||||
body = driver(litellm_params("metadata"), {"requestMetadata": {forged_key: "FORGED"}})
|
||||
|
||||
assert "requestMetadata" not in body
|
||||
assert "FORGED" not in json.dumps(body)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
|
||||
def test_benign_caller_body_metadata_still_survives_when_no_identity_resolves(driver):
|
||||
"""Removing the field must be scoped to the reserved keys being the only thing left, not a
|
||||
blanket drop of the caller's own attribution pairs."""
|
||||
litellm.bedrock_request_metadata_fields = ALL_FIELDS
|
||||
|
||||
body = driver(
|
||||
litellm_params("metadata"),
|
||||
{"requestMetadata": {"cost_center": "cc-9", "user_api_key_team_alias": "FORGED"}},
|
||||
)
|
||||
|
||||
assert body["requestMetadata"] == {"cost_center": "cc-9"}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", CONVERSE_DRIVERS)
|
||||
def test_caller_body_metadata_is_left_alone_when_forwarding_is_off(driver):
|
||||
"""With the feature off the proxy does not own the field, so the pre-existing pass-through
|
||||
behaviour for a caller-supplied `requestMetadata` must be unchanged."""
|
||||
litellm.bedrock_request_metadata_fields = None
|
||||
caller_supplied = {"user_api_key_team_alias": "caller-set", "cost_center": "cc-9"}
|
||||
|
||||
body = driver(litellm_params("metadata", **IDENTITY), {"requestMetadata": caller_supplied})
|
||||
|
||||
assert body["requestMetadata"] == caller_supplied
|
||||
|
||||
|
||||
@pytest.mark.parametrize("driver", HEADER_DRIVERS)
|
||||
def test_a_caller_header_is_left_alone_when_forwarding_is_off(driver):
|
||||
"""The proxy only claims the name when the operator turned forwarding on; with the feature
|
||||
off this is an ordinary passthrough header and stripping it would be a regression."""
|
||||
litellm.bedrock_request_metadata_fields = None
|
||||
|
||||
headers = driver({BEDROCK_REQUEST_METADATA_HEADER: FORGED}, litellm_params("metadata", **IDENTITY))
|
||||
|
||||
assert metadata_header_values(headers) == [FORGED]
|
||||
Loading…
Add table
Reference in a new issue