Skip to content
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 49 additions & 3 deletions crates/goose/src/providers/bedrock.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ use std::collections::HashMap;

use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::retry::{ProviderRetry, RetryConfig};
use crate::conversation::message::Message;
use crate::impl_provider_default;
use crate::model::ModelConfig;
Expand All @@ -23,17 +23,28 @@ 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 = "anthropic.claude-3-5-sonnet-20240620-v1:0";
pub const BEDROCK_DEFAULT_MODEL: &str = "anthropic.claude-sonnet-4-20250514-v1:0";
pub const BEDROCK_KNOWN_MODELS: &[&str] = &[
"anthropic.claude-3-5-sonnet-20240620-v1:0",
"anthropic.claude-3-5-sonnet-20241022-v2:0",
"anthropic.claude-3-7-sonnet-20250219-v1:0",
"anthropic.claude-sonnet-4-20250514-v1:0",
"anthropic.claude-opus-4-20250514-v1:0",
"anthropic.claude-opus-4-1-20250805-v1:0"
];

pub const BEDROCK_DEFAULT_MAX_RETRIES: usize = 6;
pub const BEDROCK_DEFAULT_INITIAL_RETRY_INTERVAL_MS: u64 = 2000;
pub const BEDROCK_DEFAULT_BACKOFF_MULTIPLIER: f64 = 2.0;
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,
}

impl BedrockProvider {
Expand Down Expand Up @@ -65,7 +76,38 @@ impl BedrockProvider {
)?;
let client = Client::new(&sdk_config);

Ok(Self { client, model })
let retry_config = Self::load_retry_config(config);

Ok(Self {
client,
model,
retry_config,
})
}

fn load_retry_config(config: &crate::config::Config) -> RetryConfig {
let max_retries = config
.get_param::<usize>("BEDROCK_MAX_RETRIES")
.unwrap_or(BEDROCK_DEFAULT_MAX_RETRIES);

let initial_interval_ms = config
.get_param::<u64>("BEDROCK_INITIAL_RETRY_INTERVAL_MS")
.unwrap_or(BEDROCK_DEFAULT_INITIAL_RETRY_INTERVAL_MS);

let backoff_multiplier = config
.get_param::<f64>("BEDROCK_BACKOFF_MULTIPLIER")
.unwrap_or(BEDROCK_DEFAULT_BACKOFF_MULTIPLIER);

let max_interval_ms = config
.get_param::<u64>("BEDROCK_MAX_RETRY_INTERVAL_MS")
.unwrap_or(BEDROCK_DEFAULT_MAX_RETRY_INTERVAL_MS);

RetryConfig {
max_retries,
initial_interval_ms,
backoff_multiplier,
max_interval_ms,
}
}

async fn converse(
Expand Down Expand Up @@ -147,6 +189,10 @@ impl Provider for BedrockProvider {
)
}

fn retry_config(&self) -> RetryConfig {
self.retry_config.clone()
}

fn get_model_config(&self) -> ModelConfig {
self.model.clone()
}
Expand Down
7 changes: 6 additions & 1 deletion crates/goose/src/providers/retry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -114,4 +114,9 @@ pub trait ProviderRetry {
}
}

impl<P: Provider> ProviderRetry for P {}
// Let specific providers define their retry config if desired
impl<P: Provider> ProviderRetry for P {
fn retry_config(&self) -> RetryConfig {
Provider::retry_config(self)
}
}
4 changes: 2 additions & 2 deletions documentation/docs/getting-started/providers.md
Original file line number Diff line number Diff line change
Expand Up @@ -21,11 +21,11 @@ Goose is compatible with a wide range of LLM providers, allowing you to choose a

| Provider | Description | Parameters |
|-----------------------------------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| [Amazon Bedrock](https://aws.amazon.com/bedrock/) | Offers a variety of foundation models, including Claude, Jurassic-2, and others. **AWS environment variables must be set in advance, not configured through `goose configure`** | `AWS_PROFILE`, or `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_REGION`, ... |
| [Amazon Bedrock](https://aws.amazon.com/bedrock/) | Offers a variety of foundation models, including Claude, Jurassic-2, and others. **AWS environment variables must be set in advance, not configured through `goose configure`** | `AWS_PROFILE`, or `AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, `AWS_REGION`, ... and optionally `BEDROCK_MAX_RETRIES` (6), `BEDROCK_INITIAL_RETRY_INTERVAL_MS` (2000), `BEDROCK_BACKOFF_MULTIPLIER` (2.0), `BEDROCK_MAX_RETRY_INTERVAL_MS` (120000) |
| [Amazon SageMaker TGI](https://docs.aws.amazon.com/sagemaker/latest/dg/realtime-endpoints.html) | Run Text Generation Inference models through Amazon SageMaker endpoints. **AWS credentials must be configured in advance.** | `SAGEMAKER_ENDPOINT_NAME`, `AWS_REGION` (optional), `AWS_PROFILE` (optional) |
| [Anthropic](https://www.anthropic.com/) | Offers Claude, an advanced AI model for natural language tasks. | `ANTHROPIC_API_KEY`, `ANTHROPIC_HOST` (optional) |
| [Azure OpenAI](https://learn.microsoft.com/en-us/azure/ai-services/openai/) | Access Azure-hosted OpenAI models, including GPT-4 and GPT-3.5. Supports both API key and Azure credential chain authentication. | `AZURE_OPENAI_ENDPOINT`, `AZURE_OPENAI_DEPLOYMENT_NAME`, `AZURE_OPENAI_API_KEY` (optional) |
| [Databricks](https://www.databricks.com/) | Unified data analytics and AI platform for building and deploying models. | `DATABRICKS_HOST`, `DATABRICKS_TOKEN` |
| [Databricks](https://www.databricks.com/) | Unified data analytics and AI platform for building and deploying models. | `DATABRICKS_HOST`, `DATABRICKS_TOKEN` and optionally `DATABRICKS_MAX_RETRIES` (3), `DATABRICKS_INITIAL_RETRY_INTERVAL_MS` (1000), `DATABRICKS_BACKOFF_MULTIPLIER` (2.0), `DATABRICKS_MAX_RETRY_INTERVAL_MS` (30000) |
| [Docker Model Runner](https://docs.docker.com/ai/model-runner/) | Local models running in Docker Desktop or Docker CE with OpenAI-compatible API endpoints. **Because this provider runs locally, you must first [download a model](#local-llms).** | `OPENAI_HOST`, `OPENAI_BASE_PATH` |
| [Gemini](https://ai.google.dev/gemini-api/docs) | Advanced LLMs by Google with multimodal capabilities (text, images). | `GOOGLE_API_KEY` |
| [GCP Vertex AI](https://cloud.google.com/vertex-ai) | Google Cloud's Vertex AI platform, supporting Gemini and Claude models. **Credentials must be [configured in advance](https://cloud.google.com/vertex-ai/docs/authentication).** | `GCP_PROJECT_ID`, `GCP_LOCATION` and optionally `GCP_MAX_RATE_LIMIT_RETRIES` (5), `GCP_MAX_OVERLOADED_RETRIES` (5), `GCP_INITIAL_RETRY_INTERVAL_MS` (5000), `GCP_BACKOFF_MULTIPLIER` (2.0), `GCP_MAX_RETRY_INTERVAL_MS` (320_000). |
Expand Down
Loading