Improve AWS credential loading and configuration handling in BedrockProvider (#5699)

Signed-off-by: fbalicchia <fbalicchia@gmail.com>
This commit is contained in:
Filippo Balicchia
2025-11-26 00:56:50 +01:00
committed by GitHub
parent 9c52611946
commit f6a2e2ec83
+37 -11
View File
@@ -22,8 +22,9 @@ use super::formats::bedrock::{
pub const BEDROCK_DOC_LINK: &str =
"https://docs.aws.amazon.com/bedrock/latest/userguide/models-supported.html";
pub const BEDROCK_DEFAULT_MODEL: &str = "us.anthropic.claude-sonnet-4-20250514-v1:0";
pub const BEDROCK_DEFAULT_MODEL: &str = "us.anthropic.claude-sonnet-4-5-20250929-v1:0";
pub const BEDROCK_KNOWN_MODELS: &[&str] = &[
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
"us.anthropic.claude-sonnet-4-20250514-v1:0",
"us.anthropic.claude-3-7-sonnet-20250219-v1:0",
"us.anthropic.claude-opus-4-20250514-v1:0",
@@ -37,6 +38,8 @@ pub const BEDROCK_DEFAULT_MAX_RETRY_INTERVAL_MS: u64 = 120_000;
#[derive(Debug, serde::Serialize)]
pub struct BedrockProvider {
#[serde(skip)]
client: Client,
model: ModelConfig,
#[serde(skip)]
retry_config: RetryConfig,
@@ -49,7 +52,7 @@ impl BedrockProvider {
let config = crate::config::Config::global();
// Attempt to load config and secrets to get AWS_ prefixed keys
// to re-export them into the environment for aws_config::load_from_env()
// to re-export them into the environment for aws_config to use as fallback
let set_aws_env_vars = |res: Result<HashMap<String, Value>, _>| {
if let Ok(map) = res {
map.into_iter()
@@ -62,16 +65,38 @@ impl BedrockProvider {
set_aws_env_vars(config.all_values());
set_aws_env_vars(config.all_secrets());
let sdk_config = aws_config::load_from_env().await;
// Use load_defaults() which supports AWS SSO, profiles, and environment variables
let mut loader = aws_config::defaults(aws_config::BehaviorVersion::latest());
if let Ok(profile_name) = config.get_param::<String>("AWS_PROFILE") {
if !profile_name.is_empty() {
loader = loader.profile_name(&profile_name);
}
}
// Check for AWS_REGION configuration
if let Ok(region) = config.get_param::<String>("AWS_REGION") {
if !region.is_empty() {
loader = loader.region(aws_config::Region::new(region));
}
}
let sdk_config = loader.load().await;
// Validate credentials or return error back up
sdk_config
.credentials_provider()
.unwrap()
.ok_or_else(|| anyhow::anyhow!("No AWS credentials provider configured"))?
.provide_credentials()
.await?;
.await
.map_err(|e| anyhow::anyhow!("Failed to load AWS credentials: {}. Make sure to run 'aws sso login --profile <your-profile>' if using SSO", e))?;
let client = Client::new(&sdk_config);
let retry_config = Self::load_retry_config(config);
Ok(Self {
client,
model,
retry_config,
name: Self::metadata().name,
@@ -111,10 +136,8 @@ impl BedrockProvider {
) -> Result<(bedrock::Message, Option<bedrock::TokenUsage>), ProviderError> {
let model_name = &self.model.model_name;
let sdk_config = aws_config::load_from_env().await;
let client = Client::new(&sdk_config);
let mut request = client
let mut request = self
.client
.converse()
.system(bedrock::SystemContentBlock::Text(system.to_string()))
.model_id(model_name.to_string())
@@ -175,11 +198,14 @@ impl Provider for BedrockProvider {
ProviderMetadata::new(
"aws_bedrock",
"Amazon Bedrock",
"Run models through Amazon Bedrock. You may have to set 'AWS_' environment variables to configure authentication.",
"Run models through Amazon Bedrock. Supports AWS SSO profiles - run 'aws sso login --profile <profile-name>' before using. Configure with AWS_PROFILE and AWS_REGION, or use environment variables/credentials.",
BEDROCK_DEFAULT_MODEL,
BEDROCK_KNOWN_MODELS.to_vec(),
BEDROCK_DOC_LINK,
vec![ConfigKey::new("AWS_PROFILE", true, false, Some("default"))],
vec![
ConfigKey::new("AWS_PROFILE", true, false, Some("default")),
ConfigKey::new("AWS_REGION", true, false, None),
],
)
}