mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
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>
This commit is contained in:
norkelbyte
GitHub
parent
1d5e0c4be5
commit
845f109524
13 changed files with 1081 additions and 43 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")], {})]
|
||||
|
|
|
|||
|
|
@ -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))]
|
||||
]
|
||||
|
|
|
|||
78
tests/unit/passthrough/test_utils.py
Normal file
78
tests/unit/passthrough/test_utils.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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."}'
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue