ai-space/goose
3
0
Fork 0
mirror of https://github.com/aaif-goose/goose.git synced 2026-07-03 14:10:03 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat: Add configurable Bedrock retry parameters (#4316)

Browse Source 
Signed-off-by: Dan Wuensch <dan.wuensch@ncino.com>
Co-authored-by: angiejones <jones.angie@gmail.com>
This commit is contained in:
Dan Wuensch
2025-09-04 12:26:29 -04:00
committed by GitHub
parent 429aeffc77
commit 627696a68a
4 changed files with 98 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/providers/bedrock.rs
+49 -3
View File
@@ -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;
@@ -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 {
@@ -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(
@@ -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()
}
crates/goose/src/providers/retry.rs
+6 -1
View File
@@ -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)
}
}
documentation/docs/getting-started/providers.md
+2 -2
View File
@@ -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` |
| [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` |
| [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). |
documentation/docs/guides/environment-variables.md
+41
View File
@@ -99,6 +99,47 @@ export GOOSE_PLANNER_PROVIDER="openai"
export GOOSE_PLANNER_MODEL="gpt-4"
```
### Provider Retries
Configurable retry parameters for LLM providers.
#### AWS Bedrock
| Variable | Purpose | Default |
|---------------------|-------------|---------|
| `BEDROCK_MAX_RETRIES` | The max number of retry attempts before giving up | 6 |
| `BEDROCK_INITIAL_RETRY_INTERVAL_MS` | How long to wait (in milliseconds) before the first retry | 2000 |
| `BEDROCK_BACKOFF_MULTIPLIER` | The factor by which the retry interval increases after each attempt | 2 (doubles every time) |
| `BEDROCK_MAX_RETRY_INTERVAL_MS` | The cap on the retry interval in milliseconds | 120000 |
**Examples**
```bash
export BEDROCK_MAX_RETRIES=10 # 10 retry attempts
export BEDROCK_INITIAL_RETRY_INTERVAL_MS=1000 # start with 1 second before first retry
export BEDROCK_BACKOFF_MULTIPLIER=3 # each retry waits 3x longer than the previous
export BEDROCK_MAX_RETRY_INTERVAL_MS=300000 # cap the maximum retry delay at 5 min
```
#### Databricks
| Variable | Purpose | Default |
|---------------------|-------------|---------|
| `DATABRICKS_MAX_RETRIES` | The max number of retry attempts before giving up | 3 |
| `DATABRICKS_INITIAL_RETRY_INTERVAL_MS` | How long to wait (in milliseconds) before the first retry | 1000 |
| `DATABRICKS_BACKOFF_MULTIPLIER` | The factor by which the retry interval increases after each attempt | 2 (doubles every time) |
| `DATABRICKS_MAX_RETRY_INTERVAL_MS` | The cap on the retry interval in milliseconds | 30000 |
**Examples**
```bash
export DATABRICKS_MAX_RETRIES=5 # 5 retry attempts
export DATABRICKS_INITIAL_RETRY_INTERVAL_MS=500 # start with 0.5 second before first retry
export DATABRICKS_BACKOFF_MULTIPLIER=2 # each retry waits 2x longer than the previous
export DATABRICKS_MAX_RETRY_INTERVAL_MS=60000 # cap the maximum retry delay at 1 min
```
## Session Management
These variables control how Goose manages conversation sessions and context.