diff --git a/litellm/llms/base_aws_llm.py b/litellm/llms/base_aws_llm.py new file mode 100644 index 0000000000..6b298d5be1 --- /dev/null +++ b/litellm/llms/base_aws_llm.py @@ -0,0 +1,196 @@ +import json +from typing import List, Optional + +import httpx + +from litellm._logging import verbose_logger +from litellm.caching import DualCache, InMemoryCache +from litellm.utils import get_secret + +from .base import BaseLLM + + +class AwsAuthError(Exception): + def __init__(self, status_code, message): + self.status_code = status_code + self.message = message + self.request = httpx.Request( + method="POST", url="https://us-west-2.console.aws.amazon.com/bedrock" + ) + self.response = httpx.Response(status_code=status_code, request=self.request) + super().__init__( + self.message + ) # Call the base class constructor with the parameters it needs + + +class BaseAWSLLM(BaseLLM): + def __init__(self) -> None: + self.iam_cache = DualCache() + super().__init__() + + def get_credentials( + self, + aws_access_key_id: Optional[str] = None, + aws_secret_access_key: Optional[str] = None, + aws_session_token: Optional[str] = None, + aws_region_name: Optional[str] = None, + aws_session_name: Optional[str] = None, + aws_profile_name: Optional[str] = None, + aws_role_name: Optional[str] = None, + aws_web_identity_token: Optional[str] = None, + aws_sts_endpoint: Optional[str] = None, + ): + """ + Return a boto3.Credentials object + """ + import boto3 + + ## CHECK IS 'os.environ/' passed in + params_to_check: List[Optional[str]] = [ + aws_access_key_id, + aws_secret_access_key, + aws_session_token, + aws_region_name, + aws_session_name, + aws_profile_name, + aws_role_name, + aws_web_identity_token, + aws_sts_endpoint, + ] + + # Iterate over parameters and update if needed + for i, param in enumerate(params_to_check): + if param and param.startswith("os.environ/"): + _v = get_secret(param) + if _v is not None and isinstance(_v, str): + params_to_check[i] = _v + # Assign updated values back to parameters + ( + aws_access_key_id, + aws_secret_access_key, + aws_session_token, + aws_region_name, + aws_session_name, + aws_profile_name, + aws_role_name, + aws_web_identity_token, + aws_sts_endpoint, + ) = params_to_check + + ### CHECK STS ### + if ( + aws_web_identity_token is not None + and aws_role_name is not None + and aws_session_name is not None + ): + verbose_logger.debug( + f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" + ) + + if aws_sts_endpoint is None: + sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com" + else: + sts_endpoint = aws_sts_endpoint + + iam_creds_cache_key = json.dumps( + { + "aws_web_identity_token": aws_web_identity_token, + "aws_role_name": aws_role_name, + "aws_session_name": aws_session_name, + "aws_region_name": aws_region_name, + "aws_sts_endpoint": sts_endpoint, + } + ) + + iam_creds_dict = self.iam_cache.get_cache(iam_creds_cache_key) + if iam_creds_dict is None: + oidc_token = get_secret(aws_web_identity_token) + + if oidc_token is None: + raise AwsAuthError( + message="OIDC token could not be retrieved from secret manager.", + status_code=401, + ) + + sts_client = boto3.client( + "sts", + region_name=aws_region_name, + endpoint_url=sts_endpoint, + ) + + # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html + sts_response = sts_client.assume_role_with_web_identity( + RoleArn=aws_role_name, + RoleSessionName=aws_session_name, + WebIdentityToken=oidc_token, + DurationSeconds=3600, + ) + + iam_creds_dict = { + "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], + "aws_secret_access_key": sts_response["Credentials"][ + "SecretAccessKey" + ], + "aws_session_token": sts_response["Credentials"]["SessionToken"], + "region_name": aws_region_name, + } + + self.iam_cache.set_cache( + key=iam_creds_cache_key, + value=json.dumps(iam_creds_dict), + ttl=3600 - 60, + ) + + session = boto3.Session(**iam_creds_dict) + + iam_creds = session.get_credentials() + + return iam_creds + elif aws_role_name is not None and aws_session_name is not None: + sts_client = boto3.client( + "sts", + aws_access_key_id=aws_access_key_id, # [OPTIONAL] + aws_secret_access_key=aws_secret_access_key, # [OPTIONAL] + ) + + sts_response = sts_client.assume_role( + RoleArn=aws_role_name, RoleSessionName=aws_session_name + ) + + # Extract the credentials from the response and convert to Session Credentials + sts_credentials = sts_response["Credentials"] + from botocore.credentials import Credentials + + credentials = Credentials( + access_key=sts_credentials["AccessKeyId"], + secret_key=sts_credentials["SecretAccessKey"], + token=sts_credentials["SessionToken"], + ) + return credentials + elif aws_profile_name is not None: ### CHECK SESSION ### + # uses auth values from AWS profile usually stored in ~/.aws/credentials + client = boto3.Session(profile_name=aws_profile_name) + + return client.get_credentials() + elif ( + aws_access_key_id is not None + and aws_secret_access_key is not None + and aws_session_token is not None + ): ### CHECK FOR AWS SESSION TOKEN ### + from botocore.credentials import Credentials + + credentials = Credentials( + access_key=aws_access_key_id, + secret_key=aws_secret_access_key, + token=aws_session_token, + ) + return credentials + else: + session = boto3.Session( + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + region_name=aws_region_name, + ) + + return session.get_credentials() diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index c433c32b7d..73387212ff 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -57,6 +57,7 @@ from litellm.utils import ( ) from .base import BaseLLM +from .base_aws_llm import BaseAWSLLM from .bedrock import BedrockError, ModelResponseIterator, convert_messages_to_prompt from .prompt_templates.factory import ( _bedrock_converse_messages_pt, @@ -87,7 +88,6 @@ BEDROCK_CONVERSE_MODELS = [ ] -iam_cache = DualCache() _response_stream_shape_cache = None bedrock_tool_name_mappings: InMemoryCache = InMemoryCache( max_size_in_memory=50, default_ttl=600 @@ -312,7 +312,7 @@ def make_sync_call( return completion_stream -class BedrockLLM(BaseLLM): +class BedrockLLM(BaseAWSLLM): """ Example call @@ -380,183 +380,6 @@ class BedrockLLM(BaseLLM): prompt += f"{message['content']}" return prompt, chat_history # type: ignore - def get_credentials( - self, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, - aws_region_name: Optional[str] = None, - aws_session_name: Optional[str] = None, - aws_profile_name: Optional[str] = None, - aws_role_name: Optional[str] = None, - aws_web_identity_token: Optional[str] = None, - aws_sts_endpoint: Optional[str] = None, - ): - """ - Return a boto3.Credentials object - """ - import boto3 - - print_verbose( - f"Boto3 get_credentials called variables passed to function {locals()}" - ) - - ## CHECK IS 'os.environ/' passed in - params_to_check: List[Optional[str]] = [ - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region_name, - aws_session_name, - aws_profile_name, - aws_role_name, - aws_web_identity_token, - aws_sts_endpoint, - ] - - # Iterate over parameters and update if needed - for i, param in enumerate(params_to_check): - if param and param.startswith("os.environ/"): - _v = get_secret(param) - if _v is not None and isinstance(_v, str): - params_to_check[i] = _v - # Assign updated values back to parameters - ( - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region_name, - aws_session_name, - aws_profile_name, - aws_role_name, - aws_web_identity_token, - aws_sts_endpoint, - ) = params_to_check - - ### CHECK STS ### - if ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ): - print_verbose( - f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" - ) - - if aws_sts_endpoint is None: - sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com" - else: - sts_endpoint = aws_sts_endpoint - - iam_creds_cache_key = json.dumps( - { - "aws_web_identity_token": aws_web_identity_token, - "aws_role_name": aws_role_name, - "aws_session_name": aws_session_name, - "aws_region_name": aws_region_name, - "aws_sts_endpoint": sts_endpoint, - } - ) - - iam_creds_dict = iam_cache.get_cache(iam_creds_cache_key) - if iam_creds_dict is None: - oidc_token = get_secret(aws_web_identity_token) - - if oidc_token is None: - raise BedrockError( - message="OIDC token could not be retrieved from secret manager.", - status_code=401, - ) - - sts_client = boto3.client( - "sts", - region_name=aws_region_name, - endpoint_url=sts_endpoint, - ) - - # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html - # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html - sts_response = sts_client.assume_role_with_web_identity( - RoleArn=aws_role_name, - RoleSessionName=aws_session_name, - WebIdentityToken=oidc_token, - DurationSeconds=3600, - ) - - iam_creds_dict = { - "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], - "aws_secret_access_key": sts_response["Credentials"][ - "SecretAccessKey" - ], - "aws_session_token": sts_response["Credentials"]["SessionToken"], - "region_name": aws_region_name, - } - - iam_cache.set_cache( - key=iam_creds_cache_key, - value=json.dumps(iam_creds_dict), - ttl=3600 - 60, - ) - - session = boto3.Session(**iam_creds_dict) - - iam_creds = session.get_credentials() - - return iam_creds - elif aws_role_name is not None and aws_session_name is not None: - print_verbose( - f"Using STS Client AWS aws_role_name: {aws_role_name} aws_session_name: {aws_session_name}" - ) - sts_client = boto3.client( - "sts", - aws_access_key_id=aws_access_key_id, # [OPTIONAL] - aws_secret_access_key=aws_secret_access_key, # [OPTIONAL] - ) - - sts_response = sts_client.assume_role( - RoleArn=aws_role_name, RoleSessionName=aws_session_name - ) - - # Extract the credentials from the response and convert to Session Credentials - sts_credentials = sts_response["Credentials"] - from botocore.credentials import Credentials - - credentials = Credentials( - access_key=sts_credentials["AccessKeyId"], - secret_key=sts_credentials["SecretAccessKey"], - token=sts_credentials["SessionToken"], - ) - return credentials - elif aws_profile_name is not None: ### CHECK SESSION ### - # uses auth values from AWS profile usually stored in ~/.aws/credentials - print_verbose(f"Using AWS profile: {aws_profile_name}") - client = boto3.Session(profile_name=aws_profile_name) - - return client.get_credentials() - elif ( - aws_access_key_id is not None - and aws_secret_access_key is not None - and aws_session_token is not None - ): ### CHECK FOR AWS SESSION TOKEN ### - print_verbose(f"Using AWS Session Token: {aws_session_token}") - from botocore.credentials import Credentials - - credentials = Credentials( - access_key=aws_access_key_id, - secret_key=aws_secret_access_key, - token=aws_session_token, - ) - return credentials - else: - print_verbose("Using Default AWS Session") - session = boto3.Session( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - region_name=aws_region_name, - ) - - return session.get_credentials() - def process_response( self, model: str, @@ -1414,7 +1237,7 @@ class AmazonConverseConfig: return optional_params -class BedrockConverseLLM(BaseLLM): +class BedrockConverseLLM(BaseAWSLLM): def __init__(self) -> None: super().__init__() @@ -1554,173 +1377,6 @@ class BedrockConverseLLM(BaseLLM): """ return urllib.parse.quote(model_id, safe="") - def get_credentials( - self, - aws_access_key_id: Optional[str] = None, - aws_secret_access_key: Optional[str] = None, - aws_session_token: Optional[str] = None, - aws_region_name: Optional[str] = None, - aws_session_name: Optional[str] = None, - aws_profile_name: Optional[str] = None, - aws_role_name: Optional[str] = None, - aws_web_identity_token: Optional[str] = None, - aws_sts_endpoint: Optional[str] = None, - ): - """ - Return a boto3.Credentials object - """ - import boto3 - - ## CHECK IS 'os.environ/' passed in - params_to_check: List[Optional[str]] = [ - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region_name, - aws_session_name, - aws_profile_name, - aws_role_name, - aws_web_identity_token, - aws_sts_endpoint, - ] - - # Iterate over parameters and update if needed - for i, param in enumerate(params_to_check): - if param and param.startswith("os.environ/"): - _v = get_secret(param) - if _v is not None and isinstance(_v, str): - params_to_check[i] = _v - # Assign updated values back to parameters - ( - aws_access_key_id, - aws_secret_access_key, - aws_session_token, - aws_region_name, - aws_session_name, - aws_profile_name, - aws_role_name, - aws_web_identity_token, - aws_sts_endpoint, - ) = params_to_check - - ### CHECK STS ### - if ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ): - print_verbose( - f"IN Web Identity Token: {aws_web_identity_token} | Role Name: {aws_role_name} | Session Name: {aws_session_name}" - ) - - if aws_sts_endpoint is None: - sts_endpoint = f"https://sts.{aws_region_name}.amazonaws.com" - else: - sts_endpoint = aws_sts_endpoint - - iam_creds_cache_key = json.dumps( - { - "aws_web_identity_token": aws_web_identity_token, - "aws_role_name": aws_role_name, - "aws_session_name": aws_session_name, - "aws_region_name": aws_region_name, - "aws_sts_endpoint": sts_endpoint, - } - ) - - iam_creds_dict = iam_cache.get_cache(iam_creds_cache_key) - if iam_creds_dict is None: - oidc_token = get_secret(aws_web_identity_token) - - if oidc_token is None: - raise BedrockError( - message="OIDC token could not be retrieved from secret manager.", - status_code=401, - ) - - sts_client = boto3.client( - "sts", - region_name=aws_region_name, - endpoint_url=sts_endpoint, - ) - - # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html - # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html - sts_response = sts_client.assume_role_with_web_identity( - RoleArn=aws_role_name, - RoleSessionName=aws_session_name, - WebIdentityToken=oidc_token, - DurationSeconds=3600, - ) - - iam_creds_dict = { - "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], - "aws_secret_access_key": sts_response["Credentials"][ - "SecretAccessKey" - ], - "aws_session_token": sts_response["Credentials"]["SessionToken"], - "region_name": aws_region_name, - } - - iam_cache.set_cache( - key=iam_creds_cache_key, - value=json.dumps(iam_creds_dict), - ttl=3600 - 60, - ) - - session = boto3.Session(**iam_creds_dict) - - iam_creds = session.get_credentials() - - return iam_creds - elif aws_role_name is not None and aws_session_name is not None: - sts_client = boto3.client( - "sts", - aws_access_key_id=aws_access_key_id, # [OPTIONAL] - aws_secret_access_key=aws_secret_access_key, # [OPTIONAL] - ) - - sts_response = sts_client.assume_role( - RoleArn=aws_role_name, RoleSessionName=aws_session_name - ) - - # Extract the credentials from the response and convert to Session Credentials - sts_credentials = sts_response["Credentials"] - from botocore.credentials import Credentials - - credentials = Credentials( - access_key=sts_credentials["AccessKeyId"], - secret_key=sts_credentials["SecretAccessKey"], - token=sts_credentials["SessionToken"], - ) - return credentials - elif aws_profile_name is not None: ### CHECK SESSION ### - # uses auth values from AWS profile usually stored in ~/.aws/credentials - client = boto3.Session(profile_name=aws_profile_name) - - return client.get_credentials() - elif ( - aws_access_key_id is not None - and aws_secret_access_key is not None - and aws_session_token is not None - ): ### CHECK FOR AWS SESSION TOKEN ### - from botocore.credentials import Credentials - - credentials = Credentials( - access_key=aws_access_key_id, - secret_key=aws_secret_access_key, - token=aws_session_token, - ) - return credentials - else: - session = boto3.Session( - aws_access_key_id=aws_access_key_id, - aws_secret_access_key=aws_secret_access_key, - region_name=aws_region_name, - ) - - return session.get_credentials() - async def async_streaming( self, model: str,