mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
fix(bedrock): skip the SigV4 credential chain when a bearer token is configured
A deployment authenticating with api_key or AWS_BEARER_TOKEN_BEDROCK still ran boto3's credential chain before every call, so an unloadable default profile (a login_session profile without botocore[crt]) made Converse, embeddings, image generation, image edit, and the Bedrock guardrail hook fail with MissingDependencyException even though the bearer token alone signs the request. The chain now runs only when no bearer token is configured
This commit is contained in:
parent
ba2e5d2d8e
commit
2aa005fed2
13 changed files with 250 additions and 92 deletions
|
|
@ -49,11 +49,16 @@ SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz
|
|||
|
||||
|
||||
class Boto3CredentialsInfo(BaseModel):
|
||||
credentials: Credentials
|
||||
credentials: Credentials | None
|
||||
aws_region_name: str
|
||||
aws_bedrock_runtime_endpoint: str | None
|
||||
|
||||
|
||||
def bedrock_bearer_token(api_key: str | None) -> str | None:
|
||||
token: Final = api_key if api_key is not None else get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
return token or None
|
||||
|
||||
|
||||
class _WebIdentityTokenClaims(BaseModel):
|
||||
aud: str | list[str] | None = None
|
||||
iss: str | None = None
|
||||
|
|
@ -1388,7 +1393,7 @@ class BaseAWSLLM:
|
|||
return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}"
|
||||
|
||||
def _get_boto_credentials_from_optional_params(
|
||||
self, optional_params: dict, model: str | None = None
|
||||
self, optional_params: dict, model: str | None = None, bearer_token: str | None = None
|
||||
) -> Boto3CredentialsInfo:
|
||||
"""
|
||||
Get boto3 credentials from optional params
|
||||
|
|
@ -1420,17 +1425,21 @@ class BaseAWSLLM:
|
|||
) # https://bedrock-runtime.{region_name}.amazonaws.com
|
||||
aws_external_id: Final = optional_params.pop("aws_external_id", None)
|
||||
|
||||
credentials: Final[Credentials] = self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
credentials: Final[Credentials | None] = (
|
||||
None
|
||||
if bearer_token is not None
|
||||
else self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
)
|
||||
)
|
||||
|
||||
return Boto3CredentialsInfo(
|
||||
|
|
@ -1451,14 +1460,9 @@ class BaseAWSLLM:
|
|||
api_key: str | None = None,
|
||||
supports_bearer_token: bool = True,
|
||||
) -> AWSPreparedRequest:
|
||||
if not supports_bearer_token:
|
||||
aws_bearer_token: str | None = None
|
||||
elif api_key is not None:
|
||||
aws_bearer_token = api_key
|
||||
else:
|
||||
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
aws_bearer_token: Final = bedrock_bearer_token(api_key) if supports_bearer_token else None
|
||||
|
||||
if aws_bearer_token:
|
||||
if aws_bearer_token is not None:
|
||||
try:
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
|
|
@ -1555,13 +1559,9 @@ class BaseAWSLLM:
|
|||
Returns:
|
||||
Tuple[dict, Optional[str]]: A tuple containing the headers and the json str body of the request
|
||||
"""
|
||||
if api_key is not None:
|
||||
aws_bearer_token: str | None = api_key
|
||||
else:
|
||||
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
aws_bearer_token: Final = bedrock_bearer_token(api_key)
|
||||
|
||||
# If aws bearer token is set, use it directly in the header
|
||||
if aws_bearer_token:
|
||||
if aws_bearer_token is not None:
|
||||
headers = headers or {}
|
||||
headers["Content-Type"] = "application/json"
|
||||
headers["Authorization"] = f"Bearer {aws_bearer_token}"
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ from litellm.rust_bridge.chat_completions import rust_chat_completions_accepts
|
|||
from litellm.types.utils import ModelResponse
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
|
||||
from ..common_utils import BedrockError, _get_all_bedrock_regions
|
||||
from .invoke_handler import AWSEventStreamDecoder, MockResponseIterator, make_call
|
||||
|
||||
|
|
@ -349,17 +349,21 @@ class BedrockConverseLLM(BaseAWSLLM):
|
|||
|
||||
litellm_params["aws_region_name"] = aws_region_name # [DO NOT DELETE] important for async calls
|
||||
|
||||
credentials: Final[Credentials | None] = self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
credentials: Final[Credentials | None] = (
|
||||
None
|
||||
if bedrock_bearer_token(api_key) is not None
|
||||
else self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
)
|
||||
)
|
||||
|
||||
### SET RUNTIME ENDPOINT ###
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import copy
|
|||
import json
|
||||
import urllib.parse
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, Final, get_args
|
||||
from typing import TYPE_CHECKING, Final, get_args
|
||||
|
||||
import httpx
|
||||
|
||||
|
|
@ -26,7 +26,7 @@ from litellm.types.llms.bedrock import (
|
|||
)
|
||||
from litellm.types.utils import EmbeddingResponse, LlmProviders
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..base_aws_llm import BaseAWSLLM, Credentials, bedrock_bearer_token
|
||||
from ..common_utils import BedrockError
|
||||
from .amazon_nova_transformation import AmazonNovaEmbeddingConfig
|
||||
from .amazon_titan_g1_transformation import AmazonTitanG1Config
|
||||
|
|
@ -45,11 +45,8 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
def _load_credentials(
|
||||
self,
|
||||
optional_params: dict,
|
||||
) -> tuple[Any, str]:
|
||||
try:
|
||||
from botocore.credentials import Credentials
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
bearer_token: str | None = None,
|
||||
) -> tuple[Credentials | None, str]:
|
||||
## CREDENTIALS ##
|
||||
# pop aws_secret_access_key, aws_access_key_id, aws_session_token, aws_region_name from kwargs, since completion calls fail with them
|
||||
aws_secret_access_key: Final = optional_params.pop("aws_secret_access_key", None)
|
||||
|
|
@ -78,17 +75,21 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
if aws_region_name is None:
|
||||
aws_region_name = "us-west-2"
|
||||
|
||||
credentials: Final[Credentials] = self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
credentials: Final[Credentials | None] = (
|
||||
None
|
||||
if bearer_token is not None
|
||||
else self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
)
|
||||
)
|
||||
return credentials, aws_region_name
|
||||
|
||||
|
|
@ -233,7 +234,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
client: HTTPHandler | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
batch_data: list[dict],
|
||||
credentials: Any,
|
||||
credentials: Credentials | None,
|
||||
extra_headers: dict | None,
|
||||
endpoint_url: str,
|
||||
aws_region_name: str,
|
||||
|
|
@ -301,7 +302,7 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
client: AsyncHTTPHandler | None,
|
||||
timeout: float | httpx.Timeout | None,
|
||||
batch_data: list[dict],
|
||||
credentials: Any,
|
||||
credentials: Credentials | None,
|
||||
extra_headers: dict | None,
|
||||
endpoint_url: str,
|
||||
aws_region_name: str,
|
||||
|
|
@ -383,7 +384,9 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
litellm_params: dict,
|
||||
api_key: str | None = None,
|
||||
) -> EmbeddingResponse:
|
||||
credentials, aws_region_name = self._load_credentials(optional_params)
|
||||
credentials, aws_region_name = self._load_credentials(
|
||||
optional_params, bearer_token=bedrock_bearer_token(api_key)
|
||||
)
|
||||
|
||||
### TRANSFORMATION ###
|
||||
unencoded_model_id: Final = optional_params.pop("model_id", None) or model # default to model if not passed
|
||||
|
|
@ -595,8 +598,11 @@ class BedrockEmbedding(BaseAWSLLM):
|
|||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.exceptions import NoCredentialsError
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
if credentials is None:
|
||||
raise NoCredentialsError()
|
||||
|
||||
# Create AWSRequest with GET method and encoded URL
|
||||
request: Final = AWSRequest(
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token
|
||||
from ..common_utils import BedrockError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -198,7 +198,9 @@ class BedrockImageEdit(BaseAWSLLM):
|
|||
Returns:
|
||||
BedrockImageEditPreparedRequest: The prepared request object
|
||||
"""
|
||||
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)
|
||||
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(
|
||||
optional_params, model, bearer_token=bedrock_bearer_token(api_key)
|
||||
)
|
||||
|
||||
# Use the existing ARN-aware provider detection method
|
||||
bedrock_provider: Final = self.get_bedrock_invoke_provider(model)
|
||||
|
|
|
|||
|
|
@ -29,7 +29,7 @@ from litellm.llms.custom_httpx.http_handler import (
|
|||
)
|
||||
from litellm.types.utils import ImageResponse
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM
|
||||
from ..base_aws_llm import BaseAWSLLM, bedrock_bearer_token
|
||||
from ..common_utils import BedrockError
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -220,7 +220,9 @@ class BedrockImageGeneration(BaseAWSLLM):
|
|||
prepped (httpx.Request): The prepared request object
|
||||
body (bytes): The request body
|
||||
"""
|
||||
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(optional_params, model)
|
||||
boto3_credentials_info: Final = self._get_boto_credentials_from_optional_params(
|
||||
optional_params, model, bearer_token=bedrock_bearer_token(api_key)
|
||||
)
|
||||
|
||||
# Use the existing ARN-aware provider detection method
|
||||
bedrock_provider: Final = self.get_bedrock_invoke_provider(model)
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ from litellm.llms.anthropic.chat.guardrail_translation.handler import AnthropicM
|
|||
from litellm.llms.base_llm.guardrail_translation.utils import (
|
||||
effective_scan_only_tool_results_for_guardrail,
|
||||
)
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
|
||||
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, bedrock_bearer_token
|
||||
from litellm.llms.custom_httpx.http_handler import (
|
||||
get_async_httpx_client,
|
||||
httpxSpecialProvider,
|
||||
|
|
@ -56,7 +56,6 @@ from litellm.proxy.guardrails.anthropic_sse import (
|
|||
is_raw_sse_stream,
|
||||
model_response_text,
|
||||
)
|
||||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.guardrails import (
|
||||
BedrockChecksConfigModel,
|
||||
BedrockGuardrailStreamingParams,
|
||||
|
|
@ -713,9 +712,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# logic becomes shared across providers.
|
||||
|
||||
#### CALL HOOKS - proxy only ####
|
||||
def _load_credentials(
|
||||
self,
|
||||
):
|
||||
def _load_credentials(self, bearer_token: str | None = None):
|
||||
try:
|
||||
from botocore.credentials import Credentials
|
||||
except ImportError:
|
||||
|
|
@ -737,17 +734,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
aws_region_name=aws_region_name,
|
||||
)
|
||||
|
||||
credentials: Final[Credentials] = self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
credentials: Final[Credentials | None] = (
|
||||
None
|
||||
if bearer_token is not None
|
||||
else self.get_credentials(
|
||||
aws_access_key_id=aws_access_key_id,
|
||||
aws_secret_access_key=aws_secret_access_key,
|
||||
aws_session_token=aws_session_token,
|
||||
aws_region_name=aws_region_name,
|
||||
aws_session_name=aws_session_name,
|
||||
aws_profile_name=aws_profile_name,
|
||||
aws_role_name=aws_role_name,
|
||||
aws_web_identity_token=aws_web_identity_token,
|
||||
aws_sts_endpoint=aws_sts_endpoint,
|
||||
aws_external_id=aws_external_id,
|
||||
)
|
||||
)
|
||||
return credentials, aws_region_name
|
||||
|
||||
|
|
@ -779,13 +780,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
proxy_endpoint_url = f"{proxy_endpoint_url}{request_path}"
|
||||
encoded_data: Final = json.dumps(data).encode("utf-8")
|
||||
|
||||
# first check api-key, if none, fall back to sigV4
|
||||
if api_key is not None:
|
||||
aws_bearer_token: str | None = api_key
|
||||
else:
|
||||
aws_bearer_token = get_secret_str("AWS_BEARER_TOKEN_BEDROCK")
|
||||
aws_bearer_token: Final = bedrock_bearer_token(api_key)
|
||||
|
||||
if aws_bearer_token:
|
||||
if aws_bearer_token is not None:
|
||||
try:
|
||||
from botocore.awsrequest import AWSRequest
|
||||
except ImportError:
|
||||
|
|
@ -916,7 +913,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
source,
|
||||
)
|
||||
return BedrockGuardrailResponse()
|
||||
credentials, aws_region_name = self._load_credentials()
|
||||
credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
|
||||
allow_chunking: Final = not self._content_uses_contextual_grounding(content)
|
||||
|
||||
completed_chunk_usages: Final[list[BedrockGuardrailUsage]] = [] # mutable-ok: billed-chunk usage accumulator
|
||||
|
|
@ -958,7 +955,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self,
|
||||
content: Sequence[BedrockContentItem],
|
||||
base_request_data: Mapping[str, object],
|
||||
credentials: "Credentials",
|
||||
credentials: "Credentials | None",
|
||||
aws_region_name: str,
|
||||
api_key: str | None,
|
||||
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
|
||||
|
|
@ -1096,7 +1093,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self,
|
||||
content: Sequence[BedrockContentItem],
|
||||
base_request_data: Mapping[str, object],
|
||||
credentials: "Credentials",
|
||||
credentials: "Credentials | None",
|
||||
aws_region_name: str,
|
||||
api_key: str | None,
|
||||
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
|
||||
|
|
@ -1146,7 +1143,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
self,
|
||||
content: Sequence[BedrockContentItem],
|
||||
base_request_data: Mapping[str, object],
|
||||
credentials: "Credentials",
|
||||
credentials: "Credentials | None",
|
||||
aws_region_name: str,
|
||||
api_key: str | None,
|
||||
request_data: dict | None, # mutable-ok: proxy request body dict, mutated by the logging helper
|
||||
|
|
@ -1873,9 +1870,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
# Nothing to scan (e.g. tool-only turn) -> allow, like ApplyGuardrail does.
|
||||
return BedrockGuardrailResponse()
|
||||
|
||||
credentials, aws_region_name = self._load_credentials()
|
||||
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
|
||||
api_key: Final[str | None] = request_data.get("api_key") if request_data else None
|
||||
credentials, aws_region_name = self._load_credentials(bearer_token=bedrock_bearer_token(api_key))
|
||||
body: Final[dict[str, object]] = {"messages": checks_messages, "checks": self.checks}
|
||||
|
||||
prepared_request: Final = self._prepare_request(
|
||||
credentials=credentials,
|
||||
|
|
|
|||
|
|
@ -535,6 +535,7 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
try:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.exceptions import NoCredentialsError
|
||||
except ImportError:
|
||||
raise ImportError("Missing boto3 to call bedrock. Run 'pip install boto3'.")
|
||||
optional_params = optional_params or {}
|
||||
|
|
@ -582,10 +583,14 @@ class AWSSecretsManagerV2(BaseAWSLLM, BaseSecretManager):
|
|||
"X-Amz-Target": f"secretsmanager.{action}",
|
||||
}
|
||||
|
||||
credentials: Final = boto3_credentials_info.credentials
|
||||
if credentials is None:
|
||||
raise NoCredentialsError()
|
||||
|
||||
# Sign request
|
||||
request: Final = AWSRequest(method="POST", url=endpoint_url, data=body, headers=headers)
|
||||
SigV4Auth(
|
||||
boto3_credentials_info.credentials,
|
||||
credentials,
|
||||
"secretsmanager",
|
||||
boto3_credentials_info.aws_region_name,
|
||||
).add_auth(request)
|
||||
|
|
|
|||
|
|
@ -513,3 +513,27 @@ def test_the_rust_opt_in_needs_no_sigv4_principal():
|
|||
assert not {"aws_access_key_id", "aws_secret_access_key", "aws_session_token"} & params.keys()
|
||||
assert params["aws_region_name"] == "us-east-1"
|
||||
assert seen["call"][0]["api_key"] == "bedrock-bearer-token"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("configured_through", ["env_var", "api_key"])
|
||||
def test_bearer_token_auth_never_runs_the_sigv4_credential_chain(monkeypatch, configured_through):
|
||||
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
|
||||
raises; a bearer-token deployment must still serve the request, since the
|
||||
bearer token alone signs it."""
|
||||
if configured_through == "env_var":
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "bedrock-bearer-token")
|
||||
else:
|
||||
monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False)
|
||||
client = _sync_client_returning_converse_response()
|
||||
|
||||
response = BedrockConverseLLM().completion(
|
||||
**_completion_kwargs(
|
||||
optional_params={"maxTokens": 16, "aws_profile_name": "litellm-no-such-aws-profile"},
|
||||
litellm_params={},
|
||||
client=client,
|
||||
api_key="bedrock-bearer-token" if configured_through == "api_key" else None,
|
||||
)
|
||||
)
|
||||
|
||||
assert response.choices[0].message.content == "hi"
|
||||
assert client.post.call_args.kwargs["headers"]["Authorization"] == "Bearer bedrock-bearer-token"
|
||||
|
|
|
|||
|
|
@ -1033,3 +1033,29 @@ def test_load_credentials_assumes_role_with_external_id(monkeypatch):
|
|||
assert credentials.token == "assumed-session-token"
|
||||
assert aws_region_name == "us-east-1"
|
||||
assert "aws_external_id" not in optional_params
|
||||
|
||||
|
||||
def test_bedrock_embedding_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
||||
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
|
||||
raises; a bearer-token deployment must still serve the request, since the
|
||||
bearer token alone signs it."""
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
|
||||
client = HTTPHandler()
|
||||
|
||||
with patch.object(client, "post") as mock_post:
|
||||
mock_response = Mock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.text = json.dumps(titan_embedding_response)
|
||||
mock_response.json = lambda: json.loads(mock_response.text)
|
||||
mock_post.return_value = mock_response
|
||||
|
||||
response = litellm.embedding(
|
||||
model="bedrock/amazon.titan-embed-text-v1",
|
||||
input=test_input,
|
||||
client=client,
|
||||
aws_region_name="us-west-2",
|
||||
aws_profile_name="litellm-no-such-aws-profile",
|
||||
)
|
||||
|
||||
assert response.data[0]["embedding"] == titan_embedding_response["embedding"]
|
||||
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
|
|
|||
|
|
@ -135,3 +135,24 @@ class TestBedrockImageGeneration:
|
|||
assert response is not None
|
||||
assert len(response.data) > 0
|
||||
mock_bedrock_image_gen.assert_called_once()
|
||||
|
||||
|
||||
def test_image_generation_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
||||
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
|
||||
raises; a bearer-token deployment must still sign the request with the
|
||||
bearer token alone."""
|
||||
from litellm.llms.bedrock.image_generation.image_handler import BedrockImageGeneration
|
||||
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
|
||||
|
||||
request = BedrockImageGeneration()._prepare_request(
|
||||
model="amazon.nova-canvas-v1:0",
|
||||
prompt="A cute baby sea otter",
|
||||
optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"},
|
||||
api_base=None,
|
||||
extra_headers=None,
|
||||
api_key=None,
|
||||
logging_obj=Mock(),
|
||||
)
|
||||
|
||||
assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@
|
|||
import base64
|
||||
import io
|
||||
from typing import cast
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
|
@ -655,3 +656,23 @@ def test_transform_response_empty_images_without_error_raises():
|
|||
raw_response=resp,
|
||||
logging_obj=None, # type: ignore[arg-type]
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_request_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
||||
"""The deployment's AWS profile does not exist, so resolving SigV4 credentials
|
||||
raises; a bearer-token deployment must still sign the request with the
|
||||
bearer token alone."""
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
|
||||
|
||||
request = BedrockImageEdit()._prepare_request(
|
||||
model="amazon.nova-canvas-v1:0",
|
||||
image=[io.BytesIO(b"fake-png")],
|
||||
prompt="make it warmer",
|
||||
optional_params={"aws_region_name": "us-west-2", "aws_profile_name": "litellm-no-such-aws-profile"},
|
||||
api_base=None,
|
||||
extra_headers=None,
|
||||
logging_obj=Mock(),
|
||||
api_key=None,
|
||||
)
|
||||
|
||||
assert request.prepped.headers["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
|
|
|||
|
|
@ -5792,3 +5792,25 @@ async def test_apply_guardrail_debug_log_masks_signed_request_headers():
|
|||
assert header_lines, "expected the signed-request debug line to be logged"
|
||||
assert any("X-Amz-Security-Token" in message for message in header_lines)
|
||||
assert all(session_token not in message for message in rendered_messages)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
||||
"""The guardrail's AWS profile does not exist, so resolving SigV4 credentials
|
||||
raises; with a bearer token configured the guardrail must still run, since
|
||||
the bearer token alone signs the request."""
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
|
||||
guardrail = BedrockGuardrail(
|
||||
guardrailIdentifier="test-guardrail",
|
||||
guardrailVersion="DRAFT",
|
||||
aws_profile_name="litellm-no-such-aws-profile",
|
||||
)
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.json.return_value = {"action": "NONE", "assessments": []}
|
||||
|
||||
with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response) as mock_post:
|
||||
response = await guardrail.make_bedrock_api_request(source="INPUT", messages=[{"role": "user", "content": "hello"}])
|
||||
|
||||
assert response["action"] == "NONE"
|
||||
assert mock_post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
|
|
|||
|
|
@ -833,3 +833,31 @@ async def test_many_blocks_scanned_at_request_level_and_can_block():
|
|||
sent_texts = [c["text"] for m in body_messages for c in m["content"]]
|
||||
assert sent_texts == [f"b{i}" for i in range(25)]
|
||||
assert all(len(m["content"]) <= 10 for m in body_messages)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checks_bearer_token_never_runs_the_sigv4_credential_chain(monkeypatch):
|
||||
"""Same bearer-token rule as ApplyGuardrail: the guardrail's AWS profile does
|
||||
not exist, yet the InvokeGuardrailChecks call still goes out on the bearer
|
||||
token and its verdict is enforced."""
|
||||
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-12345")
|
||||
g = BedrockGuardrail(
|
||||
checks=CONTENT_FILTER_CHECKS,
|
||||
content_filter_threshold=0.5,
|
||||
aws_profile_name="litellm-no-such-aws-profile",
|
||||
)
|
||||
payload = {"results": {"contentFilter": {"results": [{"category": "VIOLENCE", "severityScore": 0.8}]}}}
|
||||
post = AsyncMock(return_value=_mock_http_response(200, payload))
|
||||
|
||||
with patch.object(g.async_handler, "post", new=post):
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await g.make_bedrock_api_request(
|
||||
source="INPUT",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
request_data={"messages": []},
|
||||
)
|
||||
|
||||
assert exc.value.detail["bedrock_guardrail_checks"] == [
|
||||
{"check": "contentFilter", "category": "VIOLENCE", "severityScore": 0.8}
|
||||
]
|
||||
assert post.call_args.kwargs["headers"]["Authorization"] == "Bearer env-bearer-token-12345"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue