, current: Point
, contour_count: usize, - base_scale: f32, } impl Path { @@ -702,35 +707,25 @@ impl Path { content_mask: Default::default(), color: Default::default(), contour_count: 0, - base_scale: 1.0, } } - /// Set the base scale of the path. - pub fn scale(mut self, factor: f32) -> Self { - self.base_scale = factor; - self - } - - /// Apply a scale to the path. - pub(crate) fn apply_scale(&self, factor: f32) -> Path { + /// Scale this path by the given factor. + pub fn scale(&self, factor: f32) -> Path { Path { id: self.id, order: self.order, - bounds: self.bounds.scale(self.base_scale * factor), - content_mask: self.content_mask.scale(self.base_scale * factor), + bounds: self.bounds.scale(factor), + content_mask: self.content_mask.scale(factor), vertices: self .vertices .iter() - .map(|vertex| vertex.scale(self.base_scale * factor)) + .map(|vertex| vertex.scale(factor)) .collect(), - start: self - .start - .map(|start| start.scale(self.base_scale * factor)), - current: self.current.scale(self.base_scale * factor), + start: self.start.map(|start| start.scale(factor)), + current: self.current.scale(factor), contour_count: self.contour_count, color: self.color, - base_scale: 1.0, } } @@ -745,7 +740,10 @@ impl Path { pub fn line_to(&mut self, to: Point) { self.contour_count += 1; if self.contour_count > 1 { - self.push_triangle((self.start, self.current, to)); + self.push_triangle( + (self.start, self.current, to), + (point(0., 1.), point(0., 1.), point(0., 1.)), + ); } self.current = to; } @@ -754,15 +752,25 @@ impl Path { pub fn curve_to(&mut self, to: Point, ctrl: Point) { self.contour_count += 1; if self.contour_count > 1 { - self.push_triangle((self.start, self.current, to)); + self.push_triangle( + (self.start, self.current, to), + (point(0., 1.), point(0., 1.), point(0., 1.)), + ); } - self.push_triangle((self.current, ctrl, to)); + self.push_triangle( + (self.current, ctrl, to), + (point(0., 0.), point(0.5, 0.), point(1., 1.)), + ); self.current = to; } /// Push a triangle to the Path. - pub fn push_triangle(&mut self, xy: (Point, Point, Point)) { + pub fn push_triangle( + &mut self, + xy: (Point, Point, Point), + st: (Point, Point, Point), + ) { self.bounds = self .bounds .union(&Bounds { @@ -780,14 +788,17 @@ impl Path { self.vertices.push(PathVertex { xy_position: xy.0, + st_position: st.0, content_mask: Default::default(), }); self.vertices.push(PathVertex { xy_position: xy.1, + st_position: st.1, content_mask: Default::default(), }); self.vertices.push(PathVertex { xy_position: xy.2, + st_position: st.2, content_mask: Default::default(), }); } @@ -803,6 +814,7 @@ impl From> for Primitive { #[repr(C)] pub(crate) struct PathVertex { pub(crate) xy_position: Point, + pub(crate) st_position: Point, pub(crate) content_mask: ContentMask, } @@ -810,6 +822,7 @@ impl PathVertex { pub fn scale(&self, factor: f32) -> PathVertex { PathVertex { xy_position: self.xy_position.scale(factor), + st_position: self.st_position, content_mask: self.content_mask.scale(factor), } } diff --git a/crates/gpui/src/window.rs b/crates/gpui/src/window.rs index 8c01b8afcf..be3b753d6a 100644 --- a/crates/gpui/src/window.rs +++ b/crates/gpui/src/window.rs @@ -2633,7 +2633,7 @@ impl Window { path.color = color.opacity(opacity); self.next_frame .scene - .insert_primitive(path.apply_scale(scale_factor)); + .insert_primitive(path.scale(scale_factor)); } /// Paint an underline into the scene for the next frame at the current z-index. diff --git a/crates/icons/src/icons.rs b/crates/icons/src/icons.rs index 332e38b038..332a8d5791 100644 --- a/crates/icons/src/icons.rs +++ b/crates/icons/src/icons.rs @@ -20,6 +20,7 @@ pub enum IconName { AiOpenAi, AiOpenRouter, AiVZero, + AiXAi, AiZed, ArrowCircle, ArrowDown, diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index 81a0f7d8a1..6bd33fcdf5 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -116,6 +116,12 @@ pub enum LanguageModelCompletionError { provider: LanguageModelProviderName, message: String, }, + #[error("{message}")] + UpstreamProviderError { + message: String, + status: StatusCode, + retry_after: Option, + }, #[error("HTTP response error from {provider}'s API: status {status_code} - {message:?}")] HttpResponseError { provider: LanguageModelProviderName, @@ -178,6 +184,21 @@ pub enum LanguageModelCompletionError { } impl LanguageModelCompletionError { + fn parse_upstream_error_json(message: &str) -> Option<(StatusCode, String)> { + let error_json = serde_json::from_str::(message).ok()?; + let upstream_status = error_json + .get("upstream_status") + .and_then(|v| v.as_u64()) + .and_then(|status| u16::try_from(status).ok()) + .and_then(|status| StatusCode::from_u16(status).ok())?; + let inner_message = error_json + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or(message) + .to_string(); + Some((upstream_status, inner_message)) + } + pub fn from_cloud_failure( upstream_provider: LanguageModelProviderName, code: String, @@ -191,6 +212,18 @@ impl LanguageModelCompletionError { Self::PromptTooLarge { tokens: Some(tokens), } + } else if code == "upstream_http_error" { + if let Some((upstream_status, inner_message)) = + Self::parse_upstream_error_json(&message) + { + return Self::from_http_status( + upstream_provider, + upstream_status, + inner_message, + retry_after, + ); + } + anyhow!("completion request failed, code: {code}, message: {message}").into() } else if let Some(status_code) = code .strip_prefix("upstream_http_") .and_then(|code| StatusCode::from_str(code).ok()) @@ -701,3 +734,104 @@ impl From for LanguageModelProviderName { Self(SharedString::from(value)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_from_cloud_failure_with_upstream_http_error() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":503}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!( + "Expected ServerOverloaded error for 503 status, got: {:?}", + error + ), + } + + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Internal server error","upstream_status":500}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider.0, "anthropic"); + assert_eq!(message, "Internal server error"); + } + _ => panic!( + "Expected ApiInternalServerError for 500 status, got: {:?}", + error + ), + } + } + + #[test] + fn test_from_cloud_failure_with_standard_format() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_503".to_string(), + "Service unavailable".to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!("Expected ServerOverloaded error for upstream_http_503"), + } + } + + #[test] + fn test_upstream_http_error_connection_timeout() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":503}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!( + "Expected ServerOverloaded error for connection timeout with 503 status, got: {:?}", + error + ), + } + + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":500}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider.0, "anthropic"); + assert_eq!( + message, + "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout" + ); + } + _ => panic!( + "Expected ApiInternalServerError for connection timeout with 500 status, got: {:?}", + error + ), + } + } +} diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 514443ddec..e928df8a74 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -44,6 +44,7 @@ ollama = { workspace = true, features = ["schemars"] } open_ai = { workspace = true, features = ["schemars"] } open_router = { workspace = true, features = ["schemars"] } vercel = { workspace = true, features = ["schemars"] } +x_ai = { workspace = true, features = ["schemars"] } partial-json-fixer.workspace = true proto.workspace = true release_channel.workspace = true diff --git a/crates/language_models/src/language_models.rs b/crates/language_models/src/language_models.rs index c7324732c9..192f5a5fae 100644 --- a/crates/language_models/src/language_models.rs +++ b/crates/language_models/src/language_models.rs @@ -20,6 +20,7 @@ use crate::provider::ollama::OllamaLanguageModelProvider; use crate::provider::open_ai::OpenAiLanguageModelProvider; use crate::provider::open_router::OpenRouterLanguageModelProvider; use crate::provider::vercel::VercelLanguageModelProvider; +use crate::provider::x_ai::XAiLanguageModelProvider; pub use crate::settings::*; pub fn init(user_store: Entity, client: Arc, cx: &mut App) { @@ -81,5 +82,6 @@ fn register_language_model_providers( VercelLanguageModelProvider::new(client.http_client(), cx), cx, ); + registry.register_provider(XAiLanguageModelProvider::new(client.http_client(), cx), cx); registry.register_provider(CopilotChatLanguageModelProvider::new(cx), cx); } diff --git a/crates/language_models/src/provider.rs b/crates/language_models/src/provider.rs index 6bc93bd366..c717be7c90 100644 --- a/crates/language_models/src/provider.rs +++ b/crates/language_models/src/provider.rs @@ -10,3 +10,4 @@ pub mod ollama; pub mod open_ai; pub mod open_router; pub mod vercel; +pub mod x_ai; diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs index 9b7fee228a..8d25af1a49 100644 --- a/crates/language_models/src/provider/cloud.rs +++ b/crates/language_models/src/provider/cloud.rs @@ -166,46 +166,9 @@ impl State { } let response = Self::fetch_models(client, llm_api_token, use_cloud).await?; - cx.update(|cx| { - this.update(cx, |this, cx| { - let mut models = Vec::new(); - - for model in response.models { - models.push(Arc::new(model.clone())); - - // Right now we represent thinking variants of models as separate models on the client, - // so we need to insert variants for any model that supports thinking. - if model.supports_thinking { - models.push(Arc::new(zed_llm_client::LanguageModel { - id: zed_llm_client::LanguageModelId( - format!("{}-thinking", model.id).into(), - ), - display_name: format!("{} Thinking", model.display_name), - ..model - })); - } - } - - this.default_model = models - .iter() - .find(|model| model.id == response.default_model) - .cloned(); - this.default_fast_model = models - .iter() - .find(|model| model.id == response.default_fast_model) - .cloned(); - this.recommended_models = response - .recommended_models - .iter() - .filter_map(|id| models.iter().find(|model| &model.id == id)) - .cloned() - .collect(); - this.models = models; - cx.notify(); - }) - })??; - - anyhow::Ok(()) + this.update(cx, |this, cx| { + this.update_models(response, cx); + }) }) .await .context("failed to fetch Zed models") @@ -216,12 +179,15 @@ impl State { }), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, - |this, _listener, _event, cx| { + move |this, _listener, _event, cx| { let client = this.client.clone(); let llm_api_token = this.llm_api_token.clone(); - cx.spawn(async move |_this, _cx| { + cx.spawn(async move |this, cx| { llm_api_token.refresh(&client).await?; - anyhow::Ok(()) + let response = Self::fetch_models(client, llm_api_token, use_cloud).await?; + this.update(cx, |this, cx| { + this.update_models(response, cx); + }) }) .detach_and_log_err(cx); }, @@ -264,6 +230,41 @@ impl State { })); } + fn update_models(&mut self, response: ListModelsResponse, cx: &mut Context) { + let mut models = Vec::new(); + + for model in response.models { + models.push(Arc::new(model.clone())); + + // Right now we represent thinking variants of models as separate models on the client, + // so we need to insert variants for any model that supports thinking. + if model.supports_thinking { + models.push(Arc::new(zed_llm_client::LanguageModel { + id: zed_llm_client::LanguageModelId(format!("{}-thinking", model.id).into()), + display_name: format!("{} Thinking", model.display_name), + ..model + })); + } + } + + self.default_model = models + .iter() + .find(|model| model.id == response.default_model) + .cloned(); + self.default_fast_model = models + .iter() + .find(|model| model.id == response.default_fast_model) + .cloned(); + self.recommended_models = response + .recommended_models + .iter() + .filter_map(|id| models.iter().find(|model| &model.id == id)) + .cloned() + .collect(); + self.models = models; + cx.notify(); + } + async fn fetch_models( client: Arc, llm_api_token: LlmApiToken, @@ -653,8 +654,62 @@ struct ApiError { headers: HeaderMap, } +/// Represents error responses from Zed's cloud API. +/// +/// Example JSON for an upstream HTTP error: +/// ```json +/// { +/// "code": "upstream_http_error", +/// "message": "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout", +/// "upstream_status": 503 +/// } +/// ``` +#[derive(Debug, serde::Deserialize)] +struct CloudApiError { + code: String, + message: String, + #[serde(default)] + #[serde(deserialize_with = "deserialize_optional_status_code")] + upstream_status: Option, + #[serde(default)] + retry_after: Option, +} + +fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let opt: Option = Option::deserialize(deserializer)?; + Ok(opt.and_then(|code| StatusCode::from_u16(code).ok())) +} + impl From for LanguageModelCompletionError { fn from(error: ApiError) -> Self { + if let Ok(cloud_error) = serde_json::from_str::(&error.body) { + if cloud_error.code.starts_with("upstream_http_") { + let status = if let Some(status) = cloud_error.upstream_status { + status + } else if cloud_error.code.ends_with("_error") { + error.status + } else { + // If there's a status code in the code string (e.g. "upstream_http_429") + // then use that; otherwise, see if the JSON contains a status code. + cloud_error + .code + .strip_prefix("upstream_http_") + .and_then(|code_str| code_str.parse::().ok()) + .and_then(|code| StatusCode::from_u16(code).ok()) + .unwrap_or(error.status) + }; + + return LanguageModelCompletionError::UpstreamProviderError { + message: cloud_error.message, + status, + retry_after: cloud_error.retry_after.map(Duration::from_secs_f64), + }; + } + } + let retry_after = None; LanguageModelCompletionError::from_http_status( PROVIDER_NAME, @@ -1294,3 +1349,155 @@ impl Component for ZedAiConfiguration { ) } } + +#[cfg(test)] +mod tests { + use super::*; + use http_client::http::{HeaderMap, StatusCode}; + use language_model::LanguageModelCompletionError; + + #[test] + fn test_api_error_conversion_with_upstream_http_error() { + // upstream_http_error with 503 status should become ServerOverloaded + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout","upstream_status":503}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 503, got: {:?}", + completion_error + ), + } + + // upstream_http_error with 500 status should become ApiInternalServerError + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the OpenAI API: internal server error","upstream_status":500}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the OpenAI API: internal server error" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 500, got: {:?}", + completion_error + ), + } + + // upstream_http_error with 429 status should become RateLimitExceeded + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Google API: rate limit exceeded","upstream_status":429}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the Google API: rate limit exceeded" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 429, got: {:?}", + completion_error + ), + } + + // Regular 500 error without upstream_http_error should remain ApiInternalServerError for Zed + let error_body = "Regular internal server error"; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider, PROVIDER_NAME); + assert_eq!(message, "Regular internal server error"); + } + _ => panic!( + "Expected ApiInternalServerError for regular 500, got: {:?}", + completion_error + ), + } + + // upstream_http_429 format should be converted to UpstreamProviderError + let error_body = r#"{"code":"upstream_http_429","message":"Upstream Anthropic rate limit exceeded.","retry_after":30.5}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { + message, + status, + retry_after, + } => { + assert_eq!(message, "Upstream Anthropic rate limit exceeded."); + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + assert_eq!(retry_after, Some(Duration::from_secs_f64(30.5))); + } + _ => panic!( + "Expected UpstreamProviderError for upstream_http_429, got: {:?}", + completion_error + ), + } + + // Invalid JSON in error body should fall back to regular error handling + let error_body = "Not JSON at all"; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::ApiInternalServerError { provider, .. } => { + assert_eq!(provider, PROVIDER_NAME); + } + _ => panic!( + "Expected ApiInternalServerError for invalid JSON, got: {:?}", + completion_error + ), + } + } +} diff --git a/crates/language_models/src/provider/open_router.rs b/crates/language_models/src/provider/open_router.rs index 5883da1e2f..90d375a6b2 100644 --- a/crates/language_models/src/provider/open_router.rs +++ b/crates/language_models/src/provider/open_router.rs @@ -376,7 +376,7 @@ impl LanguageModel for OpenRouterLanguageModel { fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { let model_id = self.model.id().trim().to_lowercase(); - if model_id.contains("gemini") { + if model_id.contains("gemini") || model_id.contains("grok-4") { LanguageModelToolSchemaFormat::JsonSchemaSubset } else { LanguageModelToolSchemaFormat::JsonSchema diff --git a/crates/language_models/src/provider/x_ai.rs b/crates/language_models/src/provider/x_ai.rs new file mode 100644 index 0000000000..5f6034571b --- /dev/null +++ b/crates/language_models/src/provider/x_ai.rs @@ -0,0 +1,571 @@ +use anyhow::{Context as _, Result, anyhow}; +use collections::BTreeMap; +use credentials_provider::CredentialsProvider; +use futures::{FutureExt, StreamExt, future::BoxFuture}; +use gpui::{AnyView, App, AsyncApp, Context, Entity, Subscription, Task, Window}; +use http_client::HttpClient; +use language_model::{ + AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, + LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, + LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, + LanguageModelToolChoice, LanguageModelToolSchemaFormat, RateLimiter, Role, +}; +use menu; +use open_ai::ResponseStreamEvent; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use settings::{Settings, SettingsStore}; +use std::sync::Arc; +use strum::IntoEnumIterator; +use x_ai::Model; + +use ui::{ElevationIndex, List, Tooltip, prelude::*}; +use ui_input::SingleLineInput; +use util::ResultExt; + +use crate::{AllLanguageModelSettings, ui::InstructionListItem}; + +const PROVIDER_ID: &str = "x_ai"; +const PROVIDER_NAME: &str = "xAI"; + +#[derive(Default, Clone, Debug, PartialEq)] +pub struct XAiSettings { + pub api_url: String, + pub available_models: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct AvailableModel { + pub name: String, + pub display_name: Option, + pub max_tokens: u64, + pub max_output_tokens: Option, + pub max_completion_tokens: Option, +} + +pub struct XAiLanguageModelProvider { + http_client: Arc, + state: gpui::Entity, +} + +pub struct State { + api_key: Option, + api_key_from_env: bool, + _subscription: Subscription, +} + +const XAI_API_KEY_VAR: &str = "XAI_API_KEY"; + +impl State { + fn is_authenticated(&self) -> bool { + self.api_key.is_some() + } + + fn reset_api_key(&self, cx: &mut Context) -> Task> { + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + credentials_provider + .delete_credentials(&api_url, &cx) + .await + .log_err(); + this.update(cx, |this, cx| { + this.api_key = None; + this.api_key_from_env = false; + cx.notify(); + }) + }) + } + + fn set_api_key(&mut self, api_key: String, cx: &mut Context) -> Task> { + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + credentials_provider + .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx) + .await + .log_err(); + this.update(cx, |this, cx| { + this.api_key = Some(api_key); + cx.notify(); + }) + }) + } + + fn authenticate(&self, cx: &mut Context) -> Task> { + if self.is_authenticated() { + return Task::ready(Ok(())); + } + + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + let (api_key, from_env) = if let Ok(api_key) = std::env::var(XAI_API_KEY_VAR) { + (api_key, true) + } else { + let (_, api_key) = credentials_provider + .read_credentials(&api_url, &cx) + .await? + .ok_or(AuthenticateError::CredentialsNotFound)?; + ( + String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?, + false, + ) + }; + this.update(cx, |this, cx| { + this.api_key = Some(api_key); + this.api_key_from_env = from_env; + cx.notify(); + })?; + + Ok(()) + }) + } +} + +impl XAiLanguageModelProvider { + pub fn new(http_client: Arc, cx: &mut App) -> Self { + let state = cx.new(|cx| State { + api_key: None, + api_key_from_env: false, + _subscription: cx.observe_global::(|_this: &mut State, cx| { + cx.notify(); + }), + }); + + Self { http_client, state } + } + + fn create_language_model(&self, model: x_ai::Model) -> Arc { + Arc::new(XAiLanguageModel { + id: LanguageModelId::from(model.id().to_string()), + model, + state: self.state.clone(), + http_client: self.http_client.clone(), + request_limiter: RateLimiter::new(4), + }) + } +} + +impl LanguageModelProviderState for XAiLanguageModelProvider { + type ObservableEntity = State; + + fn observable_entity(&self) -> Option> { + Some(self.state.clone()) + } +} + +impl LanguageModelProvider for XAiLanguageModelProvider { + fn id(&self) -> LanguageModelProviderId { + LanguageModelProviderId(PROVIDER_ID.into()) + } + + fn name(&self) -> LanguageModelProviderName { + LanguageModelProviderName(PROVIDER_NAME.into()) + } + + fn icon(&self) -> IconName { + IconName::AiXAi + } + + fn default_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(x_ai::Model::default())) + } + + fn default_fast_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(x_ai::Model::default_fast())) + } + + fn provided_models(&self, cx: &App) -> Vec> { + let mut models = BTreeMap::default(); + + for model in x_ai::Model::iter() { + if !matches!(model, x_ai::Model::Custom { .. }) { + models.insert(model.id().to_string(), model); + } + } + + for model in &AllLanguageModelSettings::get_global(cx) + .x_ai + .available_models + { + models.insert( + model.name.clone(), + x_ai::Model::Custom { + name: model.name.clone(), + display_name: model.display_name.clone(), + max_tokens: model.max_tokens, + max_output_tokens: model.max_output_tokens, + max_completion_tokens: model.max_completion_tokens, + }, + ); + } + + models + .into_values() + .map(|model| self.create_language_model(model)) + .collect() + } + + fn is_authenticated(&self, cx: &App) -> bool { + self.state.read(cx).is_authenticated() + } + + fn authenticate(&self, cx: &mut App) -> Task
, + pub(crate) st_position: Point, pub(crate) content_mask: ContentMask, } @@ -810,6 +822,7 @@ impl PathVertex { pub fn scale(&self, factor: f32) -> PathVertex { PathVertex { xy_position: self.xy_position.scale(factor), + st_position: self.st_position, content_mask: self.content_mask.scale(factor), } } diff --git a/crates/gpui/src/window.rs b/crates/gpui/src/window.rs index 8c01b8afcf..be3b753d6a 100644 --- a/crates/gpui/src/window.rs +++ b/crates/gpui/src/window.rs @@ -2633,7 +2633,7 @@ impl Window { path.color = color.opacity(opacity); self.next_frame .scene - .insert_primitive(path.apply_scale(scale_factor)); + .insert_primitive(path.scale(scale_factor)); } /// Paint an underline into the scene for the next frame at the current z-index. diff --git a/crates/icons/src/icons.rs b/crates/icons/src/icons.rs index 332e38b038..332a8d5791 100644 --- a/crates/icons/src/icons.rs +++ b/crates/icons/src/icons.rs @@ -20,6 +20,7 @@ pub enum IconName { AiOpenAi, AiOpenRouter, AiVZero, + AiXAi, AiZed, ArrowCircle, ArrowDown, diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index 81a0f7d8a1..6bd33fcdf5 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -116,6 +116,12 @@ pub enum LanguageModelCompletionError { provider: LanguageModelProviderName, message: String, }, + #[error("{message}")] + UpstreamProviderError { + message: String, + status: StatusCode, + retry_after: Option, + }, #[error("HTTP response error from {provider}'s API: status {status_code} - {message:?}")] HttpResponseError { provider: LanguageModelProviderName, @@ -178,6 +184,21 @@ pub enum LanguageModelCompletionError { } impl LanguageModelCompletionError { + fn parse_upstream_error_json(message: &str) -> Option<(StatusCode, String)> { + let error_json = serde_json::from_str::(message).ok()?; + let upstream_status = error_json + .get("upstream_status") + .and_then(|v| v.as_u64()) + .and_then(|status| u16::try_from(status).ok()) + .and_then(|status| StatusCode::from_u16(status).ok())?; + let inner_message = error_json + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or(message) + .to_string(); + Some((upstream_status, inner_message)) + } + pub fn from_cloud_failure( upstream_provider: LanguageModelProviderName, code: String, @@ -191,6 +212,18 @@ impl LanguageModelCompletionError { Self::PromptTooLarge { tokens: Some(tokens), } + } else if code == "upstream_http_error" { + if let Some((upstream_status, inner_message)) = + Self::parse_upstream_error_json(&message) + { + return Self::from_http_status( + upstream_provider, + upstream_status, + inner_message, + retry_after, + ); + } + anyhow!("completion request failed, code: {code}, message: {message}").into() } else if let Some(status_code) = code .strip_prefix("upstream_http_") .and_then(|code| StatusCode::from_str(code).ok()) @@ -701,3 +734,104 @@ impl From for LanguageModelProviderName { Self(SharedString::from(value)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_from_cloud_failure_with_upstream_http_error() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":503}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!( + "Expected ServerOverloaded error for 503 status, got: {:?}", + error + ), + } + + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Internal server error","upstream_status":500}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider.0, "anthropic"); + assert_eq!(message, "Internal server error"); + } + _ => panic!( + "Expected ApiInternalServerError for 500 status, got: {:?}", + error + ), + } + } + + #[test] + fn test_from_cloud_failure_with_standard_format() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_503".to_string(), + "Service unavailable".to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!("Expected ServerOverloaded error for upstream_http_503"), + } + } + + #[test] + fn test_upstream_http_error_connection_timeout() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":503}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!( + "Expected ServerOverloaded error for connection timeout with 503 status, got: {:?}", + error + ), + } + + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":500}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider.0, "anthropic"); + assert_eq!( + message, + "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout" + ); + } + _ => panic!( + "Expected ApiInternalServerError for connection timeout with 500 status, got: {:?}", + error + ), + } + } +} diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 514443ddec..e928df8a74 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -44,6 +44,7 @@ ollama = { workspace = true, features = ["schemars"] } open_ai = { workspace = true, features = ["schemars"] } open_router = { workspace = true, features = ["schemars"] } vercel = { workspace = true, features = ["schemars"] } +x_ai = { workspace = true, features = ["schemars"] } partial-json-fixer.workspace = true proto.workspace = true release_channel.workspace = true diff --git a/crates/language_models/src/language_models.rs b/crates/language_models/src/language_models.rs index c7324732c9..192f5a5fae 100644 --- a/crates/language_models/src/language_models.rs +++ b/crates/language_models/src/language_models.rs @@ -20,6 +20,7 @@ use crate::provider::ollama::OllamaLanguageModelProvider; use crate::provider::open_ai::OpenAiLanguageModelProvider; use crate::provider::open_router::OpenRouterLanguageModelProvider; use crate::provider::vercel::VercelLanguageModelProvider; +use crate::provider::x_ai::XAiLanguageModelProvider; pub use crate::settings::*; pub fn init(user_store: Entity, client: Arc, cx: &mut App) { @@ -81,5 +82,6 @@ fn register_language_model_providers( VercelLanguageModelProvider::new(client.http_client(), cx), cx, ); + registry.register_provider(XAiLanguageModelProvider::new(client.http_client(), cx), cx); registry.register_provider(CopilotChatLanguageModelProvider::new(cx), cx); } diff --git a/crates/language_models/src/provider.rs b/crates/language_models/src/provider.rs index 6bc93bd366..c717be7c90 100644 --- a/crates/language_models/src/provider.rs +++ b/crates/language_models/src/provider.rs @@ -10,3 +10,4 @@ pub mod ollama; pub mod open_ai; pub mod open_router; pub mod vercel; +pub mod x_ai; diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs index 9b7fee228a..8d25af1a49 100644 --- a/crates/language_models/src/provider/cloud.rs +++ b/crates/language_models/src/provider/cloud.rs @@ -166,46 +166,9 @@ impl State { } let response = Self::fetch_models(client, llm_api_token, use_cloud).await?; - cx.update(|cx| { - this.update(cx, |this, cx| { - let mut models = Vec::new(); - - for model in response.models { - models.push(Arc::new(model.clone())); - - // Right now we represent thinking variants of models as separate models on the client, - // so we need to insert variants for any model that supports thinking. - if model.supports_thinking { - models.push(Arc::new(zed_llm_client::LanguageModel { - id: zed_llm_client::LanguageModelId( - format!("{}-thinking", model.id).into(), - ), - display_name: format!("{} Thinking", model.display_name), - ..model - })); - } - } - - this.default_model = models - .iter() - .find(|model| model.id == response.default_model) - .cloned(); - this.default_fast_model = models - .iter() - .find(|model| model.id == response.default_fast_model) - .cloned(); - this.recommended_models = response - .recommended_models - .iter() - .filter_map(|id| models.iter().find(|model| &model.id == id)) - .cloned() - .collect(); - this.models = models; - cx.notify(); - }) - })??; - - anyhow::Ok(()) + this.update(cx, |this, cx| { + this.update_models(response, cx); + }) }) .await .context("failed to fetch Zed models") @@ -216,12 +179,15 @@ impl State { }), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, - |this, _listener, _event, cx| { + move |this, _listener, _event, cx| { let client = this.client.clone(); let llm_api_token = this.llm_api_token.clone(); - cx.spawn(async move |_this, _cx| { + cx.spawn(async move |this, cx| { llm_api_token.refresh(&client).await?; - anyhow::Ok(()) + let response = Self::fetch_models(client, llm_api_token, use_cloud).await?; + this.update(cx, |this, cx| { + this.update_models(response, cx); + }) }) .detach_and_log_err(cx); }, @@ -264,6 +230,41 @@ impl State { })); } + fn update_models(&mut self, response: ListModelsResponse, cx: &mut Context) { + let mut models = Vec::new(); + + for model in response.models { + models.push(Arc::new(model.clone())); + + // Right now we represent thinking variants of models as separate models on the client, + // so we need to insert variants for any model that supports thinking. + if model.supports_thinking { + models.push(Arc::new(zed_llm_client::LanguageModel { + id: zed_llm_client::LanguageModelId(format!("{}-thinking", model.id).into()), + display_name: format!("{} Thinking", model.display_name), + ..model + })); + } + } + + self.default_model = models + .iter() + .find(|model| model.id == response.default_model) + .cloned(); + self.default_fast_model = models + .iter() + .find(|model| model.id == response.default_fast_model) + .cloned(); + self.recommended_models = response + .recommended_models + .iter() + .filter_map(|id| models.iter().find(|model| &model.id == id)) + .cloned() + .collect(); + self.models = models; + cx.notify(); + } + async fn fetch_models( client: Arc, llm_api_token: LlmApiToken, @@ -653,8 +654,62 @@ struct ApiError { headers: HeaderMap, } +/// Represents error responses from Zed's cloud API. +/// +/// Example JSON for an upstream HTTP error: +/// ```json +/// { +/// "code": "upstream_http_error", +/// "message": "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout", +/// "upstream_status": 503 +/// } +/// ``` +#[derive(Debug, serde::Deserialize)] +struct CloudApiError { + code: String, + message: String, + #[serde(default)] + #[serde(deserialize_with = "deserialize_optional_status_code")] + upstream_status: Option, + #[serde(default)] + retry_after: Option, +} + +fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let opt: Option = Option::deserialize(deserializer)?; + Ok(opt.and_then(|code| StatusCode::from_u16(code).ok())) +} + impl From for LanguageModelCompletionError { fn from(error: ApiError) -> Self { + if let Ok(cloud_error) = serde_json::from_str::(&error.body) { + if cloud_error.code.starts_with("upstream_http_") { + let status = if let Some(status) = cloud_error.upstream_status { + status + } else if cloud_error.code.ends_with("_error") { + error.status + } else { + // If there's a status code in the code string (e.g. "upstream_http_429") + // then use that; otherwise, see if the JSON contains a status code. + cloud_error + .code + .strip_prefix("upstream_http_") + .and_then(|code_str| code_str.parse::().ok()) + .and_then(|code| StatusCode::from_u16(code).ok()) + .unwrap_or(error.status) + }; + + return LanguageModelCompletionError::UpstreamProviderError { + message: cloud_error.message, + status, + retry_after: cloud_error.retry_after.map(Duration::from_secs_f64), + }; + } + } + let retry_after = None; LanguageModelCompletionError::from_http_status( PROVIDER_NAME, @@ -1294,3 +1349,155 @@ impl Component for ZedAiConfiguration { ) } } + +#[cfg(test)] +mod tests { + use super::*; + use http_client::http::{HeaderMap, StatusCode}; + use language_model::LanguageModelCompletionError; + + #[test] + fn test_api_error_conversion_with_upstream_http_error() { + // upstream_http_error with 503 status should become ServerOverloaded + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout","upstream_status":503}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 503, got: {:?}", + completion_error + ), + } + + // upstream_http_error with 500 status should become ApiInternalServerError + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the OpenAI API: internal server error","upstream_status":500}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the OpenAI API: internal server error" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 500, got: {:?}", + completion_error + ), + } + + // upstream_http_error with 429 status should become RateLimitExceeded + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Google API: rate limit exceeded","upstream_status":429}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the Google API: rate limit exceeded" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 429, got: {:?}", + completion_error + ), + } + + // Regular 500 error without upstream_http_error should remain ApiInternalServerError for Zed + let error_body = "Regular internal server error"; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider, PROVIDER_NAME); + assert_eq!(message, "Regular internal server error"); + } + _ => panic!( + "Expected ApiInternalServerError for regular 500, got: {:?}", + completion_error + ), + } + + // upstream_http_429 format should be converted to UpstreamProviderError + let error_body = r#"{"code":"upstream_http_429","message":"Upstream Anthropic rate limit exceeded.","retry_after":30.5}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { + message, + status, + retry_after, + } => { + assert_eq!(message, "Upstream Anthropic rate limit exceeded."); + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + assert_eq!(retry_after, Some(Duration::from_secs_f64(30.5))); + } + _ => panic!( + "Expected UpstreamProviderError for upstream_http_429, got: {:?}", + completion_error + ), + } + + // Invalid JSON in error body should fall back to regular error handling + let error_body = "Not JSON at all"; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::ApiInternalServerError { provider, .. } => { + assert_eq!(provider, PROVIDER_NAME); + } + _ => panic!( + "Expected ApiInternalServerError for invalid JSON, got: {:?}", + completion_error + ), + } + } +} diff --git a/crates/language_models/src/provider/open_router.rs b/crates/language_models/src/provider/open_router.rs index 5883da1e2f..90d375a6b2 100644 --- a/crates/language_models/src/provider/open_router.rs +++ b/crates/language_models/src/provider/open_router.rs @@ -376,7 +376,7 @@ impl LanguageModel for OpenRouterLanguageModel { fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { let model_id = self.model.id().trim().to_lowercase(); - if model_id.contains("gemini") { + if model_id.contains("gemini") || model_id.contains("grok-4") { LanguageModelToolSchemaFormat::JsonSchemaSubset } else { LanguageModelToolSchemaFormat::JsonSchema diff --git a/crates/language_models/src/provider/x_ai.rs b/crates/language_models/src/provider/x_ai.rs new file mode 100644 index 0000000000..5f6034571b --- /dev/null +++ b/crates/language_models/src/provider/x_ai.rs @@ -0,0 +1,571 @@ +use anyhow::{Context as _, Result, anyhow}; +use collections::BTreeMap; +use credentials_provider::CredentialsProvider; +use futures::{FutureExt, StreamExt, future::BoxFuture}; +use gpui::{AnyView, App, AsyncApp, Context, Entity, Subscription, Task, Window}; +use http_client::HttpClient; +use language_model::{ + AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, + LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, + LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, + LanguageModelToolChoice, LanguageModelToolSchemaFormat, RateLimiter, Role, +}; +use menu; +use open_ai::ResponseStreamEvent; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use settings::{Settings, SettingsStore}; +use std::sync::Arc; +use strum::IntoEnumIterator; +use x_ai::Model; + +use ui::{ElevationIndex, List, Tooltip, prelude::*}; +use ui_input::SingleLineInput; +use util::ResultExt; + +use crate::{AllLanguageModelSettings, ui::InstructionListItem}; + +const PROVIDER_ID: &str = "x_ai"; +const PROVIDER_NAME: &str = "xAI"; + +#[derive(Default, Clone, Debug, PartialEq)] +pub struct XAiSettings { + pub api_url: String, + pub available_models: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct AvailableModel { + pub name: String, + pub display_name: Option, + pub max_tokens: u64, + pub max_output_tokens: Option, + pub max_completion_tokens: Option, +} + +pub struct XAiLanguageModelProvider { + http_client: Arc, + state: gpui::Entity, +} + +pub struct State { + api_key: Option, + api_key_from_env: bool, + _subscription: Subscription, +} + +const XAI_API_KEY_VAR: &str = "XAI_API_KEY"; + +impl State { + fn is_authenticated(&self) -> bool { + self.api_key.is_some() + } + + fn reset_api_key(&self, cx: &mut Context) -> Task> { + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + credentials_provider + .delete_credentials(&api_url, &cx) + .await + .log_err(); + this.update(cx, |this, cx| { + this.api_key = None; + this.api_key_from_env = false; + cx.notify(); + }) + }) + } + + fn set_api_key(&mut self, api_key: String, cx: &mut Context) -> Task> { + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + credentials_provider + .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx) + .await + .log_err(); + this.update(cx, |this, cx| { + this.api_key = Some(api_key); + cx.notify(); + }) + }) + } + + fn authenticate(&self, cx: &mut Context) -> Task> { + if self.is_authenticated() { + return Task::ready(Ok(())); + } + + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + let (api_key, from_env) = if let Ok(api_key) = std::env::var(XAI_API_KEY_VAR) { + (api_key, true) + } else { + let (_, api_key) = credentials_provider + .read_credentials(&api_url, &cx) + .await? + .ok_or(AuthenticateError::CredentialsNotFound)?; + ( + String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?, + false, + ) + }; + this.update(cx, |this, cx| { + this.api_key = Some(api_key); + this.api_key_from_env = from_env; + cx.notify(); + })?; + + Ok(()) + }) + } +} + +impl XAiLanguageModelProvider { + pub fn new(http_client: Arc, cx: &mut App) -> Self { + let state = cx.new(|cx| State { + api_key: None, + api_key_from_env: false, + _subscription: cx.observe_global::(|_this: &mut State, cx| { + cx.notify(); + }), + }); + + Self { http_client, state } + } + + fn create_language_model(&self, model: x_ai::Model) -> Arc { + Arc::new(XAiLanguageModel { + id: LanguageModelId::from(model.id().to_string()), + model, + state: self.state.clone(), + http_client: self.http_client.clone(), + request_limiter: RateLimiter::new(4), + }) + } +} + +impl LanguageModelProviderState for XAiLanguageModelProvider { + type ObservableEntity = State; + + fn observable_entity(&self) -> Option> { + Some(self.state.clone()) + } +} + +impl LanguageModelProvider for XAiLanguageModelProvider { + fn id(&self) -> LanguageModelProviderId { + LanguageModelProviderId(PROVIDER_ID.into()) + } + + fn name(&self) -> LanguageModelProviderName { + LanguageModelProviderName(PROVIDER_NAME.into()) + } + + fn icon(&self) -> IconName { + IconName::AiXAi + } + + fn default_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(x_ai::Model::default())) + } + + fn default_fast_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(x_ai::Model::default_fast())) + } + + fn provided_models(&self, cx: &App) -> Vec> { + let mut models = BTreeMap::default(); + + for model in x_ai::Model::iter() { + if !matches!(model, x_ai::Model::Custom { .. }) { + models.insert(model.id().to_string(), model); + } + } + + for model in &AllLanguageModelSettings::get_global(cx) + .x_ai + .available_models + { + models.insert( + model.name.clone(), + x_ai::Model::Custom { + name: model.name.clone(), + display_name: model.display_name.clone(), + max_tokens: model.max_tokens, + max_output_tokens: model.max_output_tokens, + max_completion_tokens: model.max_completion_tokens, + }, + ); + } + + models + .into_values() + .map(|model| self.create_language_model(model)) + .collect() + } + + fn is_authenticated(&self, cx: &App) -> bool { + self.state.read(cx).is_authenticated() + } + + fn authenticate(&self, cx: &mut App) -> Task
, } @@ -810,6 +822,7 @@ impl PathVertex { pub fn scale(&self, factor: f32) -> PathVertex { PathVertex { xy_position: self.xy_position.scale(factor), + st_position: self.st_position, content_mask: self.content_mask.scale(factor), } } diff --git a/crates/gpui/src/window.rs b/crates/gpui/src/window.rs index 8c01b8afcf..be3b753d6a 100644 --- a/crates/gpui/src/window.rs +++ b/crates/gpui/src/window.rs @@ -2633,7 +2633,7 @@ impl Window { path.color = color.opacity(opacity); self.next_frame .scene - .insert_primitive(path.apply_scale(scale_factor)); + .insert_primitive(path.scale(scale_factor)); } /// Paint an underline into the scene for the next frame at the current z-index. diff --git a/crates/icons/src/icons.rs b/crates/icons/src/icons.rs index 332e38b038..332a8d5791 100644 --- a/crates/icons/src/icons.rs +++ b/crates/icons/src/icons.rs @@ -20,6 +20,7 @@ pub enum IconName { AiOpenAi, AiOpenRouter, AiVZero, + AiXAi, AiZed, ArrowCircle, ArrowDown, diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index 81a0f7d8a1..6bd33fcdf5 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -116,6 +116,12 @@ pub enum LanguageModelCompletionError { provider: LanguageModelProviderName, message: String, }, + #[error("{message}")] + UpstreamProviderError { + message: String, + status: StatusCode, + retry_after: Option, + }, #[error("HTTP response error from {provider}'s API: status {status_code} - {message:?}")] HttpResponseError { provider: LanguageModelProviderName, @@ -178,6 +184,21 @@ pub enum LanguageModelCompletionError { } impl LanguageModelCompletionError { + fn parse_upstream_error_json(message: &str) -> Option<(StatusCode, String)> { + let error_json = serde_json::from_str::(message).ok()?; + let upstream_status = error_json + .get("upstream_status") + .and_then(|v| v.as_u64()) + .and_then(|status| u16::try_from(status).ok()) + .and_then(|status| StatusCode::from_u16(status).ok())?; + let inner_message = error_json + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or(message) + .to_string(); + Some((upstream_status, inner_message)) + } + pub fn from_cloud_failure( upstream_provider: LanguageModelProviderName, code: String, @@ -191,6 +212,18 @@ impl LanguageModelCompletionError { Self::PromptTooLarge { tokens: Some(tokens), } + } else if code == "upstream_http_error" { + if let Some((upstream_status, inner_message)) = + Self::parse_upstream_error_json(&message) + { + return Self::from_http_status( + upstream_provider, + upstream_status, + inner_message, + retry_after, + ); + } + anyhow!("completion request failed, code: {code}, message: {message}").into() } else if let Some(status_code) = code .strip_prefix("upstream_http_") .and_then(|code| StatusCode::from_str(code).ok()) @@ -701,3 +734,104 @@ impl From for LanguageModelProviderName { Self(SharedString::from(value)) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_from_cloud_failure_with_upstream_http_error() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":503}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!( + "Expected ServerOverloaded error for 503 status, got: {:?}", + error + ), + } + + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Internal server error","upstream_status":500}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider.0, "anthropic"); + assert_eq!(message, "Internal server error"); + } + _ => panic!( + "Expected ApiInternalServerError for 500 status, got: {:?}", + error + ), + } + } + + #[test] + fn test_from_cloud_failure_with_standard_format() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_503".to_string(), + "Service unavailable".to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!("Expected ServerOverloaded error for upstream_http_503"), + } + } + + #[test] + fn test_upstream_http_error_connection_timeout() { + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":503}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ServerOverloaded { provider, .. } => { + assert_eq!(provider.0, "anthropic"); + } + _ => panic!( + "Expected ServerOverloaded error for connection timeout with 503 status, got: {:?}", + error + ), + } + + let error = LanguageModelCompletionError::from_cloud_failure( + String::from("anthropic").into(), + "upstream_http_error".to_string(), + r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout","upstream_status":500}"#.to_string(), + None, + ); + + match error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider.0, "anthropic"); + assert_eq!( + message, + "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers. reset reason: connection timeout" + ); + } + _ => panic!( + "Expected ApiInternalServerError for connection timeout with 500 status, got: {:?}", + error + ), + } + } +} diff --git a/crates/language_models/Cargo.toml b/crates/language_models/Cargo.toml index 514443ddec..e928df8a74 100644 --- a/crates/language_models/Cargo.toml +++ b/crates/language_models/Cargo.toml @@ -44,6 +44,7 @@ ollama = { workspace = true, features = ["schemars"] } open_ai = { workspace = true, features = ["schemars"] } open_router = { workspace = true, features = ["schemars"] } vercel = { workspace = true, features = ["schemars"] } +x_ai = { workspace = true, features = ["schemars"] } partial-json-fixer.workspace = true proto.workspace = true release_channel.workspace = true diff --git a/crates/language_models/src/language_models.rs b/crates/language_models/src/language_models.rs index c7324732c9..192f5a5fae 100644 --- a/crates/language_models/src/language_models.rs +++ b/crates/language_models/src/language_models.rs @@ -20,6 +20,7 @@ use crate::provider::ollama::OllamaLanguageModelProvider; use crate::provider::open_ai::OpenAiLanguageModelProvider; use crate::provider::open_router::OpenRouterLanguageModelProvider; use crate::provider::vercel::VercelLanguageModelProvider; +use crate::provider::x_ai::XAiLanguageModelProvider; pub use crate::settings::*; pub fn init(user_store: Entity, client: Arc, cx: &mut App) { @@ -81,5 +82,6 @@ fn register_language_model_providers( VercelLanguageModelProvider::new(client.http_client(), cx), cx, ); + registry.register_provider(XAiLanguageModelProvider::new(client.http_client(), cx), cx); registry.register_provider(CopilotChatLanguageModelProvider::new(cx), cx); } diff --git a/crates/language_models/src/provider.rs b/crates/language_models/src/provider.rs index 6bc93bd366..c717be7c90 100644 --- a/crates/language_models/src/provider.rs +++ b/crates/language_models/src/provider.rs @@ -10,3 +10,4 @@ pub mod ollama; pub mod open_ai; pub mod open_router; pub mod vercel; +pub mod x_ai; diff --git a/crates/language_models/src/provider/cloud.rs b/crates/language_models/src/provider/cloud.rs index 9b7fee228a..8d25af1a49 100644 --- a/crates/language_models/src/provider/cloud.rs +++ b/crates/language_models/src/provider/cloud.rs @@ -166,46 +166,9 @@ impl State { } let response = Self::fetch_models(client, llm_api_token, use_cloud).await?; - cx.update(|cx| { - this.update(cx, |this, cx| { - let mut models = Vec::new(); - - for model in response.models { - models.push(Arc::new(model.clone())); - - // Right now we represent thinking variants of models as separate models on the client, - // so we need to insert variants for any model that supports thinking. - if model.supports_thinking { - models.push(Arc::new(zed_llm_client::LanguageModel { - id: zed_llm_client::LanguageModelId( - format!("{}-thinking", model.id).into(), - ), - display_name: format!("{} Thinking", model.display_name), - ..model - })); - } - } - - this.default_model = models - .iter() - .find(|model| model.id == response.default_model) - .cloned(); - this.default_fast_model = models - .iter() - .find(|model| model.id == response.default_fast_model) - .cloned(); - this.recommended_models = response - .recommended_models - .iter() - .filter_map(|id| models.iter().find(|model| &model.id == id)) - .cloned() - .collect(); - this.models = models; - cx.notify(); - }) - })??; - - anyhow::Ok(()) + this.update(cx, |this, cx| { + this.update_models(response, cx); + }) }) .await .context("failed to fetch Zed models") @@ -216,12 +179,15 @@ impl State { }), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, - |this, _listener, _event, cx| { + move |this, _listener, _event, cx| { let client = this.client.clone(); let llm_api_token = this.llm_api_token.clone(); - cx.spawn(async move |_this, _cx| { + cx.spawn(async move |this, cx| { llm_api_token.refresh(&client).await?; - anyhow::Ok(()) + let response = Self::fetch_models(client, llm_api_token, use_cloud).await?; + this.update(cx, |this, cx| { + this.update_models(response, cx); + }) }) .detach_and_log_err(cx); }, @@ -264,6 +230,41 @@ impl State { })); } + fn update_models(&mut self, response: ListModelsResponse, cx: &mut Context) { + let mut models = Vec::new(); + + for model in response.models { + models.push(Arc::new(model.clone())); + + // Right now we represent thinking variants of models as separate models on the client, + // so we need to insert variants for any model that supports thinking. + if model.supports_thinking { + models.push(Arc::new(zed_llm_client::LanguageModel { + id: zed_llm_client::LanguageModelId(format!("{}-thinking", model.id).into()), + display_name: format!("{} Thinking", model.display_name), + ..model + })); + } + } + + self.default_model = models + .iter() + .find(|model| model.id == response.default_model) + .cloned(); + self.default_fast_model = models + .iter() + .find(|model| model.id == response.default_fast_model) + .cloned(); + self.recommended_models = response + .recommended_models + .iter() + .filter_map(|id| models.iter().find(|model| &model.id == id)) + .cloned() + .collect(); + self.models = models; + cx.notify(); + } + async fn fetch_models( client: Arc, llm_api_token: LlmApiToken, @@ -653,8 +654,62 @@ struct ApiError { headers: HeaderMap, } +/// Represents error responses from Zed's cloud API. +/// +/// Example JSON for an upstream HTTP error: +/// ```json +/// { +/// "code": "upstream_http_error", +/// "message": "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout", +/// "upstream_status": 503 +/// } +/// ``` +#[derive(Debug, serde::Deserialize)] +struct CloudApiError { + code: String, + message: String, + #[serde(default)] + #[serde(deserialize_with = "deserialize_optional_status_code")] + upstream_status: Option, + #[serde(default)] + retry_after: Option, +} + +fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result, D::Error> +where + D: serde::Deserializer<'de>, +{ + let opt: Option = Option::deserialize(deserializer)?; + Ok(opt.and_then(|code| StatusCode::from_u16(code).ok())) +} + impl From for LanguageModelCompletionError { fn from(error: ApiError) -> Self { + if let Ok(cloud_error) = serde_json::from_str::(&error.body) { + if cloud_error.code.starts_with("upstream_http_") { + let status = if let Some(status) = cloud_error.upstream_status { + status + } else if cloud_error.code.ends_with("_error") { + error.status + } else { + // If there's a status code in the code string (e.g. "upstream_http_429") + // then use that; otherwise, see if the JSON contains a status code. + cloud_error + .code + .strip_prefix("upstream_http_") + .and_then(|code_str| code_str.parse::().ok()) + .and_then(|code| StatusCode::from_u16(code).ok()) + .unwrap_or(error.status) + }; + + return LanguageModelCompletionError::UpstreamProviderError { + message: cloud_error.message, + status, + retry_after: cloud_error.retry_after.map(Duration::from_secs_f64), + }; + } + } + let retry_after = None; LanguageModelCompletionError::from_http_status( PROVIDER_NAME, @@ -1294,3 +1349,155 @@ impl Component for ZedAiConfiguration { ) } } + +#[cfg(test)] +mod tests { + use super::*; + use http_client::http::{HeaderMap, StatusCode}; + use language_model::LanguageModelCompletionError; + + #[test] + fn test_api_error_conversion_with_upstream_http_error() { + // upstream_http_error with 503 status should become ServerOverloaded + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout","upstream_status":503}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the Anthropic API: upstream connect error or disconnect/reset before headers, reset reason: connection timeout" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 503, got: {:?}", + completion_error + ), + } + + // upstream_http_error with 500 status should become ApiInternalServerError + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the OpenAI API: internal server error","upstream_status":500}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the OpenAI API: internal server error" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 500, got: {:?}", + completion_error + ), + } + + // upstream_http_error with 429 status should become RateLimitExceeded + let error_body = r#"{"code":"upstream_http_error","message":"Received an error from the Google API: rate limit exceeded","upstream_status":429}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { message, .. } => { + assert_eq!( + message, + "Received an error from the Google API: rate limit exceeded" + ); + } + _ => panic!( + "Expected UpstreamProviderError for upstream 429, got: {:?}", + completion_error + ), + } + + // Regular 500 error without upstream_http_error should remain ApiInternalServerError for Zed + let error_body = "Regular internal server error"; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::ApiInternalServerError { provider, message } => { + assert_eq!(provider, PROVIDER_NAME); + assert_eq!(message, "Regular internal server error"); + } + _ => panic!( + "Expected ApiInternalServerError for regular 500, got: {:?}", + completion_error + ), + } + + // upstream_http_429 format should be converted to UpstreamProviderError + let error_body = r#"{"code":"upstream_http_429","message":"Upstream Anthropic rate limit exceeded.","retry_after":30.5}"#; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::UpstreamProviderError { + message, + status, + retry_after, + } => { + assert_eq!(message, "Upstream Anthropic rate limit exceeded."); + assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); + assert_eq!(retry_after, Some(Duration::from_secs_f64(30.5))); + } + _ => panic!( + "Expected UpstreamProviderError for upstream_http_429, got: {:?}", + completion_error + ), + } + + // Invalid JSON in error body should fall back to regular error handling + let error_body = "Not JSON at all"; + + let api_error = ApiError { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: error_body.to_string(), + headers: HeaderMap::new(), + }; + + let completion_error: LanguageModelCompletionError = api_error.into(); + + match completion_error { + LanguageModelCompletionError::ApiInternalServerError { provider, .. } => { + assert_eq!(provider, PROVIDER_NAME); + } + _ => panic!( + "Expected ApiInternalServerError for invalid JSON, got: {:?}", + completion_error + ), + } + } +} diff --git a/crates/language_models/src/provider/open_router.rs b/crates/language_models/src/provider/open_router.rs index 5883da1e2f..90d375a6b2 100644 --- a/crates/language_models/src/provider/open_router.rs +++ b/crates/language_models/src/provider/open_router.rs @@ -376,7 +376,7 @@ impl LanguageModel for OpenRouterLanguageModel { fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { let model_id = self.model.id().trim().to_lowercase(); - if model_id.contains("gemini") { + if model_id.contains("gemini") || model_id.contains("grok-4") { LanguageModelToolSchemaFormat::JsonSchemaSubset } else { LanguageModelToolSchemaFormat::JsonSchema diff --git a/crates/language_models/src/provider/x_ai.rs b/crates/language_models/src/provider/x_ai.rs new file mode 100644 index 0000000000..5f6034571b --- /dev/null +++ b/crates/language_models/src/provider/x_ai.rs @@ -0,0 +1,571 @@ +use anyhow::{Context as _, Result, anyhow}; +use collections::BTreeMap; +use credentials_provider::CredentialsProvider; +use futures::{FutureExt, StreamExt, future::BoxFuture}; +use gpui::{AnyView, App, AsyncApp, Context, Entity, Subscription, Task, Window}; +use http_client::HttpClient; +use language_model::{ + AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, + LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, + LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, + LanguageModelToolChoice, LanguageModelToolSchemaFormat, RateLimiter, Role, +}; +use menu; +use open_ai::ResponseStreamEvent; +use schemars::JsonSchema; +use serde::{Deserialize, Serialize}; +use settings::{Settings, SettingsStore}; +use std::sync::Arc; +use strum::IntoEnumIterator; +use x_ai::Model; + +use ui::{ElevationIndex, List, Tooltip, prelude::*}; +use ui_input::SingleLineInput; +use util::ResultExt; + +use crate::{AllLanguageModelSettings, ui::InstructionListItem}; + +const PROVIDER_ID: &str = "x_ai"; +const PROVIDER_NAME: &str = "xAI"; + +#[derive(Default, Clone, Debug, PartialEq)] +pub struct XAiSettings { + pub api_url: String, + pub available_models: Vec, +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] +pub struct AvailableModel { + pub name: String, + pub display_name: Option, + pub max_tokens: u64, + pub max_output_tokens: Option, + pub max_completion_tokens: Option, +} + +pub struct XAiLanguageModelProvider { + http_client: Arc, + state: gpui::Entity, +} + +pub struct State { + api_key: Option, + api_key_from_env: bool, + _subscription: Subscription, +} + +const XAI_API_KEY_VAR: &str = "XAI_API_KEY"; + +impl State { + fn is_authenticated(&self) -> bool { + self.api_key.is_some() + } + + fn reset_api_key(&self, cx: &mut Context) -> Task> { + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + credentials_provider + .delete_credentials(&api_url, &cx) + .await + .log_err(); + this.update(cx, |this, cx| { + this.api_key = None; + this.api_key_from_env = false; + cx.notify(); + }) + }) + } + + fn set_api_key(&mut self, api_key: String, cx: &mut Context) -> Task> { + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + credentials_provider + .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx) + .await + .log_err(); + this.update(cx, |this, cx| { + this.api_key = Some(api_key); + cx.notify(); + }) + }) + } + + fn authenticate(&self, cx: &mut Context) -> Task> { + if self.is_authenticated() { + return Task::ready(Ok(())); + } + + let credentials_provider = ::global(cx); + let settings = &AllLanguageModelSettings::get_global(cx).x_ai; + let api_url = if settings.api_url.is_empty() { + x_ai::XAI_API_URL.to_string() + } else { + settings.api_url.clone() + }; + cx.spawn(async move |this, cx| { + let (api_key, from_env) = if let Ok(api_key) = std::env::var(XAI_API_KEY_VAR) { + (api_key, true) + } else { + let (_, api_key) = credentials_provider + .read_credentials(&api_url, &cx) + .await? + .ok_or(AuthenticateError::CredentialsNotFound)?; + ( + String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?, + false, + ) + }; + this.update(cx, |this, cx| { + this.api_key = Some(api_key); + this.api_key_from_env = from_env; + cx.notify(); + })?; + + Ok(()) + }) + } +} + +impl XAiLanguageModelProvider { + pub fn new(http_client: Arc, cx: &mut App) -> Self { + let state = cx.new(|cx| State { + api_key: None, + api_key_from_env: false, + _subscription: cx.observe_global::(|_this: &mut State, cx| { + cx.notify(); + }), + }); + + Self { http_client, state } + } + + fn create_language_model(&self, model: x_ai::Model) -> Arc { + Arc::new(XAiLanguageModel { + id: LanguageModelId::from(model.id().to_string()), + model, + state: self.state.clone(), + http_client: self.http_client.clone(), + request_limiter: RateLimiter::new(4), + }) + } +} + +impl LanguageModelProviderState for XAiLanguageModelProvider { + type ObservableEntity = State; + + fn observable_entity(&self) -> Option> { + Some(self.state.clone()) + } +} + +impl LanguageModelProvider for XAiLanguageModelProvider { + fn id(&self) -> LanguageModelProviderId { + LanguageModelProviderId(PROVIDER_ID.into()) + } + + fn name(&self) -> LanguageModelProviderName { + LanguageModelProviderName(PROVIDER_NAME.into()) + } + + fn icon(&self) -> IconName { + IconName::AiXAi + } + + fn default_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(x_ai::Model::default())) + } + + fn default_fast_model(&self, _cx: &App) -> Option> { + Some(self.create_language_model(x_ai::Model::default_fast())) + } + + fn provided_models(&self, cx: &App) -> Vec> { + let mut models = BTreeMap::default(); + + for model in x_ai::Model::iter() { + if !matches!(model, x_ai::Model::Custom { .. }) { + models.insert(model.id().to_string(), model); + } + } + + for model in &AllLanguageModelSettings::get_global(cx) + .x_ai + .available_models + { + models.insert( + model.name.clone(), + x_ai::Model::Custom { + name: model.name.clone(), + display_name: model.display_name.clone(), + max_tokens: model.max_tokens, + max_output_tokens: model.max_output_tokens, + max_completion_tokens: model.max_completion_tokens, + }, + ); + } + + models + .into_values() + .map(|model| self.create_language_model(model)) + .collect() + } + + fn is_authenticated(&self, cx: &App) -> bool { + self.state.read(cx).is_authenticated() + } + + fn authenticate(&self, cx: &mut App) -> Task