tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 6 Packages Projects Releases Wiki Activity

Generic retry and error parsing (#3558)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: Michael Neale <michael.neale@gmail.com>
This commit is contained in:
Douwe Osinga
2025-08-05 10:37:04 +02:00
committed by GitHub
parent 33f57da4c9
commit 918faddf03
26 changed files with 1190 additions and 1748 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/providers/databricks.rs
+81 -288
View File
@@ -2,7 +2,6 @@ use anyhow::Result;
use async_stream::try_stream;
use async_trait::async_trait;
use futures::TryStreamExt;
use reqwest::{Client, StatusCode};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::io;
@@ -10,43 +9,34 @@ use std::time::Duration;
use tokio::pin;
use tokio_util::io::StreamReader;
use super::api_client::{ApiClient, AuthMethod, AuthProvider};
use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::embedding::EmbeddingCapable;
use super::errors::ProviderError;
use super::formats::databricks::{create_request, response_to_message};
use super::oauth;
use super::utils::{get_model, ImageFormat};
use super::retry::ProviderRetry;
use super::utils::{get_model, handle_response_openai_compat, ImageFormat};
use crate::config::ConfigError;
use crate::impl_provider_default;
use crate::message::Message;
use crate::model::ModelConfig;
use crate::providers::formats::openai::{get_usage, response_to_streaming_message};
use crate::providers::retry::{
RetryConfig, DEFAULT_BACKOFF_MULTIPLIER, DEFAULT_INITIAL_RETRY_INTERVAL_MS,
DEFAULT_MAX_RETRIES, DEFAULT_MAX_RETRY_INTERVAL_MS,
};
use rmcp::model::Tool;
use serde_json::json;
use tokio::time::sleep;
use tokio_stream::StreamExt;
use tokio_util::codec::{FramedRead, LinesCodec};
use url::Url;
const DEFAULT_CLIENT_ID: &str = "databricks-cli";
const DEFAULT_REDIRECT_URL: &str = "http://localhost:8020";
// "offline_access" scope is used to request an OAuth 2.0 Refresh Token
// https://openid.net/specs/openid-connect-core-1_0.html#OfflineAccess
const DEFAULT_SCOPES: &[&str] = &["all-apis", "offline_access"];
/// Default timeout for API requests in seconds
const DEFAULT_TIMEOUT_SECS: u64 = 600;
/// Default initial interval for retry (in milliseconds)
const DEFAULT_INITIAL_RETRY_INTERVAL_MS: u64 = 5000;
/// Default maximum number of retries
const DEFAULT_MAX_RETRIES: usize = 6;
/// Default retry backoff multiplier
const DEFAULT_BACKOFF_MULTIPLIER: f64 = 2.0;
/// Default maximum interval for retry (in milliseconds)
const DEFAULT_MAX_RETRY_INTERVAL_MS: u64 = 320_000;
pub const DATABRICKS_DEFAULT_MODEL: &str = "databricks-claude-3-7-sonnet";
// Databricks can passthrough to a wide range of models, we only provide the default
pub const DATABRICKS_KNOWN_MODELS: &[&str] = &[
"databricks-meta-llama-3-3-70b-instruct",
"databricks-meta-llama-3-1-405b-instruct",
@@ -57,53 +47,6 @@ pub const DATABRICKS_KNOWN_MODELS: &[&str] = &[
pub const DATABRICKS_DOC_URL: &str =
"https://docs.databricks.com/en/generative-ai/external-models/index.html";
/// Retry configuration for handling rate limit errors
#[derive(Debug, Clone)]
struct RetryConfig {
/// Maximum number of retry attempts
max_retries: usize,
/// Initial interval between retries in milliseconds
initial_interval_ms: u64,
/// Multiplier for backoff (exponential)
backoff_multiplier: f64,
/// Maximum interval between retries in milliseconds
max_interval_ms: u64,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
max_retries: DEFAULT_MAX_RETRIES,
initial_interval_ms: DEFAULT_INITIAL_RETRY_INTERVAL_MS,
backoff_multiplier: DEFAULT_BACKOFF_MULTIPLIER,
max_interval_ms: DEFAULT_MAX_RETRY_INTERVAL_MS,
}
}
}
impl RetryConfig {
/// Calculate the delay for a specific retry attempt (with jitter)
fn delay_for_attempt(&self, attempt: usize) -> Duration {
if attempt == 0 {
return Duration::from_millis(0);
}
// Calculate exponential backoff
let exponent = (attempt - 1) as u32;
let base_delay_ms = (self.initial_interval_ms as f64
* self.backoff_multiplier.powi(exponent as i32)) as u64;
// Apply max limit
let capped_delay_ms = std::cmp::min(base_delay_ms, self.max_interval_ms);
// Add jitter (+/-20% randomness) to avoid thundering herd problem
let jitter_factor = 0.8 + (rand::random::<f64>() * 0.4); // Between 0.8 and 1.2
let jittered_delay_ms = (capped_delay_ms as f64 * jitter_factor) as u64;
Duration::from_millis(jittered_delay_ms)
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum DatabricksAuth {
Token(String),
@@ -116,7 +59,6 @@ pub enum DatabricksAuth {
}
impl DatabricksAuth {
/// Create a new OAuth configuration with default values
pub fn oauth(host: String) -> Self {
Self::OAuth {
host,
@@ -125,16 +67,36 @@ impl DatabricksAuth {
scopes: DEFAULT_SCOPES.iter().map(|s| s.to_string()).collect(),
}
}
pub fn token(token: String) -> Self {
Self::Token(token)
}
}
struct DatabricksAuthProvider {
auth: DatabricksAuth,
}
#[async_trait]
impl AuthProvider for DatabricksAuthProvider {
async fn get_auth_header(&self) -> Result<(String, String)> {
let token = match &self.auth {
DatabricksAuth::Token(token) => token.clone(),
DatabricksAuth::OAuth {
host,
client_id,
redirect_url,
scopes,
} => oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?,
};
Ok(("Authorization".to_string(), format!("Bearer {}", token)))
}
}
#[derive(Debug, serde::Serialize)]
pub struct DatabricksProvider {
#[serde(skip)]
client: Client,
host: String,
api_client: ApiClient,
auth: DatabricksAuth,
model: ModelConfig,
image_format: ImageFormat,
@@ -148,8 +110,6 @@ impl DatabricksProvider {
pub fn from_env(model: ModelConfig) -> Result<Self> {
let config = crate::config::Config::global();
// For compatibility for now we check both config and secret for databricks host
// but it is not actually a secret value
let mut host: Result<String, ConfigError> = config.get_param("DATABRICKS_HOST");
if host.is_err() {
host = config.get_secret("DATABRICKS_HOST")
@@ -163,38 +123,29 @@ impl DatabricksProvider {
}
let host = host?;
let client = Client::builder()
.timeout(Duration::from_secs(DEFAULT_TIMEOUT_SECS))
.build()?;
// Load optional retry configuration from environment
let retry_config = Self::load_retry_config(config);
// If we find a databricks token we prefer that
if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") {
return Ok(Self {
client,
host,
auth: DatabricksAuth::token(api_key),
model,
image_format: ImageFormat::OpenAi,
retry_config,
});
}
let auth = if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") {
DatabricksAuth::token(api_key)
} else {
DatabricksAuth::oauth(host.clone())
};
let auth_method =
AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone() }));
let api_client =
ApiClient::with_timeout(host, auth_method, Duration::from_secs(DEFAULT_TIMEOUT_SECS))?;
// Otherwise use Oauth flow
Ok(Self {
client,
auth: DatabricksAuth::oauth(host.clone()),
host,
api_client,
auth,
model,
image_format: ImageFormat::OpenAi,
retry_config,
})
}
/// Loads retry configuration from environment variables or uses defaults.
fn load_retry_config(config: &crate::config::Config) -> RetryConfig {
let max_retries = config
.get_param("DATABRICKS_MAX_RETRIES")
@@ -228,184 +179,36 @@ impl DatabricksProvider {
}
}
/// Create a new DatabricksProvider with the specified host and token
///
/// # Arguments
///
/// * `host` - The Databricks host URL
/// * `token` - The Databricks API token
///
/// # Returns
///
/// Returns a Result containing the new DatabricksProvider instance
pub fn from_params(host: String, api_key: String, model: ModelConfig) -> Result<Self> {
let client = Client::builder()
.timeout(Duration::from_secs(600))
.build()?;
let auth = DatabricksAuth::token(api_key);
let auth_method =
AuthMethod::Custom(Box::new(DatabricksAuthProvider { auth: auth.clone() }));
let api_client = ApiClient::with_timeout(host, auth_method, Duration::from_secs(600))?;
Ok(Self {
client,
host,
auth: DatabricksAuth::token(api_key),
api_client,
auth,
model,
image_format: ImageFormat::OpenAi,
retry_config: RetryConfig::default(),
})
}
async fn ensure_auth_header(&self) -> Result<String> {
match &self.auth {
DatabricksAuth::Token(token) => Ok(format!("Bearer {}", token)),
DatabricksAuth::OAuth {
host,
client_id,
redirect_url,
scopes,
} => {
let token =
oauth::get_oauth_token_async(host, client_id, redirect_url, scopes).await?;
Ok(format!("Bearer {}", token))
}
}
}
async fn post(&self, payload: &Value) -> Result<Value, ProviderError> {
// Check if this is an embedding request by looking at the payload structure
let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none();
let path = if is_embedding {
// For embeddings, use the embeddings endpoint
format!("serving-endpoints/{}/invocations", "text-embedding-3-small")
fn get_endpoint_path(&self, is_embedding: bool) -> String {
if is_embedding {
"serving-endpoints/text-embedding-3-small/invocations".to_string()
} else {
// For chat completions, use the model name in the path
format!("serving-endpoints/{}/invocations", self.model.model_name)
};
match self.post_with_retry(path.as_str(), payload).await {
Ok(res) => res.json().await.map_err(|_| {
ProviderError::RequestFailed("Response body is not valid JSON".to_string())
}),
Err(e) => Err(e),
}
}
async fn post_with_retry(
&self,
path: &str,
payload: &Value,
) -> Result<reqwest::Response, ProviderError> {
let base_url = Url::parse(&self.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let url = base_url.join(path).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
async fn post(&self, payload: Value) -> Result<Value, ProviderError> {
let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none();
let path = self.get_endpoint_path(is_embedding);
let mut attempts = 0;
loop {
let auth_header = self.ensure_auth_header().await?;
let response = self
.client
.post(url.clone())
.header("Authorization", auth_header)
.json(payload)
.send()
.await?;
let status = response.status();
break match status {
StatusCode::OK => Ok(response),
StatusCode::TOO_MANY_REQUESTS
| StatusCode::INTERNAL_SERVER_ERROR
| StatusCode::SERVICE_UNAVAILABLE => {
if attempts < self.retry_config.max_retries {
attempts += 1;
tracing::warn!(
"{}: retrying ({}/{})",
status,
attempts,
self.retry_config.max_retries
);
let delay = self.retry_config.delay_for_attempt(attempts);
tracing::info!("Backing off for {:?} before retry", delay);
sleep(delay).await;
continue;
}
Err(match status {
StatusCode::TOO_MANY_REQUESTS => {
ProviderError::RateLimitExceeded("Rate limit exceeded".to_string())
}
_ => ProviderError::ServerError("Server error".to_string()),
})
}
StatusCode::BAD_REQUEST => {
// Databricks provides a generic 'error' but also includes 'external_model_message' which is provider specific
// We try to extract the error message from the payload and check for phrases that indicate context length exceeded
let bytes = response.bytes().await?;
let payload_str = String::from_utf8_lossy(&bytes).to_lowercase();
let check_phrases = [
"too long",
"context length",
"context_length_exceeded",
"reduce the length",
"token count",
"exceeds",
"exceed context limit",
"input length",
"max_tokens",
"decrease input length",
"context limit",
];
if check_phrases.iter().any(|c| payload_str.contains(c)) {
return Err(ProviderError::ContextLengthExceeded(payload_str));
}
let mut error_msg = "Unknown error".to_string();
if let Ok(response_json) = serde_json::from_slice::<Value>(&bytes) {
// try to convert message to string, if that fails use external_model_message
error_msg = response_json
.get("message")
.and_then(|m| m.as_str())
.or_else(|| {
response_json
.get("external_model_message")
.and_then(|ext| ext.get("message"))
.and_then(|m| m.as_str())
})
.unwrap_or("Unknown error")
.to_string();
}
tracing::debug!(
"{}",
format!(
"Provider request failed with status: {}. Payload: {:?}",
status, payload_str
)
);
return Err(ProviderError::RequestFailed(format!(
"Request failed with status: {}. Message: {}",
status, error_msg
)));
}
_ => {
tracing::debug!(
"{}",
format!(
"Provider request failed with status: {}. Payload: {:?}",
status,
response.text().await.ok().unwrap_or_default()
)
);
return Err(ProviderError::RequestFailed(format!(
"Request failed with status: {}",
status
)));
}
};
}
let response = self.api_client.response_post(&path, &payload).await?;
handle_response_openai_compat(response).await
}
}
@@ -426,6 +229,10 @@ impl Provider for DatabricksProvider {
)
}
fn retry_config(&self) -> RetryConfig {
self.retry_config.clone()
}
fn get_model_config(&self) -> ModelConfig {
self.model.clone()
}
@@ -441,15 +248,13 @@ impl Provider for DatabricksProvider {
tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
let mut payload = create_request(&self.model, system, messages, tools, &self.image_format)?;
// Remove the model key which is part of the url with databricks
payload
.as_object_mut()
.expect("payload should have model key")
.remove("model");
let response = self.post(&payload).await?;
let response = self.with_retry(|| self.post(payload.clone())).await?;
// Parse response
let message = response_to_message(&response)?;
let usage = response.get("usage").map(get_usage).unwrap_or_else(|| {
tracing::debug!("Failed to get usage data");
@@ -468,7 +273,6 @@ impl Provider for DatabricksProvider {
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let mut payload = create_request(&self.model, system, messages, tools, &self.image_format)?;
// Remove the model key which is part of the url with databricks
payload
.as_object_mut()
.expect("payload should have model key")
@@ -479,18 +283,24 @@ impl Provider for DatabricksProvider {
.unwrap()
.insert("stream".to_string(), Value::Bool(true));
let path = self.get_endpoint_path(false);
let response = self
.post_with_retry(
format!("serving-endpoints/{}/invocations", self.model.model_name).as_str(),
&payload,
)
.with_retry(|| async {
let resp = self.api_client.response_post(&path, &payload).await?;
if !resp.status().is_success() {
return Err(ProviderError::RequestFailed(format!(
"HTTP {}: {}",
resp.status(),
resp.text().await.unwrap_or_default()
)));
}
Ok(resp)
})
.await?;
// Map reqwest error to io::Error
let stream = response.bytes_stream().map_err(io::Error::other);
let model_config = self.model.clone();
// Wrap in a line decoder and yield lines inside the stream
Ok(Box::pin(try_stream! {
let stream_reader = StreamReader::new(stream);
let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from);
@@ -519,32 +329,16 @@ impl Provider for DatabricksProvider {
.map_err(|e| ProviderError::ExecutionError(e.to_string()))
}
async fn fetch_supported_models_async(&self) -> Result<Option<Vec<String>>, ProviderError> {
let base_url = Url::parse(&self.host)
.map_err(|e| ProviderError::RequestFailed(format!("Invalid base URL: {e}")))?;
let url = base_url.join("api/2.0/serving-endpoints").map_err(|e| {
ProviderError::RequestFailed(format!("Failed to construct endpoint URL: {e}"))
})?;
let auth_header = match self.ensure_auth_header().await {
Ok(header) => header,
Err(e) => {
tracing::warn!("Failed to authorize with Databricks: {}", e);
return Ok(None); // Return None to fall back to manual input
}
};
async fn fetch_supported_models(&self) -> Result<Option<Vec<String>>, ProviderError> {
let response = match self
.client
.get(url)
.header("Authorization", auth_header)
.send()
.api_client
.response_get("api/2.0/serving-endpoints")
.await
{
Ok(resp) => resp,
Err(e) => {
tracing::warn!("Failed to fetch Databricks models: {}", e);
return Ok(None); // Return None to fall back to manual input
return Ok(None);
}
};
@@ -559,7 +353,7 @@ impl Provider for DatabricksProvider {
} else {
tracing::warn!("Failed to fetch Databricks models: {}", status);
}
return Ok(None); // Return None to fall back to manual input
return Ok(None);
}
let json: Value = match response.json().await {
@@ -610,12 +404,11 @@ impl EmbeddingCapable for DatabricksProvider {
return Ok(vec![]);
}
// Create request in Databricks format for embeddings
let request = json!({
"input": texts,
});
let response = self.post(&request).await?;
let response = self.with_retry(|| self.post(request.clone())).await?;
let embeddings = response["data"]
.as_array()