From 845f10952405dce252e0ce98d5fc96ebd825009d Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Tue, 10 Nov 2026 13:28:49 -0800 Subject: [PATCH] fix(bedrock): forward x-pass- headers on the /bedrock/model pass-through routes (#50728) * fix(bedrock): forward x-pass- headers on the /bedrock/model pass-through routes The /bedrock/model/{model}/invoke, invoke-with-response-stream, converse and converse-stream routes built the pass-through call from the request body, so x-pass- headers never reached Bedrock. The processor now hands the caller's x-pass- headers to the pass-through route, keeps them out of the logged body snapshot, and on Converse drops the invoke-only guardrail headers when the deployment sets guardrailConfig and the request metadata header when the proxy owns request metadata * fix(bedrock): bill the served service tier on invoke pass-through and scope x-pass- forwarding to the Bedrock routes * test(bedrock): assert forwarded x-pass- headers on the request Bedrock receives * test(bedrock): gate the router-model x-pass- forwarding on the request Bedrock receives * fix(bedrock): bill a caller's forwarded invoke guardrail and keep x-pass- headers out of guardrail vendor payloads * fix(bedrock): bill a caller guardrail on every invoke model and replace a trace value Bedrock rejects * fix(bedrock): keep guardrail traces out of chat responses, forward the trace header as sent, and drop credential x-pass- headers --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../base_llm/passthrough/transformation.py | 9 +- .../bedrock/passthrough/transformation.py | 95 +++- litellm/passthrough/main.py | 16 +- litellm/passthrough/utils.py | 34 +- litellm/proxy/litellm_pre_call_utils.py | 4 + .../llm_passthrough_endpoints.py | 9 + ...drock_passthrough_guardrail_config_wire.py | 26 +- ...drock_passthrough_request_metadata_wire.py | 61 ++- .../test_base_invoke_transformation.py | 136 ++++++ ...test_bedrock_passthrough_transformation.py | 444 +++++++++++++++++- tests/unit/passthrough/test_utils.py | 78 +++ .../test_llm_pass_through_endpoints.py | 180 +++++++ .../unit/proxy/test_litellm_pre_call_utils.py | 32 ++ 13 files changed, 1081 insertions(+), 43 deletions(-) create mode 100644 tests/unit/passthrough/test_utils.py diff --git a/litellm/llms/base_llm/passthrough/transformation.py b/litellm/llms/base_llm/passthrough/transformation.py index 7e58cc6e71b..898033bff8f 100644 --- a/litellm/llms/base_llm/passthrough/transformation.py +++ b/litellm/llms/base_llm/passthrough/transformation.py @@ -4,6 +4,7 @@ import re from abc import abstractmethod from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Final, Protocol, TypeAlias from pydantic import ConfigDict, TypeAdapter, ValidationError @@ -29,6 +30,7 @@ if TYPE_CHECKING: RELAYED_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) _DECODED_JSON: Final = TypeAdapter(object) +_NO_RESPONSE_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) class PassthroughMetadata(LiteLLMBaseModel): @@ -236,7 +238,12 @@ class BasePassthroughConfig(BaseLLMModelInfo): return None def create_stream_collector( - self, model: str, custom_llm_provider: str, endpoint: str + self, + model: str, + custom_llm_provider: str, + endpoint: str, + *, + response_headers: Mapping[str, str] = _NO_RESPONSE_HEADERS, ) -> PassthroughStreamCollector: return RawBytesStreamCollector( provider_config=self, model=model, custom_llm_provider=custom_llm_provider, endpoint=endpoint diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 4a2bb57ce5a..901c16afaa7 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -3,16 +3,21 @@ from collections.abc import Callable, Iterator, Mapping, Sequence from functools import partial from itertools import chain from types import MappingProxyType -from typing import TYPE_CHECKING, ClassVar, Final, Optional, cast +from typing import TYPE_CHECKING, ClassVar, Final, Optional, TypeVar, cast import httpx from httpx import Response -from pydantic import TypeAdapter +from pydantic import BaseModel, 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.litellm_core_utils.llm_cost_calc.guardrail_cost import object_mapping +from litellm.llms.base_llm.passthrough.transformation import ( + BasePassthroughConfig, + PassthroughStreamCollector, + relayed_json_object, +) from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( bedrock_invoke_guardrail_headers, @@ -21,7 +26,12 @@ from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletio from litellm.types.utils import ModelResponseStream from ..base_aws_llm import BaseAWSLLM -from ..common_utils import BedrockError, BedrockEventStreamDecoderBase, BedrockModelInfo +from ..common_utils import ( + BEDROCK_INVOKE_TRACE_FIELD, + BedrockError, + BedrockEventStreamDecoderBase, + BedrockModelInfo, +) from ..request_metadata import ( BEDROCK_REQUEST_METADATA_HEADER, BedrockRequestMetadataViolation, @@ -49,6 +59,9 @@ _BEDROCK_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stre _BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset( {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} ) +_BEDROCK_SERVICE_TIER_HEADER: Final = "x-amzn-bedrock-service-tier" +_NO_RESPONSE_HEADERS: Final[Mapping[str, str]] = MappingProxyType({}) +_TieredResponse: Final = TypeVar("_TieredResponse", bound=BaseModel) def _bedrock_endpoint_action(endpoint: str) -> str: @@ -111,7 +124,9 @@ def _with_request_metadata( 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) + return _without_owned_request_metadata_header(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 @@ -137,6 +152,10 @@ def _converse_guardrail_config(action: str, guardrail_config: object) -> object: return {key: value for key, value in fields.items() if key != "streamProcessingMode"} +def _without_invoke_guardrail_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return {name: value for name, value in headers.items() if name.lower() not in _BEDROCK_INVOKE_GUARDRAIL_HEADERS} + + def _with_deployment_guardrail( action: str, headers: Mapping[str, str], @@ -146,7 +165,10 @@ def _with_deployment_guardrail( if guardrail_config is None: return headers, body if action in _BEDROCK_CONVERSE_ACTIONS: - return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)} + return _without_invoke_guardrail_headers(headers), { + **body, + "guardrailConfig": _converse_guardrail_config(action, guardrail_config), + } if action in _BEDROCK_INVOKE_ACTIONS: guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items()) return merge_bedrock_invoke_headers( @@ -266,6 +288,19 @@ class _CoalescedChunks: return self._chunks +_BEDROCK_EVENT: Final = TypeAdapter(dict[str, object]) + + +def _with_invoke_trace(chunk: ModelResponseStream, trace: Mapping[str, object] | None) -> ModelResponseStream: + if trace is None: + return chunk + return chunk.model_copy( + update={ + "provider_specific_fields": {**(chunk.provider_specific_fields or {}), BEDROCK_INVOKE_TRACE_FIELD: trace} + } + ) + + def _translate_message(decoder: "AWSEventStreamDecoder", message: str) -> ModelResponseStream | None: from litellm.litellm_core_utils.streaming_handler import ( convert_generic_chunk_to_model_response_stream, @@ -273,11 +308,15 @@ def _translate_message(decoder: "AWSEventStreamDecoder", message: str) -> ModelR ) from litellm.types.utils import GenericStreamingChunk - translated_chunk: Final = decoder.chunk_parser(chunk_data=json.loads(message)) + event: Final = _BEDROCK_EVENT.validate_json(message) + trace: Final = object_mapping(event.get(BEDROCK_INVOKE_TRACE_FIELD)) + translated_chunk: Final = decoder.chunk_parser(chunk_data=event) if isinstance(translated_chunk, ModelResponseStream): - return translated_chunk + return _with_invoke_trace(translated_chunk, trace) if generic_chunk_has_all_required_fields(cast(dict, translated_chunk)): - return convert_generic_chunk_to_model_response_stream(cast(GenericStreamingChunk, translated_chunk)) + return _with_invoke_trace( + convert_generic_chunk_to_model_response_stream(cast(GenericStreamingChunk, translated_chunk)), trace + ) return None @@ -332,15 +371,33 @@ def _count_prompt_tokens(model: str, request_body: object) -> int: return token_counter(model=model, messages=messages) +def _invoke_guardrail_trace(httpx_response: Response) -> Mapping[str, object]: + reply: Final = relayed_json_object(httpx_response) or {} + return {BEDROCK_INVOKE_TRACE_FIELD: reply.get(BEDROCK_INVOKE_TRACE_FIELD)} + + +def _served_service_tier(response_headers: Mapping[str, str]) -> str | None: + headers: Final = httpx.Headers(response_headers) + if _BEDROCK_SERVICE_TIER_HEADER not in headers: + return None + return headers[_BEDROCK_SERVICE_TIER_HEADER] + + +def _with_served_service_tier(response: _TieredResponse, service_tier: str | None) -> _TieredResponse: + if service_tier is None: + return response + return response.model_copy(update={"service_tier": service_tier}) + + def _build_logged_response( - chunks: Sequence[ModelResponseStream], litellm_logging_obj: "LiteLLMLoggingObj" + chunks: Sequence[ModelResponseStream], litellm_logging_obj: "LiteLLMLoggingObj", service_tier: str | None ) -> Optional["CostResponseTypes"]: from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder if len(chunks) == 0: return None return stream_chunk_builder( - chunks=list(chunks), + chunks=[*chunks[:-1], _with_served_service_tier(chunks[-1], service_tier)], logging_obj=litellm_logging_obj, count_prompt_tokens=partial( _count_prompt_tokens, litellm_logging_obj.model, litellm_logging_obj.model_call_details.get("input") @@ -355,11 +412,13 @@ class BedrockEventStreamCollector: self, parse_event: Callable[["EventStreamMessage"], str | None], decoder: Optional["AWSEventStreamDecoder"], + service_tier: str | None = None, ) -> None: from botocore.eventstream import EventStreamBuffer self._parse_event = parse_event self._decoder = decoder + self._service_tier: Final = service_tier self._event_stream_buffer: Final[EventStreamBuffer] = EventStreamBuffer() self._chunks: Final = _CoalescedChunks() @@ -377,7 +436,7 @@ class BedrockEventStreamCollector: self._chunks.add(translated) def build_logged_response(self, litellm_logging_obj: "LiteLLMLoggingObj") -> Optional["CostResponseTypes"]: - return _build_logged_response(self._chunks.chunks(), litellm_logging_obj) + return _build_logged_response(self._chunks.chunks(), litellm_logging_obj, self._service_tier) class BedrockAgentRuntimeAWS(BaseAWSLLM): @@ -542,15 +601,23 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD request_data=request_data, encoding=encoding, ) + if "invoke" in endpoint: + logging_obj.record_native_bedrock_guardrail(_invoke_guardrail_trace(httpx_response)) - return litellm_model_response + return _with_served_service_tier(litellm_model_response, _served_service_tier(httpx_response.headers)) def create_stream_collector( - self, model: str, custom_llm_provider: str, endpoint: str + self, + model: str, + custom_llm_provider: str, + endpoint: str, + *, + response_headers: Mapping[str, str] = _NO_RESPONSE_HEADERS, ) -> PassthroughStreamCollector: return BedrockEventStreamCollector( parse_event=self._parse_message_from_event, decoder=self._get_event_stream_decoder(model=model, endpoint=endpoint), + service_tier=_served_service_tier(response_headers), ) def _get_event_stream_decoder(self, model: str, endpoint: str) -> Optional["AWSEventStreamDecoder"]: diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 4e94b1bb910..77b2cf9b5b5 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -57,11 +57,17 @@ def _as_generator(iterable: Iterator[bytes]) -> Generator[bytes, bytes, None]: class _SpendCollection: """Feeds relayed chunks to the provider's stream collector without letting spend tracking break the relay.""" - def __init__(self, provider_config: BasePassthroughConfig, litellm_logging_obj: LiteLLMLoggingObj) -> None: + def __init__( + self, + provider_config: BasePassthroughConfig, + litellm_logging_obj: LiteLLMLoggingObj, + response_headers: httpx.Headers, + ) -> None: self.collector: Final[PassthroughStreamCollector] = provider_config.create_stream_collector( model=litellm_logging_obj.model, custom_llm_provider=litellm_logging_obj.model_call_details.get("custom_llm_provider", ""), endpoint=litellm_logging_obj.model_call_details.get("endpoint", ""), + response_headers=response_headers, ) self.chunk_count = 0 self._failed = False @@ -97,7 +103,8 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): self._response: httpx.Response self._iterator: AsyncGenerator[bytes, bytes] self._litellm_logging_obj = litellm_logging_obj - self._spend = _SpendCollection(provider_config, litellm_logging_obj) + self._provider_config: Final = provider_config + self._spend: _SpendCollection self._flush_scheduled = False self._background_tasks: set[asyncio.Task[bool]] = set() # mutable-ok: instance set for background task tracking self._flush: asyncio.Task[bool] | None = None @@ -129,6 +136,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): self._response = await self._response_coro self.headers = self._response.headers self.status_code = self._response.status_code + self._spend = _SpendCollection(self._provider_config, self._litellm_logging_obj, self._response.headers) self._initialized = True try: raise_for_provider_status(self._response) @@ -148,7 +156,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): return _init().__await__() def _start_flush(self) -> None: - if self._flush_scheduled or not self._spend.should_flush: + if self._flush_scheduled or not self._initialized or not self._spend.should_flush: return self._flush_scheduled = True @@ -237,7 +245,7 @@ class PassthroughStreamingResponse(Generator[bytes, bytes, None]): self.status_code = response.status_code self._litellm_logging_obj = litellm_logging_obj self._iterator: Generator[bytes, bytes, None] = _as_generator(response.iter_bytes()) - self._spend = _SpendCollection(provider_config, litellm_logging_obj) + self._spend = _SpendCollection(provider_config, litellm_logging_obj, response.headers) self._flush_scheduled = False def _start_flush(self) -> None: diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index 452c9c7de9d..1bf5f08b1c1 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -25,7 +25,16 @@ _PASS_THROUGH_PROTECTED_HEADERS: Final[frozenset] = frozenset( ) # Header name prefix used to block AWS SigV4 signing headers from being overridden. -_PASS_THROUGH_PROTECTED_HEADER_PREFIXES: Final[tuple] = ("x-amz-",) +_PASS_THROUGH_PROTECTED_HEADER_PREFIXES: Final[tuple[str, ...]] = ("x-amz-",) + + +def _is_forwarded_pass_through_header(header_name: str) -> bool: + if not header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): + return False + forwarded_name: Final = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() + return forwarded_name not in _PASS_THROUGH_PROTECTED_HEADERS and not any( + forwarded_name.startswith(prefix) for prefix in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES + ) class BasePassthroughUtils: @@ -42,18 +51,11 @@ class BasePassthroughUtils: # parse_qs returns a dict where each value is a list, so let's flatten it updated_existing_query_params: Final = {k: v[0] if len(v) == 1 else v for k, v in existing_query_params.items()} - # Start with default query params (lowest priority) - merged_params: Final = {} - if default_query_params: - merged_params.update(default_query_params) + return {**(default_query_params or {}), **updated_existing_query_params, **request_query_params} - # Override with existing URL query params (medium priority) - merged_params.update(updated_existing_query_params) - - # Override with request query params (highest priority - client can override anything) - merged_params.update(request_query_params) - - return merged_params + @staticmethod + def pass_through_prefixed_headers(request_headers: Mapping[str, str]) -> Mapping[str, str]: + return {name: value for name, value in request_headers.items() if _is_forwarded_pass_through_header(name)} @staticmethod def forward_headers_from_request( @@ -88,16 +90,14 @@ class BasePassthroughUtils: # Credential and protocol-level headers are excluded from this mechanism. for header_name, header_value in request_headers.items(): if header_name.lower().startswith(PASS_THROUGH_HEADER_PREFIX): - # Strip the 'x-pass-' prefix and normalize to lowercase - actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() - if actual_header_name in _PASS_THROUGH_PROTECTED_HEADERS or any( - actual_header_name.startswith(p) for p in _PASS_THROUGH_PROTECTED_HEADER_PREFIXES - ): + if not _is_forwarded_pass_through_header(header_name): verbose_logger.debug( "x-pass- header %s maps to a protected header name; skipping", header_name, ) continue + # Strip the 'x-pass-' prefix and normalize to lowercase + actual_header_name = header_name[len(PASS_THROUGH_HEADER_PREFIX) :].lower() headers[actual_header_name] = header_value return headers diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 22399ced718..d5eac099fb6 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1656,6 +1656,10 @@ def redact_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: } +def without_credential_headers(headers: Mapping[str, str]) -> Mapping[str, str]: + return {header: value for header, value in headers.items() if not _is_credential_header(header)} + + class TenantLoggingMetadata(TypedDict): user_api_key_metadata: ReadOnly[Mapping[str, object] | None] user_api_key_team_metadata: ReadOnly[Mapping[str, object] | None] diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 6ab31e7d4b0..a2e2e36d458 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -86,6 +86,7 @@ from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, valid from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.llms.vllm.common_utils import endpoint_under_api_base from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error +from litellm.passthrough.utils import BasePassthroughUtils from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import ( can_agent_call_model, @@ -1423,6 +1424,12 @@ async def bedrock_passthrough_error_response( raise +def _bedrock_pass_through_headers(request: Request) -> Mapping[str, str]: + from litellm.proxy.litellm_pre_call_utils import without_credential_headers + + return without_credential_headers(BasePassthroughUtils.pass_through_prefixed_headers(request.headers)) + + async def handle_bedrock_passthrough_router_model( model: str, endpoint: str, @@ -1478,6 +1485,7 @@ async def handle_bedrock_passthrough_router_model( "endpoint": endpoint, "data": request_body, "custom_llm_provider": "bedrock", + "request_headers": _bedrock_pass_through_headers(request), } data: Final[dict[str, object]] = guardrail_request_data_with_streaming( MappingProxyType(bedrock_payload), @@ -1684,6 +1692,7 @@ async def bedrock_llm_proxy_route( "endpoint": endpoint, "data": request_body, "custom_llm_provider": "bedrock", + "request_headers": _bedrock_pass_through_headers(request), } data: Final[dict[str, object]] = guardrail_request_data_with_streaming( MappingProxyType(passthrough_payload), diff --git a/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py index f995a7ac83f..d630fa00138 100644 --- a/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py +++ b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py @@ -404,6 +404,13 @@ def _expected_on_the_wire(action: str, request: Request) -> Applied: return Applied("deployment-guardrail", "3", "ENABLED", None) +def _caller_guardrail_on_the_wire(request: Request) -> Applied: + assert {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} <= ( + _signed_header_names(request) + ), request.headers["authorization"] + return Applied("caller-guardrail", "9", _CALLER_INVOKE_HEADERS["x-pass-x-amzn-bedrock-trace"], None) + + @pytest.mark.parametrize("action", _ACTIONS) def test_passthrough_on_a_guardrailed_deployment_blocks_with_the_deployment_guardrail( rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str @@ -478,7 +485,10 @@ def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent( rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str ) -> None: if action.startswith("invoke"): - pytest.skip("BUG: the router-model invoke pass-through drops the caller's X-Amzn-Bedrock-Guardrail* headers") + pytest.skip( + "BUG: the router-model invoke pass-through drops the unprefixed X-Amzn-Bedrock-Guardrail* headers" + " that boto3's native guardrailIdentifier / guardrailVersion / trace args send" + ) marker: Final = uuid.uuid4().hex runtime: Final = rig.runtime(monkeypatch) @@ -491,6 +501,20 @@ def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent( ), request.body +@pytest.mark.parametrize("action", ("invoke", "invoke-with-response-stream")) +def test_deployment_without_a_guardrail_forwards_the_caller_guardrail_sent_as_x_pass_headers( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch, _CALLER_INVOKE_HEADERS) + + outcome: Final = _send(runtime, action, rig.unguarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail=None), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == _caller_guardrail_on_the_wire(request), (request.headers, request.body) + + def test_each_load_balanced_request_carries_the_guardrail_of_the_deployment_it_reached( rig: Rig, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py index 5fd12405ab2..992e18a6d5d 100644 --- a/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py +++ b/tests/integration/providers/test_bedrock_passthrough_request_metadata_wire.py @@ -5,6 +5,7 @@ from urllib.parse import unquote from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from pathlib import Path +from types import MappingProxyType from typing import Final import boto3 @@ -25,6 +26,7 @@ _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 +_LATENCY_HEADER: Final = "x-amzn-bedrock-performanceconfig-latency" _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"}]}}, @@ -105,9 +107,11 @@ class Rig: long_key: str long_identity: tuple[tuple[str, str], ...] - def runtime(self, key: str, monkeypatch: pytest.MonkeyPatch): # noqa: ANN201 # boto3 clients are untyped + def runtime( # noqa: ANN201 # boto3 clients are untyped + self, key: str, monkeypatch: pytest.MonkeyPatch, extra_headers: Mapping[str, str] = MappingProxyType({}) + ): monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", key) - return boto3.client( + client: Final = boto3.client( "bedrock-runtime", endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", region_name="us-east-1", @@ -117,6 +121,11 @@ class Rig: parameter_validation=False, ), ) + client.meta.events.register( + "before-send.bedrock-runtime.*", + lambda request, **_: request.headers.update(extra_headers), + ) + return client def received(self, marker: str) -> tuple[Request, ...]: return tuple(request for request in self.wire.drain() if marker.encode() in request.body) @@ -252,6 +261,54 @@ def test_boto3_passthrough_calls_carry_the_key_identity_in_signed_request_metada ] +def test_caller_request_metadata_header_sent_through_x_pass_never_replaces_the_key_identity( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = f"forged-header-{uuid.uuid4().hex}" + runtime: Final = rig.runtime( + rig.key, + monkeypatch, + { + f"x-pass-{_METADATA_HEADER}": json.dumps({"user_api_key_team_alias": "forged", "ticket": "T-42"}), + f"x-pass-{_LATENCY_HEADER}": "optimized", + }, + ) + model: Final = rig.strict_model + + runtime.converse(modelId=model, messages=_messages(marker)) + tuple(runtime.converse_stream(modelId=model, messages=_messages(marker))["stream"]) + runtime.invoke_model(modelId=model, body=json.dumps(_invoke_body(marker)))["body"].read() + tuple(runtime.invoke_model_with_response_stream(modelId=model, body=json.dumps(_invoke_body(marker)))["body"]) + + received: Final = rig.received(marker) + assert [request.target.rsplit("/", 1)[-1] for request in received] == [ + "converse", + "converse-stream", + "invoke", + "invoke-with-response-stream", + ], received + assert all( + request.headers.get(_LATENCY_HEADER) == "optimized" and _LATENCY_HEADER in _signed_header_names(request) + for request in received + ), [request.headers for request in received] + converse_requests: Final = received[:2] + invoke_requests: Final = received[2:] + assert [_METADATA_HEADER in request.headers for request in converse_requests] == [False, False], [ + request.headers for request in converse_requests + ] + assert [json.loads(request.body)["requestMetadata"] for request in converse_requests] == [ + rig.identity, + rig.identity, + ], [request.body for request in converse_requests] + 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: diff --git a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py index c36644ba536..f6728e2ed10 100644 --- a/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py +++ b/tests/unit/llms/bedrock/chat/invoke_transformations/test_base_invoke_transformation.py @@ -1,4 +1,7 @@ +import base64 import json +import struct +from binascii import crc32 from collections.abc import Mapping from types import MappingProxyType from typing import Final @@ -518,3 +521,136 @@ def test_imported_llama_model_arn_is_encoded_and_max_tokens_become_max_gen_len(r assert response.choices[0].message.content == "Here's a joke..." assert str(route.calls[0].request.url) == expected_url assert json.loads(route.calls[0].request.content) == {"prompt": "Tell me a joke", "max_gen_len": 100} + + +GUARDED_INVOKE_TRACE: Final = MappingProxyType( + { + "guardrail": { + "input": {"caller-guardrail": {"invocationMetrics": {"usage": {"topicPolicyUnits": 1}}}}, + "actionReason": "No action.", + } + } +) +GUARDED_INVOKE_METRICS: Final = MappingProxyType( + {"inputTokenCount": 25, "outputTokenCount": 12, "invocationLatency": 300, "firstByteLatency": 90} +) +NOVA_INVOKE_MODEL: Final = "us.amazon.nova-2-lite-v1:0" +LLAMA_INVOKE_MODEL: Final = "meta.llama3-8b-instruct-v1:0" + + +def _event_frame(payload: Mapping[str, object]) -> bytes: + def header(name: str, value: str) -> bytes: + name_b, value_b = name.encode(), value.encode() + return ( + struct.pack("!B", len(name_b)) + name_b + struct.pack("!B", 7) + struct.pack("!H", len(value_b)) + value_b + ) + + payload_b: Final = json.dumps( + {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()}, separators=(",", ":") + ).encode() + headers_b: Final = ( + header(":event-type", "chunk") + header(":content-type", "application/json") + header(":message-type", "event") + ) + prelude: Final = struct.pack("!II", 12 + len(headers_b) + len(payload_b) + 4, len(headers_b)) + prelude_crc: Final = crc32(prelude) & 0xFFFFFFFF + message: Final = struct.pack("!I", prelude_crc) + headers_b + payload_b + return prelude + message + struct.pack("!I", crc32(message, prelude_crc) & 0xFFFFFFFF) + + +def _guarded_invoke_completion(model: str, reply: httpx.Response, stream: bool) -> object: + return litellm.completion( + model=f"bedrock/invoke/{model}", + messages=[{"role": "user", "content": "hi"}], + max_tokens=20, + stream=stream, + guardrailConfig={"guardrailIdentifier": "caller-guardrail", "guardrailVersion": "1", "trace": "enabled"}, + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(lambda _request: reply))), + aws_access_key_id="fake", + aws_secret_access_key="fake", + aws_region_name="us-east-1", + ) + + +@pytest.mark.parametrize( + ("model", "reply_body"), + [ + pytest.param( + NOVA_INVOKE_MODEL, + { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 25, "outputTokens": 12, "totalTokens": 37}, + }, + id="nova", + ), + pytest.param( + LLAMA_INVOKE_MODEL, + {"generation": "hi", "prompt_token_count": 25, "generation_token_count": 12, "stop_reason": "stop"}, + id="llama", + ), + ], +) +def test_guarded_non_claude_invoke_completion_keeps_the_guardrail_trace_from_the_client( + model: str, reply_body: Mapping[str, object] +) -> None: + reply: Final = httpx.Response(200, json={**reply_body, "amazon-bedrock-trace": dict(GUARDED_INVOKE_TRACE)}) + + response: Final = _guarded_invoke_completion(model, reply, stream=False) + + assert isinstance(response, litellm.ModelResponse) + assert response.choices[0].message.content == "hi" + assert "amazon-bedrock-trace" not in response.model_dump(), response.model_dump() + + +@pytest.mark.parametrize( + ("model", "events"), + [ + pytest.param( + NOVA_INVOKE_MODEL, + ( + {"messageStart": {"role": "assistant"}}, + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hi"}}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + { + "metadata": {"usage": {"inputTokens": 25, "outputTokens": 12, "totalTokens": 37}, "metrics": {}}, + "amazon-bedrock-invocationMetrics": dict(GUARDED_INVOKE_METRICS), + "amazon-bedrock-guardrailAction": "NONE", + "amazon-bedrock-trace": dict(GUARDED_INVOKE_TRACE), + }, + ), + id="nova", + ), + pytest.param( + LLAMA_INVOKE_MODEL, + ( + {"generation": "hi", "prompt_token_count": 25, "generation_token_count": 1, "stop_reason": None}, + { + "generation": "", + "prompt_token_count": None, + "generation_token_count": 12, + "stop_reason": "stop", + "amazon-bedrock-invocationMetrics": dict(GUARDED_INVOKE_METRICS), + "amazon-bedrock-trace": dict(GUARDED_INVOKE_TRACE), + }, + ), + id="llama", + ), + ], +) +def test_guarded_non_claude_invoke_stream_yields_only_the_text_and_stop_chunks( + model: str, events: tuple[Mapping[str, object], ...] +) -> None: + reply: Final = httpx.Response( + 200, + headers={"content-type": "application/vnd.amazon.eventstream"}, + content=b"".join(_event_frame(event) for event in events), + ) + + stream: Final = _guarded_invoke_completion(model, reply, stream=True) + + assert isinstance(stream, litellm.CustomStreamWrapper) + assert [ + ([(choice.delta.content, choice.finish_reason) for choice in chunk.choices], chunk.provider_specific_fields or {}) + for chunk in stream + ] == [([("hi", None)], {}), ([(None, "stop")], {})] diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 44c0e82c006..c4706b17294 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -23,7 +23,8 @@ from litellm.llms.bedrock.passthrough.transformation import ( BedrockPassthroughConfig, is_bedrock_streaming_endpoint, ) -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.passthrough.main import AsyncPassthroughStreamingResponse, PassthroughStreamingResponse from litellm.types.utils import ModelResponse CONVERSE_MODEL = "anthropic.claude-sonnet-4-5-20250929-v1:0" @@ -597,9 +598,13 @@ def _invoke_chunk(payload: dict) -> bytes: return _event_frame("chunk", {"bytes": base64.b64encode(json.dumps(payload).encode()).decode()}) -def _stream_logging_obj(endpoint: str, dynamic_async_success_callbacks: list[CustomLogger] | None = None) -> Logging: +def _stream_logging_obj( + endpoint: str, + dynamic_async_success_callbacks: list[CustomLogger] | None = None, + model: str = CONVERSE_MODEL, +) -> Logging: logging_obj = Logging( - model=f"bedrock/{CONVERSE_MODEL}", + model=f"bedrock/{model}", messages=[], stream=True, call_type="pass_through_endpoint", @@ -902,6 +907,137 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): assert collector.build_logged_response(_converse_stream_logging_obj()) is None +TIERED_MODEL: Final = "amazon.nova-pro-v1:0" +TIERED_INVOKE_ENDPOINT: Final = f"/model/{TIERED_MODEL}/invoke" +TIERED_INVOKE_STREAM_ENDPOINT: Final = f"/model/{TIERED_MODEL}/invoke-with-response-stream" +SERVICE_TIER_HEADER: Final = "x-amzn-bedrock-service-tier" +TIERED_PROMPT_TOKENS: Final = 25 +TIERED_COMPLETION_TOKENS: Final = 12 +SERVED_TIER_PRICES: Final = ( + pytest.param({SERVICE_TIER_HEADER: "priority"}, "_priority", id="priority"), + pytest.param({}, "", id="standard"), +) + + +def _tier_priced_cost(price_suffix: str) -> float: + prices: Final = litellm.model_cost[TIERED_MODEL] + return ( + TIERED_PROMPT_TOKENS * prices[f"input_cost_per_token{price_suffix}"] + + TIERED_COMPLETION_TOKENS * prices[f"output_cost_per_token{price_suffix}"] + ) + + +def _nova_invoke_text_stream(texts: list[str]) -> bytes: + return b"".join( + _invoke_chunk(event) + for event in ( + {"messageStart": {"role": "assistant"}}, + *({"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": text}}} for text in texts), + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + { + "metadata": { + "usage": { + "inputTokens": TIERED_PROMPT_TOKENS, + "outputTokens": TIERED_COMPLETION_TOKENS, + "totalTokens": TIERED_PROMPT_TOKENS + TIERED_COMPLETION_TOKENS, + }, + "metrics": {"latencyMs": 120}, + } + }, + ) + ) + + +def _provider_stream_reply(reply_headers: Mapping[str, str], stream: bytes) -> httpx.Response: + return httpx.Response( + 200, + headers={"content-type": "application/vnd.amazon.eventstream", **reply_headers}, + content=stream, + request=httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com{TIERED_INVOKE_STREAM_ENDPOINT}"), + ) + + +@pytest.mark.parametrize(("reply_headers", "price_suffix"), SERVED_TIER_PRICES) +def test_invoke_reply_is_priced_at_the_service_tier_bedrock_served( + reply_headers: Mapping[str, str], price_suffix: str +) -> None: + reply: Final = httpx.Response( + 200, + json={ + "output": {"message": {"role": "assistant", "content": [{"text": "Invoke response"}]}}, + "stopReason": "end_turn", + "usage": { + "inputTokens": TIERED_PROMPT_TOKENS, + "outputTokens": TIERED_COMPLETION_TOKENS, + "totalTokens": TIERED_PROMPT_TOKENS + TIERED_COMPLETION_TOKENS, + }, + }, + headers=reply_headers, + ) + + logged: Final = BedrockPassthroughConfig().logging_non_streaming_response( + model=TIERED_MODEL, + custom_llm_provider="bedrock", + httpx_response=reply, + request_data={}, + logging_obj=_stream_logging_obj(TIERED_INVOKE_ENDPOINT, model=TIERED_MODEL), + endpoint=TIERED_INVOKE_ENDPOINT, + ) + + assert litellm.completion_cost(model=f"bedrock/{TIERED_MODEL}", completion_response=logged) == pytest.approx( + _tier_priced_cost(price_suffix) + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("reply_headers", "price_suffix"), SERVED_TIER_PRICES) +async def test_invoke_stream_relay_is_priced_at_the_service_tier_bedrock_served( + reply_headers: Mapping[str, str], price_suffix: str +) -> None: + callback: Final = _CaptureAsyncSuccessCallback() + stream: Final = _nova_invoke_text_stream(["Invoke ", "response"]) + reply: Final = _provider_stream_reply(reply_headers, stream) + + async def provider_reply() -> httpx.Response: + return reply + + relay: Final = AsyncPassthroughStreamingResponse( + response=provider_reply(), + litellm_logging_obj=_stream_logging_obj(TIERED_INVOKE_STREAM_ENDPOINT, [callback], model=TIERED_MODEL), + provider_config=BedrockPassthroughConfig(), + ) + relayed: Final = b"".join([chunk async for chunk in relay]) + + assert relayed == stream + assert await relay.wait_for_spend_flush() is True + assert callback.kwargs is not None + assert callback.kwargs["response_cost"] == pytest.approx(_tier_priced_cost(price_suffix)) + + +class _ImmediateExecutor: + def submit(self, fn: Callable[..., object], /, *args: object, **kwargs: object) -> None: + fn(*args, **kwargs) + + +@pytest.mark.parametrize(("reply_headers", "price_suffix"), SERVED_TIER_PRICES) +def test_sync_invoke_stream_relay_is_priced_at_the_service_tier_bedrock_served( + reply_headers: Mapping[str, str], price_suffix: str +) -> None: + logging_obj: Final = _stream_logging_obj(TIERED_INVOKE_STREAM_ENDPOINT, model=TIERED_MODEL) + stream: Final = _nova_invoke_text_stream(["Invoke ", "response"]) + + with patch("litellm.utils.executor", _ImmediateExecutor()): + relayed: Final = b"".join( + PassthroughStreamingResponse( + _provider_stream_reply(reply_headers, stream), logging_obj, BedrockPassthroughConfig() + ) + ) + + assert relayed == stream + assert logging_obj.model_call_details["response_cost"] == pytest.approx(_tier_priced_cost(price_suffix)) + + 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 = { @@ -977,7 +1113,11 @@ async def test_converse_passthrough_signs_proxy_identity_into_request_metadata( 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 + action, + caller_body, + PROXY_IDENTITY, + request_query_params=query_params, + request_headers={f"x-pass-{REQUEST_METADATA_HEADER.upper()}": '{"user_api_key_team_alias":"forged"}'}, ) assert json.loads(request.content) == { @@ -1281,12 +1421,14 @@ async def test_deployment_guardrail_replaces_the_caller_converse_guardrail( {**CONVERSE_BODY, "guardrailConfig": caller_guardrail}, {}, deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS, ) assert json.loads(request.content) == { **CONVERSE_BODY, "guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}, } + assert _invoke_guardrail_headers(request) == {} @pytest.mark.asyncio @@ -1328,6 +1470,59 @@ async def test_passthrough_without_a_deployment_guardrail_sends_the_caller_guard } +CALLER_GUARDRAIL_IDENTITY_HEADERS: Final[Mapping[str, str]] = { + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "x-pass-X-Amzn-Bedrock-GuardrailVersion": "9", +} +SENT_CALLER_GUARDRAIL_IDENTITY: Final[Mapping[str, str]] = { + "x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-amzn-bedrock-guardrailversion": "9", +} + + +SDK_INVOKE_MODEL: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0" + + +def _sdk_passthrough_call( + model: str, action: str, request_headers: Mapping[str, str], reply_body: Mapping[str, object] +) -> tuple[httpx.Request, object]: + sent: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json=reply_body, request=request) + + response: Final = litellm.llm_passthrough_route( + model=f"bedrock/{model}", + method="POST", + endpoint=f"/model/{model}/{action}", + json=INVOKE_BODY, + request_headers=dict(request_headers), + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + aws_access_key_id="AKIASDKKEY", + aws_secret_access_key="sdk-secret", + aws_region_name="us-east-1", + ) + return sent[0], response + + +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize("caller_trace", [None, "DISABLED", "enabled"]) +def test_sdk_invoke_passthrough_sends_the_caller_guardrail_trace_exactly_as_sent( + action: str, caller_trace: str | None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + caller_trace_header: Final = {} if caller_trace is None else {"x-pass-X-Amzn-Bedrock-Trace": caller_trace} + sent_trace_header: Final = {} if caller_trace is None else {"x-amzn-bedrock-trace": caller_trace} + + request, _ = _sdk_passthrough_call( + SDK_INVOKE_MODEL, action, {**CALLER_GUARDRAIL_IDENTITY_HEADERS, **caller_trace_header}, {"ok": True} + ) + + assert _invoke_guardrail_headers(request) == {**SENT_CALLER_GUARDRAIL_IDENTITY, **sent_trace_header} + assert json.loads(request.content) == INVOKE_BODY + + @pytest.mark.asyncio @pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) @pytest.mark.parametrize("guardrail_config", ["deployment-guardrail", {"guardrailVersion": "3"}]) @@ -1368,3 +1563,244 @@ async def test_passthrough_leaves_actions_without_a_guardrail_field_as_sent(monk assert json.loads(request.content) == count_tokens_body assert _invoke_guardrail_headers(request) == {} + + +GUARDRAIL_UNIT_PRICE: Final = 0.00015 +GUARDRAIL_PRICING: Final = { + "litellm_provider": "bedrock", + "mode": "guardrail", + "guardrail_cost_per_unit": {"topicPolicyUnits": GUARDRAIL_UNIT_PRICE, "contentPolicyUnits": GUARDRAIL_UNIT_PRICE}, +} +_ASSESSMENT_USAGE: Final = {"invocationMetrics": {"usage": {"topicPolicyUnits": 1, "contentPolicyUnits": 1}}} +INVOKE_GUARDRAIL_TRACE: Final = { + "guardrail": { + "input": {"caller-guardrail": _ASSESSMENT_USAGE}, + "outputs": [{"caller-guardrail": _ASSESSMENT_USAGE}], + "actionReason": "No action.", + } +} +BILLED_GUARDRAIL_USAGE: Final = {"topicPolicyUnits": 2, "contentPolicyUnits": 2} +BILLED_GUARDRAIL_COST: Final = 4 * GUARDRAIL_UNIT_PRICE +NOVA_INVOKE_MODEL: Final = "us.amazon.nova-2-lite-v1:0" +LLAMA_INVOKE_MODEL: Final = "meta.llama3-8b-instruct-v1:0" +INVOKE_METRICS: Final = { + "inputTokenCount": 25, + "outputTokenCount": 12, + "invocationLatency": 300, + "firstByteLatency": 90, +} + +INVOKE_REPLIES_WITH_A_GUARDRAIL_TRACE: Final = ( + pytest.param( + NOVA_INVOKE_MODEL, + { + "output": {"message": {"role": "assistant", "content": [{"text": "hi"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 25, "outputTokens": 12, "totalTokens": 37}, + "amazon-bedrock-guardrailAction": "NONE", + }, + id="nova", + ), + pytest.param( + LLAMA_INVOKE_MODEL, + {"generation": "hi", "prompt_token_count": 25, "generation_token_count": 12, "stop_reason": "stop"}, + id="llama", + ), + pytest.param( + "mistral.mistral-7b-instruct-v0:2", {"outputs": [{"text": "hi", "stop_reason": "stop"}]}, id="mistral" + ), + pytest.param( + CONVERSE_MODEL, + { + "id": "msg-1", + "type": "message", + "role": "assistant", + "model": CONVERSE_MODEL, + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 25, "output_tokens": 12}, + }, + id="claude", + ), +) + +INVOKE_STREAMS_WITH_A_GUARDRAIL_TRACE: Final = ( + pytest.param( + NOVA_INVOKE_MODEL, + ( + {"messageStart": {"role": "assistant"}}, + {"contentBlockDelta": {"contentBlockIndex": 0, "delta": {"text": "hi"}}}, + {"contentBlockStop": {"contentBlockIndex": 0}}, + {"messageStop": {"stopReason": "end_turn"}}, + { + "metadata": { + "usage": {"inputTokens": 25, "outputTokens": 12, "totalTokens": 37}, + "metrics": {}, + "trace": {}, + }, + "amazon-bedrock-invocationMetrics": INVOKE_METRICS, + "amazon-bedrock-guardrailAction": "NONE", + "amazon-bedrock-trace": INVOKE_GUARDRAIL_TRACE, + }, + ), + id="nova", + ), + pytest.param( + LLAMA_INVOKE_MODEL, + ( + {"generation": "hi", "prompt_token_count": 25, "generation_token_count": 1, "stop_reason": None}, + { + "generation": "", + "prompt_token_count": None, + "generation_token_count": 12, + "stop_reason": "stop", + "amazon-bedrock-invocationMetrics": INVOKE_METRICS, + "amazon-bedrock-trace": INVOKE_GUARDRAIL_TRACE, + }, + ), + id="llama", + ), + pytest.param( + CONVERSE_MODEL, + ( + { + "type": "message_start", + "message": { + "id": "msg-1", + "type": "message", + "role": "assistant", + "model": CONVERSE_MODEL, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 25, "output_tokens": 1}, + }, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hi"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 12}}, + { + "type": "message_stop", + "amazon-bedrock-invocationMetrics": INVOKE_METRICS, + "amazon-bedrock-trace": INVOKE_GUARDRAIL_TRACE, + }, + ), + id="claude", + ), +) + + +def _billed_guardrails(standard_logging_object: object) -> list[tuple[object, object]]: + assert isinstance(standard_logging_object, dict) + return [ + (entry["guardrail_usage"], entry["guardrail_cost"]) + for entry in standard_logging_object["guardrail_information"] or [] + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "reply_body"), INVOKE_REPLIES_WITH_A_GUARDRAIL_TRACE) +async def test_invoke_reply_bills_the_guardrail_bedrock_traced( + model: str, reply_body: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost, "bedrock/guardrails", GUARDRAIL_PRICING) + endpoint: Final = f"/model/{model}/invoke" + callback: Final = _CaptureAsyncSuccessCallback() + logging_obj: Final = Logging( + model=model, + messages=[], + stream=False, + call_type="allm_passthrough_route", + start_time=datetime(2025, 1, 1), + litellm_call_id="call-1", + function_id="fn-1", + dynamic_async_success_callbacks=[callback], + ) + logging_obj.update_environment_variables(litellm_params={}, optional_params={}) + logging_obj.model_call_details["custom_llm_provider"] = "bedrock" + logging_obj.model_call_details["endpoint"] = endpoint + reply: Final = httpx.Response( + 200, + json={**reply_body, "amazon-bedrock-trace": INVOKE_GUARDRAIL_TRACE}, + request=httpx.Request("POST", f"https://bedrock-runtime.us-east-1.amazonaws.com{endpoint}"), + ) + + await logging_obj.async_success_handler( + result=reply, start_time=datetime(2025, 1, 1), end_time=datetime(2025, 1, 1, 0, 0, 1) + ) + + assert callback.kwargs is not None + assert _billed_guardrails(callback.kwargs["standard_logging_object"]) == [ + (BILLED_GUARDRAIL_USAGE, pytest.approx(BILLED_GUARDRAIL_COST)) + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("model", "events"), INVOKE_STREAMS_WITH_A_GUARDRAIL_TRACE) +async def test_invoke_stream_relay_bills_the_guardrail_bedrock_traced( + model: str, events: tuple[dict[str, object], ...], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem(litellm.model_cost, "bedrock/guardrails", GUARDRAIL_PRICING) + endpoint: Final = f"/model/{model}/invoke-with-response-stream" + callback: Final = _CaptureAsyncSuccessCallback() + stream: Final = b"".join(_invoke_chunk(event) for event in events) + + async def provider_reply() -> httpx.Response: + return _provider_stream_reply({}, stream) + + relay: Final = AsyncPassthroughStreamingResponse( + response=provider_reply(), + litellm_logging_obj=_stream_logging_obj(endpoint, [callback], model=model), + provider_config=BedrockPassthroughConfig(), + ) + relayed: Final = b"".join([chunk async for chunk in relay]) + + assert relayed == stream + assert await relay.wait_for_spend_flush() is True + assert callback.kwargs is not None + assert _billed_guardrails(callback.kwargs["standard_logging_object"]) == [ + (BILLED_GUARDRAIL_USAGE, pytest.approx(BILLED_GUARDRAIL_COST)) + ] + + +class _CaptureSyncSuccessCallback(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.standard_logging_objects: list[object] = [] + + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.standard_logging_objects.append(kwargs.get("standard_logging_object")) + + +@pytest.mark.parametrize( + ("model", "reply_body"), + [param for param in INVOKE_REPLIES_WITH_A_GUARDRAIL_TRACE if param.id in {"nova", "claude"}], +) +def test_sdk_invoke_passthrough_bills_the_caller_guardrail_bedrock_traced_once( + model: str, reply_body: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "bedrock_request_metadata_fields", None) + monkeypatch.setitem(litellm.model_cost, "bedrock/guardrails", GUARDRAIL_PRICING) + callback: Final = _CaptureSyncSuccessCallback() + monkeypatch.setattr(litellm, "callbacks", [callback]) + traced_reply: Final = {**reply_body, "amazon-bedrock-trace": INVOKE_GUARDRAIL_TRACE} + + with patch("litellm.utils.executor", _ImmediateExecutor()): + request, response = _sdk_passthrough_call( + model, + "invoke", + {**CALLER_GUARDRAIL_IDENTITY_HEADERS, "x-pass-X-Amzn-Bedrock-Trace": "ENABLED"}, + traced_reply, + ) + + assert isinstance(response, httpx.Response) + assert response.json() == traced_reply + assert _invoke_guardrail_headers(request) == { + **SENT_CALLER_GUARDRAIL_IDENTITY, + "x-amzn-bedrock-trace": "ENABLED", + } + assert [_billed_guardrails(logged) for logged in callback.standard_logging_objects] == [ + [(BILLED_GUARDRAIL_USAGE, pytest.approx(BILLED_GUARDRAIL_COST))] + ] diff --git a/tests/unit/passthrough/test_utils.py b/tests/unit/passthrough/test_utils.py new file mode 100644 index 00000000000..02a1c70f8cb --- /dev/null +++ b/tests/unit/passthrough/test_utils.py @@ -0,0 +1,78 @@ +from typing import Final + +import httpx +import pytest + +from litellm.passthrough.utils import _PASS_THROUGH_PROTECTED_HEADERS, BasePassthroughUtils + + +def test_pass_through_prefixed_headers_keeps_only_the_x_pass_headers_forwarded_upstream() -> None: + inbound: Final = { + "authorization": "Bearer sk-litellm-key", + "x-litellm-api-key": "sk-litellm-key", + "content-type": "application/json", + "x-amzn-bedrock-guardrailidentifier": "native-guardrail", + "x-passthrough-id": "not-a-pass-header", + "x-pass-authorization": "Bearer upstream-secret", + "X-Pass-X-Api-Key": "upstream-api-key", + "x-pass-X-Amz-Security-Token": "forged-session-token", + "x-pass-x-amzn-bedrock-trace": "DISABLED", + "X-Pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "X-PASS-ANTHROPIC-BETA": "beta-1", + } + + selected: Final = BasePassthroughUtils.pass_through_prefixed_headers(inbound) + + assert dict(selected) == { + "x-pass-x-amzn-bedrock-trace": "DISABLED", + "X-Pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "X-PASS-ANTHROPIC-BETA": "beta-1", + } + + +@pytest.mark.parametrize( + "header_name", + [ + *(f"x-pass-{name}" for name in sorted(_PASS_THROUGH_PROTECTED_HEADERS)), + "X-Pass-Authorization", + "x-pass-X-Amz-Date", + "x-pass-x-amz-security-token", + "x-pass-X-Amzn-Bedrock-Trace", + "x-pass-anthropic-beta", + "x-pass-cookie", + "x-passthrough-id", + "authorization", + ], +) +def test_pass_through_prefixed_headers_select_exactly_the_headers_forwarded_upstream(header_name: str) -> None: + inbound: Final = {header_name: "caller-value"} + + selected: Final = BasePassthroughUtils.pass_through_prefixed_headers(inbound) + forwarded: Final = BasePassthroughUtils.forward_headers_from_request( + request_headers=dict(inbound), headers={}, forward_headers=False + ) + + assert {name[len("x-pass-") :].lower(): value for name, value in selected.items()} == forwarded + + +@pytest.mark.parametrize( + ("default_query_params", "expected"), + [ + pytest.param( + {"api-version": "default", "region": "default"}, + {"api-version": "url", "region": "default", "shared": "request", "tag": ["a", "b"]}, + id="with_defaults", + ), + pytest.param(None, {"api-version": "url", "shared": "request", "tag": ["a", "b"]}, id="without_defaults"), + ], +) +def test_merged_query_parameters_let_the_request_override_the_url_and_the_url_override_defaults( + default_query_params: dict[str, str | list] | None, expected: dict[str, str | list[str]] +) -> None: + merged: Final = BasePassthroughUtils.get_merged_query_parameters( + existing_url=httpx.URL("https://upstream.test/v1?api-version=url&shared=url&tag=a&tag=b"), + request_query_params={"shared": "request"}, + default_query_params=default_query_params, + ) + + assert merged == expected diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 6e551f8ff0a..8e88fd7449a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1350,11 +1350,23 @@ async def test_mistral_passthrough_accepts_multipart_without_json_parsing(): assert captured_kwargs["custom_headers"] == {"Authorization": "Bearer mistral-test-key"} +BEDROCK_PROXY_MODEL: Final = "anthropic.claude-sonnet-5-5" +BEDROCK_PROXY_CALLER_BEARER: Final = "sk-caller-virtual-key" +BEDROCK_PROXY_ROUTES: Final = pytest.mark.parametrize( + ("path_model", "uses_router"), + [ + pytest.param("bedrock-sonnet", True, id="router_model"), + pytest.param(BEDROCK_PROXY_MODEL, False, id="direct_model"), + ], +) + + class TestBedrockLLMProxyRoute: @pytest.mark.asyncio async def test_bedrock_llm_proxy_route_application_inference_profile(self): mock_request = Mock() mock_request.method = "POST" + mock_request.headers = {} mock_response = Mock() mock_user_api_key_dict = Mock() mock_request_body = {"messages": [{"role": "user", "content": "test"}]} @@ -1397,6 +1409,7 @@ class TestBedrockLLMProxyRoute: async def test_bedrock_llm_proxy_route_regular_model(self): mock_request = Mock() mock_request.method = "POST" + mock_request.headers = {} mock_response = Mock() mock_user_api_key_dict = Mock() mock_request_body = {"messages": [{"role": "user", "content": "test"}]} @@ -1444,6 +1457,7 @@ class TestBedrockLLMProxyRoute: ) -> None: mock_request: Final = Mock() mock_request.method = "POST" + mock_request.headers = {} mock_processor: Final = Mock() mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") request_body: Final = {"messages": [{"role": "user", "content": "test"}]} @@ -1484,6 +1498,7 @@ class TestBedrockLLMProxyRoute: ) -> None: mock_request: Final = Mock() mock_request.method = "POST" + mock_request.headers = {} mock_processor: Final = Mock() mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") request_body: Final = {"messages": [{"role": "user", "content": "test"}]} @@ -1515,6 +1530,91 @@ class TestBedrockLLMProxyRoute: processing_data: Final = processor_constructor.call_args.kwargs["data"] _assert_bedrock_processing_data_classification(processing_data, is_streaming) + @BEDROCK_PROXY_ROUTES + def test_invoke_signs_the_caller_x_pass_header_and_sends_no_caller_credential( + self, monkeypatch: pytest.MonkeyPatch, httpx_transport: None, path_model: str, uses_router: bool + ) -> None: + caller_cookie: Final = "proxy_session=caller-secret" + + sent, _ = _bedrock_invoke_through_proxy( + monkeypatch, + path_model, + uses_router, + { + "x-pass-X-Amzn-Bedrock-Service-Tier": "priority", + "Authorization": f"Bearer {BEDROCK_PROXY_CALLER_BEARER}", + "Cookie": caller_cookie, + }, + ) + + signed_headers: Final = sent.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";") + assert sent.headers["x-amzn-bedrock-service-tier"] == "priority" + assert "x-amzn-bedrock-service-tier" in signed_headers + assert [name for name in sent.headers if name.startswith("x-pass-") or name == "cookie"] == [] + assert [ + value for value in sent.headers.values() if BEDROCK_PROXY_CALLER_BEARER in value or caller_cookie in value + ] == [] + + @BEDROCK_PROXY_ROUTES + @pytest.mark.parametrize("caller_trace", [None, "DISABLED", "enabled"]) + def test_invoke_sends_the_caller_guardrail_trace_exactly_as_sent( + self, + monkeypatch: pytest.MonkeyPatch, + httpx_transport: None, + path_model: str, + uses_router: bool, + caller_trace: str | None, + ) -> None: + trace_header: Final = {} if caller_trace is None else {"x-pass-X-Amzn-Bedrock-Trace": caller_trace} + + sent, _ = _bedrock_invoke_through_proxy( + monkeypatch, + path_model, + uses_router, + { + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "x-pass-X-Amzn-Bedrock-GuardrailVersion": "9", + **trace_header, + }, + ) + + assert ( + sent.headers["x-amzn-bedrock-guardrailidentifier"], + sent.headers["x-amzn-bedrock-guardrailversion"], + sent.headers.get("x-amzn-bedrock-trace"), + ) == ("caller-guardrail", "9", caller_trace) + + @BEDROCK_PROXY_ROUTES + def test_invoke_keeps_protected_x_pass_values_from_bedrock_guardrails_and_the_logged_body( + self, monkeypatch: pytest.MonkeyPatch, httpx_transport: None, path_model: str, uses_router: bool + ) -> None: + protected_values: Final = ( + "Bearer upstream-secret", + "upstream-api-key", + "forged-session-token", + "session=upstream-cookie", + "Basic upstream-proxy-secret", + ) + + sent, guardrail_data = _bedrock_invoke_through_proxy( + monkeypatch, + path_model, + uses_router, + { + "x-pass-authorization": protected_values[0], + "x-pass-x-api-key": protected_values[1], + "x-pass-X-Amz-Security-Token": protected_values[2], + "x-pass-cookie": protected_values[3], + "x-pass-proxy-authorization": protected_values[4], + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + }, + ) + + logged_body: Final = json.dumps(guardrail_data["proxy_server_request"]["body"], default=str) + assert sent.headers["x-amzn-bedrock-guardrailidentifier"] == "caller-guardrail" + assert guardrail_data["request_headers"] == {"x-pass-x-amzn-bedrock-guardrailidentifier": "caller-guardrail"} + assert [value for value in protected_values if value in str(sent.headers.raw) or value in logged_body] == [] + @pytest.mark.asyncio @pytest.mark.parametrize( ("model_id", "action"), @@ -1526,6 +1626,7 @@ class TestBedrockLLMProxyRoute: async def test_bedrock_direct_model_id_does_not_imply_streaming(self, model_id: str, action: str) -> None: mock_request: Final = Mock() mock_request.method = "POST" + mock_request.headers = {} mock_processor: Final = Mock() mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") request_body: Final = {"messages": [{"role": "user", "content": "test"}]} @@ -1562,6 +1663,7 @@ class TestBedrockLLMProxyRoute: async def test_bedrock_router_model_id_does_not_imply_streaming(self, model_id: str, action: str) -> None: mock_request: Final = Mock() mock_request.method = "POST" + mock_request.headers = {} mock_processor: Final = Mock() mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success") request_body: Final = {"messages": [{"role": "user", "content": "test"}]} @@ -1865,6 +1967,84 @@ class TestBedrockLLMProxyRoute: assert "Blocked by guardrail" in str(exc_info.value.detail) +class _RequestDataCaptureGuardrail(CustomGuardrail): + def __init__(self) -> None: + super().__init__(guardrail_name="request-data-capture", default_on=True, event_hook="pre_call") + self.data: dict[str, object] = {} + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + self.data = data + return data + + +def _bedrock_invoke_through_proxy( + monkeypatch: pytest.MonkeyPatch, path_model: str, uses_router: bool, headers: Mapping[str, str] +) -> tuple[httpx.Request, dict[str, object]]: + from litellm.proxy import proxy_server + + for ambient in (*TestBedrockAgentRuntimePassthroughVirtualKeyLeak.AMBIENT_AWS_ENV, "AWS_BEDROCK_RUNTIME_ENDPOINT"): + monkeypatch.delenv(ambient, raising=False) + monkeypatch.setenv("AWS_ACCESS_KEY_ID", "AKIATESTACCESSKEY") + monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret") + monkeypatch.setenv("AWS_REGION_NAME", "us-east-1") + router: Final = ( + litellm.Router( + model_list=[ + { + "model_name": "bedrock-sonnet", + "litellm_params": {"model": f"bedrock/{BEDROCK_PROXY_MODEL}", "aws_region_name": "us-east-1"}, + } + ] + ) + if uses_router + else None + ) + guardrail: Final = _RequestDataCaptureGuardrail() + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + monkeypatch.setitem( + proxy_server.app.dependency_overrides, + user_api_key_auth, + lambda: UserAPIKeyAuth(api_key=BEDROCK_PROXY_CALLER_BEARER), + ) + + with respx.mock(assert_all_called=True) as upstream: + invoke: Final = upstream.post( + f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{BEDROCK_PROXY_MODEL}/invoke" + ).mock( + return_value=httpx.Response( + 200, + json={ + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": BEDROCK_PROXY_MODEL, + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 5, "output_tokens": 1}, + }, + ) + ) + response: Final = TestClient(proxy_server.app).post( + f"/bedrock/model/{path_model}/invoke", + json={ + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + }, + headers=dict(headers), + ) + + assert response.status_code == 200, response.text + return invoke.calls.last.request, guardrail.data + + _THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/" _THROTTLE_BODY: Final = b'{"message":"Too many requests, please wait before trying again."}' diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 28d2a4fa7ad..e8ee31b27f5 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -934,6 +934,38 @@ async def test_add_litellm_data_to_request_body_snapshot_excludes_secret_fields( ) +@pytest.mark.asyncio +async def test_chat_completions_body_field_named_request_headers_stays_in_the_logged_body(): + from litellm.proxy.proxy_server import ProxyConfig + + client_field: Final = {"x-client-trace": "client-supplied"} + request: Final = Request( + { + "type": "http", + "method": "POST", + "path": "/chat/completions", + "scheme": "http", + "server": ("localhost", 4000), + "query_string": b"", + "headers": [(b"authorization", b"Bearer sk-litellm-key")], + } + ) + + enriched: Final = await add_litellm_data_to_request( + data={ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hello"}], + "request_headers": client_field, + }, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + + assert enriched["proxy_server_request"]["body"]["request_headers"] == client_field + + @pytest.mark.asyncio async def test_add_litellm_data_to_request_body_snapshot_excludes_proxy_server_request(): """Regression: the body snapshot used to include the proxy_server_request