From ffe204c20991a0943d15905686cbec2fd39e0ed1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=20Houl=C3=A9?= <13155277+tomhoule@users.noreply.github.com> Date: Thu, 20 Aug 2026 17:53:44 +0000 Subject: [PATCH] cloud_llm_client: Add Baseten language model provider (#62950) Release Notes: - N/A --- .../cloud_llm_client/src/cloud_llm_client.rs | 1 + crates/language_model_core/src/provider.rs | 6 +++++ .../src/language_models_cloud.rs | 26 ++++++++++++------- 3 files changed, 23 insertions(+), 10 deletions(-) diff --git a/crates/cloud_llm_client/src/cloud_llm_client.rs b/crates/cloud_llm_client/src/cloud_llm_client.rs index 167507e86a..b4429c0aa0 100644 --- a/crates/cloud_llm_client/src/cloud_llm_client.rs +++ b/crates/cloud_llm_client/src/cloud_llm_client.rs @@ -91,6 +91,7 @@ impl FromStr for UsageLimit { #[strum(serialize_all = "snake_case")] pub enum LanguageModelProvider { Anthropic, + Baseten, OpenAi, Google, XAi, diff --git a/crates/language_model_core/src/provider.rs b/crates/language_model_core/src/provider.rs index 317a0ab280..b32d2592c8 100644 --- a/crates/language_model_core/src/provider.rs +++ b/crates/language_model_core/src/provider.rs @@ -5,6 +5,10 @@ pub const ANTHROPIC_PROVIDER_ID: LanguageModelProviderId = pub const ANTHROPIC_PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("Anthropic"); +pub const BASETEN_PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("baseten"); +pub const BASETEN_PROVIDER_NAME: LanguageModelProviderName = + LanguageModelProviderName::new("Baseten"); + pub const OPEN_AI_PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openai"); pub const OPEN_AI_PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("OpenAI"); @@ -25,6 +29,8 @@ pub fn provider_name_for_id(provider_id: &LanguageModelProviderId) -> LanguageMo OPEN_AI_PROVIDER_NAME } else if provider_id == &ANTHROPIC_PROVIDER_ID { ANTHROPIC_PROVIDER_NAME + } else if provider_id == &BASETEN_PROVIDER_ID { + BASETEN_PROVIDER_NAME } else if provider_id == &GOOGLE_PROVIDER_ID { GOOGLE_PROVIDER_NAME } else if provider_id == &X_AI_PROVIDER_ID { diff --git a/crates/language_models_cloud/src/language_models_cloud.rs b/crates/language_models_cloud/src/language_models_cloud.rs index e75d0b2536..0d00a31ec7 100644 --- a/crates/language_models_cloud/src/language_models_cloud.rs +++ b/crates/language_models_cloud/src/language_models_cloud.rs @@ -19,12 +19,12 @@ use http_client::{ AsyncBody, HttpClient, HttpClientWithUrl, HttpRequestExt, Method, Response, StatusCode, }; use language_model::{ - ANTHROPIC_PROVIDER_ID, ANTHROPIC_PROVIDER_NAME, CompactionResult, DisabledReason, - GOOGLE_PROVIDER_ID, GOOGLE_PROVIDER_NAME, LanguageModel, LanguageModelCompletionError, - LanguageModelCompletionEvent, LanguageModelEffortLevel, LanguageModelId, LanguageModelName, - LanguageModelProviderId, LanguageModelProviderName, LanguageModelRequest, - LanguageModelToolChoice, LanguageModelToolSchemaFormat, OPEN_AI_PROVIDER_ID, - OPEN_AI_PROVIDER_NAME, RateLimiter, X_AI_PROVIDER_ID, X_AI_PROVIDER_NAME, + ANTHROPIC_PROVIDER_ID, ANTHROPIC_PROVIDER_NAME, BASETEN_PROVIDER_ID, BASETEN_PROVIDER_NAME, + CompactionResult, DisabledReason, GOOGLE_PROVIDER_ID, GOOGLE_PROVIDER_NAME, LanguageModel, + LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelEffortLevel, + LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName, + LanguageModelRequest, LanguageModelToolChoice, LanguageModelToolSchemaFormat, + OPEN_AI_PROVIDER_ID, OPEN_AI_PROVIDER_NAME, RateLimiter, X_AI_PROVIDER_ID, X_AI_PROVIDER_NAME, ZED_CLOUD_PROVIDER_ID, ZED_CLOUD_PROVIDER_NAME, }; @@ -504,6 +504,7 @@ impl LanguageModel for CloudLanguageModel ANTHROPIC_PROVIDER_ID, + Baseten => BASETEN_PROVIDER_ID, OpenAi => OPEN_AI_PROVIDER_ID, Google => GOOGLE_PROVIDER_ID, XAi => X_AI_PROVIDER_ID, @@ -514,6 +515,7 @@ impl LanguageModel for CloudLanguageModel ANTHROPIC_PROVIDER_NAME, + Baseten => BASETEN_PROVIDER_NAME, OpenAi => OPEN_AI_PROVIDER_NAME, Google => GOOGLE_PROVIDER_NAME, XAi => X_AI_PROVIDER_NAME, @@ -613,7 +615,8 @@ impl LanguageModel for CloudLanguageModel async { Err(LanguageModelCompletionError::Other(anyhow::anyhow!( "this cloud model does not support explicit compaction" @@ -649,7 +652,7 @@ impl LanguageModel for CloudLanguageModel bool { use cloud_llm_client::LanguageModelProvider::*; - matches!(self.model.provider, OpenAi | XAi) + matches!(self.model.provider, Baseten | OpenAi | XAi) } fn telemetry_id(&self) -> String { @@ -659,6 +662,7 @@ impl LanguageModel for CloudLanguageModel LanguageModelToolSchemaFormat { match self.model.provider { cloud_llm_client::LanguageModelProvider::Anthropic + | cloud_llm_client::LanguageModelProvider::Baseten | cloud_llm_client::LanguageModelProvider::OpenAi => { LanguageModelToolSchemaFormat::JsonSchema } @@ -848,7 +852,8 @@ impl LanguageModel for CloudLanguageModel { + provider @ (cloud_llm_client::LanguageModelProvider::Baseten + | cloud_llm_client::LanguageModelProvider::XAi) => { let http_client = self.http_client.clone(); let token_provider = self.token_provider.clone(); let request = match into_open_ai( @@ -877,7 +882,7 @@ impl LanguageModel for CloudLanguageModel LanguageModelProviderName { match provider { cloud_llm_client::LanguageModelProvider::Anthropic => ANTHROPIC_PROVIDER_NAME, + cloud_llm_client::LanguageModelProvider::Baseten => BASETEN_PROVIDER_NAME, cloud_llm_client::LanguageModelProvider::OpenAi => OPEN_AI_PROVIDER_NAME, cloud_llm_client::LanguageModelProvider::Google => GOOGLE_PROVIDER_NAME, cloud_llm_client::LanguageModelProvider::XAi => X_AI_PROVIDER_NAME,