aisuite icon indicating copy to clipboard operation
aisuite copied to clipboard

Add support for SSO-based AWS sessions for Bedrock runtime client

Open Shaked35 opened this issue 2 months ago • 0 comments

We are currently unable to use this library when connecting to AWS Bedrock due to our SSO-based setup. The boto3.client() initialization does not allow for specifying a profile from the AWS config, preventing the use of SSO credentials. By initializing the boto3.Session with a specified profile (if provided), we can ensure the correct credentials are used.

Proposed Change: Replace the direct boto3.client() call with a session-based approach that uses the AWS config file profile if available. For example:

Existing code:

self.client = boto3.client("bedrock-runtime", region_name=self.region_name)

Proposed change:

if "profile_name" in config:
    session = boto3.Session(profile_name=config.get("profile_name"))
else:
    session = boto3.Session()

self.client = session.client("bedrock-runtime", region_name=self.region_name)

Why This Change Is Needed:

Allows the use of SSO-based credentials configured via ~/.aws/config. Ensures consistency with other AWS service integrations that rely on profile-based sessions. Improves flexibility for developers working in multi-account or SSO-enabled AWS environments. Additional Context: If running in an environment where a profile_name is not specified, the default session behavior remains unchanged, ensuring backward compatibility. If the profile_name is set, the credentials will be properly sourced from the associated profile, enabling SSO-based authentication.

Shaked35 avatar Dec 19 '24 06:12 Shaked35