mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 11:43:54 +00:00
Support aws_session_token for bedrock client. https://github.com/BerriAI/litellm/issues/4346
This commit is contained in:
Brian Schultheiss
parent
cf52d3fa00
commit
d6925b40fe
2 changed files with 68 additions and 60 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -2176,45 +2176,18 @@ 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(
|
||||
model=model,
|
||||
messages=messages,
|
||||
custom_prompt_dict=litellm.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,
|
||||
)
|
||||
|
||||
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 ("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,
|
||||
|
|
@ -2249,6 +2222,7 @@ def completion(
|
|||
acompletion=acompletion,
|
||||
client=client,
|
||||
)
|
||||
|
||||
if optional_params.get("stream", False):
|
||||
## LOGGING
|
||||
logging.post_call(
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue