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

delete embeddings support (#9865)

Browse Source
This commit is contained in:
Jack Amadeo
2026-06-18 15:22:32 -04:00
committed by GitHub
parent a8b11e39b7
commit b3cf3cbcff
6 changed files with 5 additions and 273 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-providers/src/base.rs
-15
View File
@@ -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<String>,
) -> Result<Vec<Vec<f32>>, 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.
crates/goose/src/providers/databricks.rs
+5 -104
View File
@@ -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<Mutex<Option<String>>>,
@@ -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<Self> {
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<Value, ProviderError> {
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<String>,
) -> Result<Vec<Vec<f32>>, ProviderError> {
EmbeddingCapable::create_embeddings(self, session_id, texts)
.await
.map_err(|e| ProviderError::ExecutionError(e.to_string()))
}
async fn fetch_supported_models(&self) -> Result<Vec<String>, 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<String>,
) -> Result<Vec<Vec<f32>>> {
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::<Option<Vec<f32>>>()
.ok_or_else(|| anyhow::anyhow!("Invalid embedding values"))
})
.collect::<Result<Vec<Vec<f32>>>>()?;
Ok(embeddings)
}
}
#[cfg(test)]
mod tests {
use super::*;
crates/goose/src/providers/embedding.rs
-28
View File
@@ -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<String>,
pub model: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingResponse {
pub data: Vec<EmbeddingData>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingData {
pub embedding: Vec<f32>,
}
#[async_trait]
pub trait EmbeddingCapable {
async fn create_embeddings(
&self,
session_id: &str,
texts: Vec<String>,
) -> Result<Vec<Vec<f32>>>;
}
crates/goose/src/providers/litellm.rs
-47
View File
@@ -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<String>,
) -> Result<Vec<Vec<f32>>, 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<f32> = 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 {
crates/goose/src/providers/mod.rs
-1
View File
@@ -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;
crates/goose/src/providers/openai.rs
-78
View File
@@ -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<String>,
) -> Result<Vec<Vec<f32>>, 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<String, String> {
.collect()
}
#[async_trait]
impl EmbeddingCapable for OpenAiProvider {
async fn create_embeddings(
&self,
session_id: &str,
texts: Vec<String>,
) -> Result<Vec<Vec<f32>>> {
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::*;