mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
Improve AWS credential loading and configuration handling in BedrockProvider (#5699)
Signed-off-by: fbalicchia <fbalicchia@gmail.com>
This commit is contained in:
committed by
GitHub
parent
9c52611946
commit
f6a2e2ec83
@@ -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),
|
||||
],
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user