From 7f91e5354886fa1a445d953e2a368d507c3c0c3f Mon Sep 17 00:00:00 2001 From: Brian Schultheiss Date: Sun, 23 Jun 2024 13:15:04 -0700 Subject: [PATCH] updated documentation to reference boto3.client credential extraction, and update boto3.client creation to support session_token. --- docs/my-website/docs/providers/bedrock.md | 5 +++++ litellm/llms/bedrock.py | 25 ++++++++++++++++++++++- 2 files changed, 29 insertions(+), 1 deletion(-) diff --git a/docs/my-website/docs/providers/bedrock.md b/docs/my-website/docs/providers/bedrock.md index f380a6a50e3..adbc54caa7d 100644 --- a/docs/my-website/docs/providers/bedrock.md +++ b/docs/my-website/docs/providers/bedrock.md @@ -476,6 +476,7 @@ response = completion( messages=[{ "content": "Hello, how are you?","role": "user"}], aws_access_key_id="", aws_secret_access_key="", + aws_session_token="", aws_region_name="", ) ``` @@ -549,6 +550,10 @@ response = completion( This is a deprecated flow. Boto3 is not async. And boto3.client does not let us make the http call through httpx. Pass in your aws params through the method above 👆. [See Auth Code](https://github.com/BerriAI/litellm/blob/55a20c7cce99a93d36a82bf3ae90ba3baf9a7f89/litellm/llms/bedrock_httpx.py#L284) [Add new auth flow](https://github.com/BerriAI/litellm/issues) + +Experimental - 2024-Jun-23: + aws_access_key_id, aws_secret_access_key=, and aws_session_token will be extracted from boto3.client and be passed onto the httpx client + ::: Pass an external BedrockRuntime.Client object as a parameter to litellm.completion. Useful when using an AWS credentials profile, SSO session, assumed role session, or if environment variables are not available for auth. diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py index d0d3bef6da2..6c941bb558d 100644 --- a/litellm/llms/bedrock.py +++ b/litellm/llms/bedrock.py @@ -558,6 +558,7 @@ def init_bedrock_client( region_name=None, 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_bedrock_runtime_endpoint: Optional[str] = None, aws_session_name: Optional[str] = None, @@ -591,6 +592,7 @@ def init_bedrock_client( ( aws_access_key_id, aws_secret_access_key, + aws_session_token, aws_region_name, aws_bedrock_runtime_endpoint, aws_session_name, @@ -668,6 +670,21 @@ def init_bedrock_client( endpoint_url=endpoint_url, config=config, ) + 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 ### + client = boto3.client( + service_name="bedrock-runtime", + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + region_name=region_name, + endpoint_url=endpoint_url, + config=config, + ) + elif aws_role_name is not None and aws_session_name is not None: # use sts if role name passed in sts_client = boto3.client( @@ -786,9 +803,10 @@ def completion( _is_function_call = False json_schemas: dict = {} try: - # pop aws_secret_access_key, aws_access_key_id, aws_region_name from kwargs, since completion calls fail with them + # pop aws_secret_access_key, aws_access_key_id, aws_session_token, 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) @@ -806,6 +824,7 @@ def completion( client = init_bedrock_client( 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_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_role_name=aws_role_name, @@ -1328,6 +1347,7 @@ def embedding( # 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) @@ -1340,6 +1360,7 @@ def embedding( client = init_bedrock_client( 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_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_web_identity_token=aws_web_identity_token, @@ -1419,6 +1440,7 @@ def image_generation( # 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) @@ -1431,6 +1453,7 @@ def image_generation( client = init_bedrock_client( 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_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_web_identity_token=aws_web_identity_token,