aljaz/goose
1
0
Fork 0
mirror of https://github.com/aljazceru/goose.git synced 2025-12-18 22:54:24 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix: improve AWS credential handling in Bedrock provider, make keyring optional (#1886)

Browse Source
This commit is contained in:
Kalvin C
2025-03-27 16:28:20 -07:00
committed by GitHub
parent 6bd4f36331
commit 2fa8c3e8ab
1 changed files with 27 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
crates/goose/src/providers/bedrock.rs
View File

@@ -1,8 +1,12 @@
use std::collections::HashMap;
use anyhow::Result; use anyhow::Result;
use async_trait::async_trait; use async_trait::async_trait;
use aws_sdk_bedrockruntime::config::ProvideCredentials;
use aws_sdk_bedrockruntime::operation::converse::ConverseError; use aws_sdk_bedrockruntime::operation::converse::ConverseError;
use aws_sdk_bedrockruntime::{types as bedrock, Client}; use aws_sdk_bedrockruntime::{types as bedrock, Client};
use mcp_core::Tool; use mcp_core::Tool;
use serde_json::Value;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage}; use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
use super::errors::ProviderError; use super::errors::ProviderError;
@@ -34,12 +38,30 @@ pub struct BedrockProvider {
impl BedrockProvider { impl BedrockProvider {
pub fn from_env(model: ModelConfig) -> Result<Self> { pub fn from_env(model: ModelConfig) -> Result<Self> {
let config = crate::config::Config::global(); let config = crate::config::Config::global();
for (key, value) in config.load_secrets()?.iter() {
if key.starts_with("AWS_") && value.is_string() { // Attempt to load config and secrets to get AWS_ prefixed keys
std::env::set_var(key, value.as_str().unwrap()); // to re-export them into the environment for aws_config::load_from_env()
let set_aws_env_vars = |res: Result<HashMap<String, Value>, _>| {
if let Ok(map) = res {
map.into_iter()
.filter(|(key, _)| key.starts_with("AWS_"))
.filter_map(|(key, value)| value.as_str().map(|s| (key, s.to_string())))
.for_each(|(key, s)| std::env::set_var(key, s));
} }
} };
set_aws_env_vars(config.load_values());
set_aws_env_vars(config.load_secrets());
let sdk_config = futures::executor::block_on(aws_config::load_from_env()); let sdk_config = futures::executor::block_on(aws_config::load_from_env());
// validate credentials or return error back up
futures::executor::block_on(
sdk_config
.credentials_provider()
.unwrap()
.provide_credentials(),
)?;
let client = Client::new(&sdk_config); let client = Client::new(&sdk_config);
Ok(Self { client, model }) Ok(Self { client, model })
@@ -63,7 +85,7 @@ impl Provider for BedrockProvider {
BEDROCK_DEFAULT_MODEL, BEDROCK_DEFAULT_MODEL,
BEDROCK_KNOWN_MODELS.iter().map(|s| s.to_string()).collect(), BEDROCK_KNOWN_MODELS.iter().map(|s| s.to_string()).collect(),
BEDROCK_DOC_LINK, BEDROCK_DOC_LINK,
vec![ConfigKey::new("AWS_PROFILE", true, false, Some("us-west-2"))], vec![ConfigKey::new("AWS_PROFILE", true, false, Some("default"))],
) )
} }