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

feat: add retry for google (#1854)

Browse Source
This commit is contained in:
Yingjie He
2025-03-25 12:39:32 -07:00
committed by GitHub
parent d619c95273
commit b310bb00b4
1 changed files with 42 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
crates/goose/src/providers/google.rs
View File

@@ -78,15 +78,49 @@ impl GoogleProvider {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
let response = self
.client
.post(url)
.header("CONTENT_TYPE", "application/json")
.json(&payload)
.send()
.await?;
let max_retries = 3;
let mut retries = 0;
let base_delay = Duration::from_secs(2);
handle_response_google_compat(response).await
loop {
let response = self
.client
.post(url.clone()) // Clone the URL for each retry
.header("CONTENT_TYPE", "application/json")
.json(&payload)
.send()
.await;
match response {
Ok(res) => {
match handle_response_google_compat(res).await {
Ok(result) => return Ok(result),
Err(ProviderError::RateLimitExceeded(_)) => {
retries += 1;
if retries > max_retries {
return Err(ProviderError::RateLimitExceeded(
"Max retries exceeded for rate limit error".to_string(),
));
}
let delay = 2u64.pow(retries);
let total_delay = Duration::from_secs(delay) + base_delay;
println!("Rate limit hit. Retrying in {:?}", total_delay);
tokio::time::sleep(total_delay).await;
continue;
}
Err(err) => return Err(err), // Other errors
}
}
Err(err) => {
return Err(ProviderError::RequestFailed(format!(
"Request failed: {}",
err
)));
}
}
}
}
}