diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 84ab10907cc..3eb8acd3907 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -301,6 +301,7 @@ class BedrockLLM(BaseLLM): 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, @@ -316,6 +317,7 @@ class BedrockLLM(BaseLLM): 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, @@ -333,6 +335,7 @@ class BedrockLLM(BaseLLM): ( aws_access_key_id, aws_secret_access_key, + aws_session_token, aws_region_name, aws_session_name, aws_profile_name, @@ -426,6 +429,18 @@ class BedrockLLM(BaseLLM): 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, @@ -733,6 +748,7 @@ class BedrockLLM(BaseLLM): # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) aws_access_key_id = optional_params.pop("aws_access_key_id", None) + aws_session_token = optional_params.pop("aws_session_token", None) aws_region_name = optional_params.pop("aws_region_name", None) aws_role_name = optional_params.pop("aws_role_name", None) aws_session_name = optional_params.pop("aws_session_name", None) @@ -764,6 +780,7 @@ class BedrockLLM(BaseLLM): credentials: 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, @@ -1418,6 +1435,7 @@ class BedrockConverseLLM(BaseLLM): 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, @@ -1433,6 +1451,7 @@ class BedrockConverseLLM(BaseLLM): 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, @@ -1450,6 +1469,7 @@ class BedrockConverseLLM(BaseLLM): ( aws_access_key_id, aws_secret_access_key, + aws_session_token, aws_region_name, aws_session_name, aws_profile_name, @@ -1543,6 +1563,18 @@ class BedrockConverseLLM(BaseLLM): 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, @@ -1678,6 +1710,7 @@ class BedrockConverseLLM(BaseLLM): # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them aws_secret_access_key = optional_params.pop("aws_secret_access_key", None) aws_access_key_id = optional_params.pop("aws_access_key_id", None) + aws_session_token = optional_params.pop("aws_session_token", None) aws_region_name = optional_params.pop("aws_region_name", None) aws_role_name = optional_params.pop("aws_role_name", None) aws_session_name = optional_params.pop("aws_session_name", None) @@ -1709,6 +1742,7 @@ class BedrockConverseLLM(BaseLLM): credentials: 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, diff --git a/litellm/main.py b/litellm/main.py index 307659c8a2f..92cd2fee11a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2176,13 +2176,23 @@ def completion( # boto3 reads keys from .env custom_prompt_dict = custom_prompt_dict or litellm.custom_prompt_dict - if ( - "aws_bedrock_client" in optional_params - ): # use old bedrock flow for aws_bedrock_client users. - response = bedrock.completion( + + if ("aws_bedrock_client" in optional_params): + # Extract credentials for legacy boto3 client and pass thru to httpx + aws_bedrock_client = optional_params.pop("aws_bedrock_client") + creds = aws_bedrock_client._get_credentials().get_frozen_credentials() + if creds.access_key: + optional_params["aws_access_key_id"] = creds.access_key + if creds.secret_key: + optional_params["aws_secret_access_key"] = creds.secret_key + if creds.token: + optional_params["aws_session_token"] = creds.token + + if model.startswith("anthropic"): + response = bedrock_converse_chat_completion.completion( model=model, messages=messages, - custom_prompt_dict=litellm.custom_prompt_dict, + custom_prompt_dict=custom_prompt_dict, model_response=model_response, print_verbose=print_verbose, optional_params=optional_params, @@ -2192,63 +2202,27 @@ def completion( logging_obj=logging, extra_headers=extra_headers, timeout=timeout, + acompletion=acompletion, + client=client, + ) + else: + response = bedrock_chat_completion.completion( + model=model, + messages=messages, + custom_prompt_dict=custom_prompt_dict, + model_response=model_response, + print_verbose=print_verbose, + optional_params=optional_params, + litellm_params=litellm_params, + logger_fn=logger_fn, + encoding=encoding, + logging_obj=logging, + extra_headers=extra_headers, + timeout=timeout, + acompletion=acompletion, + client=client, ) - if ( - "stream" in optional_params - and optional_params["stream"] == True - and not isinstance(response, CustomStreamWrapper) - ): - # don't try to access stream object, - if "ai21" in model: - response = CustomStreamWrapper( - response, - model, - custom_llm_provider="bedrock", - logging_obj=logging, - ) - else: - response = CustomStreamWrapper( - iter(response), - model, - custom_llm_provider="bedrock", - logging_obj=logging, - ) - else: - if model.startswith("anthropic"): - response = bedrock_converse_chat_completion.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=encoding, - logging_obj=logging, - extra_headers=extra_headers, - timeout=timeout, - acompletion=acompletion, - client=client, - ) - else: - response = bedrock_chat_completion.completion( - model=model, - messages=messages, - custom_prompt_dict=custom_prompt_dict, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=encoding, - logging_obj=logging, - extra_headers=extra_headers, - timeout=timeout, - acompletion=acompletion, - client=client, - ) if optional_params.get("stream", False): ## LOGGING logging.post_call(