Add support for AWS credentials from profile file

https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#aws-config-file
This commit is contained in:
David Leen 2024-02-08 15:06:47 -08:00
parent 59465bd612
commit 140d915adf
2 changed files with 30 additions and 3 deletions

View file

@ -197,7 +197,7 @@ response = completion(
### SSO Login (AWS Profile) ### SSO Login (AWS Profile)
- Set `AWS_PROFILE` environment variable - Set `AWS_PROFILE` environment variable
- Make bedrock completion call - Make bedrock completion call
```python ```python
import os import os
from litellm import completion from litellm import completion
@ -208,11 +208,24 @@ response = completion(
) )
``` ```
### STS based Auth or pass `aws_profile_name`:
```python
import os
from litellm import completion
response = completion(
model="bedrock/anthropic.claude-instant-v1",
messages=[{ "content": "Hello, how are you?","role": "user"}],
aws_profile_name="dev-profile",
)
```
### STS based Auth
- Set `aws_role_name` and `aws_session_name` in completion() / embedding() function - Set `aws_role_name` and `aws_session_name` in completion() / embedding() function
Make the bedrock completion call Make the bedrock completion call
```python ```python
from litellm import completion from litellm import completion

View file

@ -356,6 +356,7 @@ def init_bedrock_client(
aws_region_name: Optional[str] = None, aws_region_name: Optional[str] = None,
aws_bedrock_runtime_endpoint: Optional[str] = None, aws_bedrock_runtime_endpoint: Optional[str] = None,
aws_session_name: Optional[str] = None, aws_session_name: Optional[str] = None,
aws_profile_name: Optional[str] = None,
aws_role_name: Optional[str] = None, aws_role_name: Optional[str] = None,
timeout: Optional[int] = None, timeout: Optional[int] = None,
): ):
@ -371,6 +372,7 @@ def init_bedrock_client(
aws_region_name, aws_region_name,
aws_bedrock_runtime_endpoint, aws_bedrock_runtime_endpoint,
aws_session_name, aws_session_name,
aws_profile_name,
aws_role_name, aws_role_name,
] ]
@ -385,6 +387,7 @@ def init_bedrock_client(
aws_region_name, aws_region_name,
aws_bedrock_runtime_endpoint, aws_bedrock_runtime_endpoint,
aws_session_name, aws_session_name,
aws_profile_name,
aws_role_name, aws_role_name,
) = params_to_check ) = params_to_check
@ -450,6 +453,15 @@ def init_bedrock_client(
endpoint_url=endpoint_url, endpoint_url=endpoint_url,
config=config, config=config,
) )
elif aws_profile_name is not None:
# uses auth values from AWS profile usually stored in ~/.aws/credentials
client = boto3.Session(profile_name=aws_profile_name).client(
service_name="bedrock-runtime",
region_name=region_name,
endpoint_url=endpoint_url,
config=config,
)
else: else:
# aws_access_key_id is None, assume user is trying to auth using env variables # aws_access_key_id is None, assume user is trying to auth using env variables
# boto3 automatically reads env variables # boto3 automatically reads env variables
@ -523,6 +535,7 @@ def completion(
aws_region_name = optional_params.pop("aws_region_name", None) aws_region_name = optional_params.pop("aws_region_name", None)
aws_role_name = optional_params.pop("aws_role_name", None) aws_role_name = optional_params.pop("aws_role_name", None)
aws_session_name = optional_params.pop("aws_session_name", None) aws_session_name = optional_params.pop("aws_session_name", None)
aws_profile_name = optional_params.pop("aws_profile_name", None)
aws_bedrock_runtime_endpoint = optional_params.pop( aws_bedrock_runtime_endpoint = optional_params.pop(
"aws_bedrock_runtime_endpoint", None "aws_bedrock_runtime_endpoint", None
) )
@ -539,6 +552,7 @@ def completion(
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint, aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
aws_role_name=aws_role_name, aws_role_name=aws_role_name,
aws_session_name=aws_session_name, aws_session_name=aws_session_name,
aws_profile_name=aws_profile_name,
) )
model = model model = model