From b3cf3cbcffedb123302bba72ab25b0e26f5c507b Mon Sep 17 00:00:00 2001 From: Jack Amadeo Date: Thu, 18 Jun 2026 15:22:32 -0400 Subject: [PATCH] delete embeddings support (#9865) --- crates/goose-providers/src/base.rs | 15 ---- crates/goose/src/providers/databricks.rs | 109 ++--------------------- crates/goose/src/providers/embedding.rs | 28 ------ crates/goose/src/providers/litellm.rs | 47 ---------- crates/goose/src/providers/mod.rs | 1 - crates/goose/src/providers/openai.rs | 78 ---------------- 6 files changed, 5 insertions(+), 273 deletions(-) delete mode 100644 crates/goose/src/providers/embedding.rs diff --git a/crates/goose-providers/src/base.rs b/crates/goose-providers/src/base.rs index b14579314..644cdb934 100644 --- a/crates/goose-providers/src/base.rs +++ b/crates/goose-providers/src/base.rs @@ -354,10 +354,6 @@ pub trait Provider: Send + Sync { )) } - fn supports_embeddings(&self) -> bool { - false - } - /// Whether the provider manages its own conversation context (e.g. CLI /// wrappers like Claude Code or Gemini CLI). When true, goose-side /// context management such as tool-pair summarization is skipped because @@ -370,17 +366,6 @@ pub trait Provider: Send + Sync { false } - /// Create embeddings if supported. Default implementation returns an error. - async fn create_embeddings( - &self, - _session_id: &str, - _texts: Vec, - ) -> Result>, ProviderError> { - Err(ProviderError::ExecutionError( - "This provider does not support embeddings".to_string(), - )) - } - /// Configure OAuth authentication for this provider /// /// This method is called when a provider has configuration keys marked with oauth_flow = true. diff --git a/crates/goose/src/providers/databricks.rs b/crates/goose/src/providers/databricks.rs index 70e5fcbad..1ff851071 100644 --- a/crates/goose/src/providers/databricks.rs +++ b/crates/goose/src/providers/databricks.rs @@ -17,12 +17,11 @@ use super::base::{ DEFAULT_PROVIDER_TIMEOUT_SECS, }; use super::databricks_auth::{DatabricksAuth, DatabricksAuthProvider}; -use super::embedding::EmbeddingCapable; use super::formats::databricks::{create_request_for_provider, DATABRICKS_PROVIDER_NAME}; use super::formats::openai_responses::create_responses_request; use super::openai_compatible::{ - handle_response_openai_compat, handle_status, map_http_error_to_provider_error, sanitize_url, - stream_openai_compat, stream_responses_compat, + handle_status, map_http_error_to_provider_error, sanitize_url, stream_openai_compat, + stream_responses_compat, }; use super::retry::ProviderRetry; use super::utils::RequestLog; @@ -86,8 +85,6 @@ pub struct DatabricksProvider { #[serde(skip)] retry_config: RetryConfig, #[serde(skip)] - fast_retry_config: RetryConfig, - #[serde(skip)] name: String, #[serde(skip)] token_cache: Arc>>, @@ -117,7 +114,6 @@ impl DatabricksProvider { let host = host?; let retry_config = Self::load_retry_config(config); - let fast_retry_config = Self::load_fast_retry_config(config); let auth = if let Ok(api_key) = config.get_secret("DATABRICKS_TOKEN") { DatabricksAuth::token(api_key) @@ -148,7 +144,6 @@ impl DatabricksProvider { model: model.clone(), image_format: ImageFormat::OpenAi, retry_config, - fast_retry_config, name: DATABRICKS_PROVIDER_NAME.to_string(), token_cache, instance_id: Self::resolve_instance_id(), @@ -194,11 +189,6 @@ impl DatabricksProvider { ) } - fn load_fast_retry_config(_config: &crate::config::Config) -> RetryConfig { - // Fast models are hardcoded to 0 retries for quick failure on Databricks - RetryConfig::new(0, 0, 1.0, 0) - } - pub fn from_params(host: String, api_key: String, model: ModelConfig) -> Result { let token_cache = Arc::new(Mutex::new(Some(api_key.clone()))); let auth = DatabricksAuth::token(api_key); @@ -220,7 +210,6 @@ impl DatabricksProvider { model, image_format: ImageFormat::OpenAi, retry_config: RetryConfig::default(), - fast_retry_config: RetryConfig::new(0, 0, 1.0, 0), name: DATABRICKS_PROVIDER_NAME.to_string(), token_cache, instance_id: Self::resolve_instance_id(), @@ -540,15 +529,8 @@ impl DatabricksProvider { } } - fn get_endpoint_path( - &self, - model_name: &str, - is_embedding: bool, - is_responses_model: bool, - ) -> String { - if is_embedding { - "serving-endpoints/text-embedding-3-small/invocations".to_string() - } else if is_responses_model { + fn get_endpoint_path(&self, model_name: &str, is_responses_model: bool) -> String { + if is_responses_model { "serving-endpoints/responses".to_string() } else { let (clean_name, _) = extract_reasoning_effort(model_name); @@ -564,32 +546,6 @@ impl DatabricksProvider { .to_string() }) } - - async fn post( - &self, - session_id: Option<&str>, - mut payload: Value, - model_name: Option<&str>, - ) -> Result { - let is_embedding = payload.get("input").is_some() && payload.get("messages").is_none(); - let model_to_use = model_name.unwrap_or(&self.model.model_name); - let (endpoint_name, _) = extract_reasoning_effort(model_to_use); - let endpoint_info = self.resolve_endpoint_info_cached(&endpoint_name).await.ok(); - let is_responses_model = Self::uses_responses_api(endpoint_info.as_ref(), &[model_to_use]); - let path = self.get_endpoint_path(model_to_use, is_embedding, is_responses_model); - - if let Some(session_id) = session_id { - if let Some(client_request_id) = self.build_client_request_id(session_id) { - payload["client_request_id"] = Value::String(client_request_id); - } - } - - let response = self - .api_client - .response_post(session_id, &path, &payload) - .await?; - handle_response_openai_compat(response).await - } } impl ProviderDef for DatabricksProvider { @@ -660,7 +616,7 @@ impl Provider for DatabricksProvider { let path = if is_responses_model { "serving-endpoints/responses".to_string() } else { - self.get_endpoint_path(&model_config.model_name, false, is_responses_model) + self.get_endpoint_path(&model_config.model_name, is_responses_model) }; let client_request_id = self.build_client_request_id(session_id); @@ -818,20 +774,6 @@ impl Provider for DatabricksProvider { } } - fn supports_embeddings(&self) -> bool { - true - } - - async fn create_embeddings( - &self, - session_id: &str, - texts: Vec, - ) -> Result>, ProviderError> { - EmbeddingCapable::create_embeddings(self, session_id, texts) - .await - .map_err(|e| ProviderError::ExecutionError(e.to_string())) - } - async fn fetch_supported_models(&self) -> Result, ProviderError> { Ok(self .fetch_supported_model_info() @@ -895,47 +837,6 @@ impl Provider for DatabricksProvider { } } -#[async_trait] -impl EmbeddingCapable for DatabricksProvider { - async fn create_embeddings( - &self, - session_id: &str, - texts: Vec, - ) -> Result>> { - if texts.is_empty() { - return Ok(vec![]); - } - - let request = json!({ - "input": texts, - }); - - let response = self - .with_retry_config( - || self.post(Some(session_id), request.clone(), None), - self.fast_retry_config.clone(), - ) - .await?; - - let embeddings = response["data"] - .as_array() - .ok_or_else(|| anyhow::anyhow!("Invalid response format: missing data array"))? - .iter() - .map(|item| { - item["embedding"] - .as_array() - .ok_or_else(|| anyhow::anyhow!("Invalid embedding format"))? - .iter() - .map(|v| v.as_f64().map(|f| f as f32)) - .collect::>>() - .ok_or_else(|| anyhow::anyhow!("Invalid embedding values")) - }) - .collect::>>>()?; - - Ok(embeddings) - } -} - #[cfg(test)] mod tests { use super::*; diff --git a/crates/goose/src/providers/embedding.rs b/crates/goose/src/providers/embedding.rs deleted file mode 100644 index 3f956b11c..000000000 --- a/crates/goose/src/providers/embedding.rs +++ /dev/null @@ -1,28 +0,0 @@ -use anyhow::Result; -use async_trait::async_trait; -use serde::{Deserialize, Serialize}; - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct EmbeddingRequest { - pub input: Vec, - pub model: String, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct EmbeddingResponse { - pub data: Vec, -} - -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct EmbeddingData { - pub embedding: Vec, -} - -#[async_trait] -pub trait EmbeddingCapable { - async fn create_embeddings( - &self, - session_id: &str, - texts: Vec, - ) -> Result>>; -} diff --git a/crates/goose/src/providers/litellm.rs b/crates/goose/src/providers/litellm.rs index 46fb5eecd..45ac5bf14 100644 --- a/crates/goose/src/providers/litellm.rs +++ b/crates/goose/src/providers/litellm.rs @@ -12,7 +12,6 @@ use super::base::{ ConfigKey, MessageStream, ModelInfo, Provider, ProviderDef, ProviderMetadata, DEFAULT_PROVIDER_TIMEOUT_SECS, }; -use super::embedding::EmbeddingCapable; use super::openai_compatible::handle_response_openai_compat; use super::retry::ProviderRetry; use super::utils::{get_model, RequestLog}; @@ -259,10 +258,6 @@ impl Provider for LiteLLMProvider { )) } - fn supports_embeddings(&self) -> bool { - true - } - async fn supports_cache_control(&self) -> bool { if let Ok(models) = self.get_or_fetch_models().await { if let Some(model_info) = models.iter().find(|m| m.name == self.model.model_name) { @@ -279,48 +274,6 @@ impl Provider for LiteLLMProvider { } } -#[async_trait] -impl EmbeddingCapable for LiteLLMProvider { - async fn create_embeddings( - &self, - session_id: &str, - texts: Vec, - ) -> Result>, anyhow::Error> { - let embedding_model = std::env::var("GOOSE_EMBEDDING_MODEL") - .unwrap_or_else(|_| "text-embedding-3-small".to_string()); - - let payload = json!({ - "input": texts, - "model": embedding_model, - "encoding_format": "float" - }); - - let response = self - .api_client - .response_post(Some(session_id), "v1/embeddings", &payload) - .await?; - let response_text = response.text().await?; - let response_json: Value = serde_json::from_str(&response_text)?; - - let data = response_json["data"] - .as_array() - .ok_or_else(|| anyhow::anyhow!("Missing data field"))?; - - let mut embeddings = Vec::new(); - for item in data { - let embedding: Vec = item["embedding"] - .as_array() - .ok_or_else(|| anyhow::anyhow!("Missing embedding field"))? - .iter() - .map(|v| v.as_f64().unwrap_or(0.0) as f32) - .collect(); - embeddings.push(embedding); - } - - Ok(embeddings) - } -} - /// Updates the request payload to include cache control headers for automatic prompt caching /// Adds ephemeral cache control to the last 2 user messages, system message, and last tool pub fn update_request_for_cache_control(original_payload: &Value) -> Value { diff --git a/crates/goose/src/providers/mod.rs b/crates/goose/src/providers/mod.rs index 00fefee82..031c77c26 100644 --- a/crates/goose/src/providers/mod.rs +++ b/crates/goose/src/providers/mod.rs @@ -26,7 +26,6 @@ pub mod cursor_agent; pub mod databricks; pub mod databricks_auth; pub mod databricks_v2; -pub mod embedding; pub mod formats; mod gcpauth; pub mod gcpvertexai; diff --git a/crates/goose/src/providers/openai.rs b/crates/goose/src/providers/openai.rs index 255bd15fe..619c75c16 100644 --- a/crates/goose/src/providers/openai.rs +++ b/crates/goose/src/providers/openai.rs @@ -2,7 +2,6 @@ use super::api_client::{ApiClient, AuthMethod}; use super::base::{ ConfigKey, ModelInfo, Provider, ProviderDef, ProviderMetadata, DEFAULT_PROVIDER_TIMEOUT_SECS, }; -use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse}; use super::formats::openai_responses::{ create_responses_request, get_responses_usage, responses_api_to_message, ResponsesApiResponse, }; @@ -35,7 +34,6 @@ pub(crate) const OPEN_AI_DEFAULT_BASE_PATH: &str = "v1/chat/completions"; const OPEN_AI_VERSIONLESS_BASE_PATH: &str = "chat/completions"; const OPEN_AI_DEFAULT_RESPONSES_PATH: &str = "v1/responses"; const OPEN_AI_DEFAULT_MODELS_PATH: &str = "v1/models"; -const OPEN_AI_DEFAULT_EMBEDDINGS_PATH: &str = "v1/embeddings"; pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o"; pub const OPEN_AI_DEFAULT_FAST_MODEL: &str = "gpt-4o-mini"; pub const OPEN_AI_KNOWN_MODELS: &[(&str, usize)] = &[ @@ -813,20 +811,6 @@ impl Provider for OpenAiProvider { self.fetch_models_from_api().await } - fn supports_embeddings(&self) -> bool { - true - } - - async fn create_embeddings( - &self, - session_id: &str, - texts: Vec, - ) -> Result>, ProviderError> { - EmbeddingCapable::create_embeddings(self, session_id, texts) - .await - .map_err(|e| ProviderError::ExecutionError(e.to_string())) - } - async fn stream( &self, model_config: &ModelConfig, @@ -953,68 +937,6 @@ fn parse_custom_headers(s: String) -> HashMap { .collect() } -#[async_trait] -impl EmbeddingCapable for OpenAiProvider { - async fn create_embeddings( - &self, - session_id: &str, - texts: Vec, - ) -> Result>> { - if texts.is_empty() { - return Ok(vec![]); - } - - let embedding_model = std::env::var("GOOSE_EMBEDDING_MODEL") - .unwrap_or_else(|_| "text-embedding-3-small".to_string()); - - let request = EmbeddingRequest { - input: texts, - model: embedding_model, - }; - - let response = self - .with_retry(|| async { - let request_clone = EmbeddingRequest { - input: request.input.clone(), - model: request.model.clone(), - }; - let request_value = serde_json::to_value(request_clone) - .map_err(|e| ProviderError::ExecutionError(e.to_string()))?; - let embeddings_path = Self::map_base_path( - &self.base_path, - "embeddings", - OPEN_AI_DEFAULT_EMBEDDINGS_PATH, - ); - self.api_client - .api_post(Some(session_id), &embeddings_path, &request_value) - .await - .map_err(|e| ProviderError::ExecutionError(e.to_string())) - }) - .await?; - - if response.status != StatusCode::OK { - let error_text = response - .payload - .as_ref() - .and_then(|p| p.as_str()) - .unwrap_or("Unknown error"); - return Err(anyhow::anyhow!("Embedding API error: {}", error_text)); - } - - let embedding_response: EmbeddingResponse = serde_json::from_value( - response - .payload - .ok_or_else(|| anyhow::anyhow!("Empty response body"))?, - )?; - - Ok(embedding_response - .data - .into_iter() - .map(|d| d.embedding) - .collect()) - } -} - #[cfg(test)] mod tests { use super::*;