Support aws_session_token for bedrock client. https://github.com/BerriAI/litellm/issues/4346

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,

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
+            if ("aws_bedrock_client" in optional_params):  
-            ):  # use old bedrock flow for aws_bedrock_client users.
+                # Extract credentials for legacy boto3 client and pass thru to httpx
-                response = bedrock.completion(
+                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(