aisuite
aisuite copied to clipboard
Add support for SSO-based AWS sessions for Bedrock runtime client
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.