diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 1f00bf7792e..c3da992a904 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -7,7 +7,7 @@ import urllib.parse from collections.abc import Callable from datetime import datetime from threading import Lock -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, cast, get_args, overload import httpx from pydantic import BaseModel, ValidationError @@ -48,12 +48,19 @@ _STS_REGION_FROM_ENDPOINT_PATTERN: Final = re.compile( SIGV4_COMPUTED_HEADERS: Final = frozenset({"authorization", "x-amz-date", "x-amz-security-token", "date"}) -class Boto3CredentialsInfo(BaseModel): - credentials: Credentials | None +class BedrockRequestTarget(BaseModel): aws_region_name: str aws_bedrock_runtime_endpoint: str | None +class Boto3CredentialsInfo(BedrockRequestTarget): + credentials: Credentials + + +class BearerRequestTarget(BedrockRequestTarget): + credentials: None = 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 @@ -1392,9 +1399,26 @@ class BaseAWSLLM: else: return f"https://bedrock-runtime.{aws_region_name}.{dns_suffix}" + @overload + def _get_boto_credentials_from_optional_params( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + model: str | None = None, + bearer_token: None = None, + ) -> Boto3CredentialsInfo: ... + + @overload + def _get_boto_credentials_from_optional_params( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + model: str | None = None, + *, + bearer_token: str, + ) -> BearerRequestTarget: ... + def _get_boto_credentials_from_optional_params( self, optional_params: dict, model: str | None = None, bearer_token: str | None = None - ) -> Boto3CredentialsInfo: + ) -> Boto3CredentialsInfo | BearerRequestTarget: """ Get boto3 credentials from optional params @@ -1425,23 +1449,24 @@ class BaseAWSLLM: ) # https://bedrock-runtime.{region_name}.amazonaws.com aws_external_id: Final = optional_params.pop("aws_external_id", None) - 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, + if bearer_token is not None: + return BearerRequestTarget( 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, + aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, ) - ) + 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, + ) return Boto3CredentialsInfo( credentials=credentials, aws_region_name=aws_region_name, diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index a7b74f3752a..5fb86d476f4 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,7 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable -from typing import TYPE_CHECKING, Final, get_args +from typing import TYPE_CHECKING, Final, get_args, overload import httpx @@ -42,6 +42,20 @@ if TYPE_CHECKING: class BedrockEmbedding(BaseAWSLLM): + @overload + def _load_credentials( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + bearer_token: None = None, + ) -> tuple[Credentials, str]: ... + + @overload + def _load_credentials( + self, + optional_params: dict, # mutable-ok: the implementation pops the aws_* keys out of the caller's dict in place + bearer_token: str, + ) -> tuple[None, str]: ... + def _load_credentials( self, optional_params: dict, @@ -598,11 +612,8 @@ 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( diff --git a/litellm/secret_managers/aws_secret_manager_v2.py b/litellm/secret_managers/aws_secret_manager_v2.py index acdb83094e6..2c7f1f8389d 100644 --- a/litellm/secret_managers/aws_secret_manager_v2.py +++ b/litellm/secret_managers/aws_secret_manager_v2.py @@ -535,7 +535,6 @@ 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 {} @@ -583,14 +582,10 @@ 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( - credentials, + boto3_credentials_info.credentials, "secretsmanager", boto3_credentials_info.aws_region_name, ).add_auth(request) diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 9511732fd50..a222e22f6d0 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -771,7 +771,7 @@ async def test_bedrock_guardrail_prepare_request_with_api_key(): @pytest.mark.asyncio -async def test_bedrock_guardrail_prepare_request_without_api_key(): +async def test_bedrock_guardrail_prepare_request_without_api_key(monkeypatch): """Test _prepare_request method falls back to SigV4 when no api_key is provided""" from unittest.mock import Mock, patch @@ -789,18 +789,13 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): # Test data without api_key test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + monkeypatch.delenv("AWS_BEARER_TOKEN_BEDROCK", raising=False) with ( - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, patch("botocore.auth.SigV4Auth") as mock_sigv4_auth, patch("botocore.awsrequest.AWSRequest") as mock_aws_request, ): - # Mock no AWS_BEARER_TOKEN_BEDROCK - mock_get_secret.return_value = None - # Mock SigV4Auth mock_sigv4_instance = Mock() mock_sigv4_auth.return_value = mock_sigv4_instance @@ -826,7 +821,7 @@ async def test_bedrock_guardrail_prepare_request_without_api_key(): @pytest.mark.asyncio -async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): +async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(monkeypatch): """Test _prepare_request method uses Bearer token from environment when available""" from unittest.mock import Mock, patch @@ -844,15 +839,9 @@ async def test_bedrock_guardrail_prepare_request_with_bearer_token_env(): # Test data without api_key test_data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token-456") - with ( - patch( - "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.get_secret_str" - ) as mock_get_secret, - patch("botocore.awsrequest.AWSRequest") as mock_aws_request, - ): - - mock_get_secret.return_value = "env-bearer-token-456" + with patch("botocore.awsrequest.AWSRequest") as mock_aws_request: mock_request_instance = Mock() mock_request_instance.prepare.return_value = Mock() mock_aws_request.return_value = mock_request_instance