use anthropic::AnthropicModelMode;use anyhow::{Context as _, Result};use cloud_llm_client::{ CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME, CLIENT_SUPPORTS_X_AI_HEADER_NAME, CompletionBody, CompletionEvent, CompletionRequestStatus, EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, OUTDATED_LLM_TOKEN_HEADER_NAME, SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, ZED_VERSION_HEADER_NAME,};use futures::{ AsyncBufReadExt, AsyncReadExt as _, FutureExt, Stream, StreamExt, future::BoxFuture, io::BufReader, stream::{self, BoxStream},};use google_ai::GoogleModelMode;use gpui::{AppContext, AsyncApp, Context, Task};use http_client::http::{HeaderMap, HeaderValue};use http_client::{ AsyncBody, HttpClient, HttpClientWithUrl, HttpRequestExt, Method, Response, StatusCode,};use language_model::{ ANTHROPIC_PROVIDER_ID, ANTHROPIC_PROVIDER_NAME, 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,}; use schemars::JsonSchema;use semver::Version;use serde::{Deserialize, Serialize, de::DeserializeOwned};use std::collections::VecDeque;use std::pin::Pin;use std::str::FromStr;use std::sync::Arc;use std::task::Poll;use std::time::Duration;use thiserror::Error; use anthropic::completion::{AnthropicEventMapper, AnthropicPromptCacheMode, into_anthropic};use google_ai::completion::{GoogleEventMapper, into_google};use open_ai::completion::{ ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response,}; const PROVIDER_ID: LanguageModelProviderId = ZED_CLOUD_PROVIDER_ID;const PROVIDER_NAME: LanguageModelProviderName = ZED_CLOUD_PROVIDER_NAME; /// Trait for acquiring and refreshing LLM authentication tokens.pub trait CloudLlmTokenProvider: Send + Sync { type AuthContext: Clone + Send + 'static; fn auth_context(&self, cx: &impl AppContext) -> Self::AuthContext; fn cached_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>; fn refresh_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>; /// Whether the user has consented to upstream providers retaining /// inference logs for models that require it (see /// [`LanguageModel::requires_data_retention`]). fn has_data_retention_consent(&self, cx: &impl AppContext) -> bool;} /// Sends an authenticated request to the Zed LLM service, retrying once with/// a refreshed token if the server signals that the cached LLM token is/// expired or otherwise rejected. Returns the raw response so callers can/// inspect headers and stream the body.pub async fn authenticated_llm_request<TP: CloudLlmTokenProvider>( http_client: &HttpClientWithUrl, token_provider: &TP, auth_context: TP::AuthContext, build_request: impl Fn(&str) -> Result<http_client::Request<AsyncBody>>,) -> Result<Response<AsyncBody>> { let token = token_provider.cached_token(auth_context.clone()).await?; let response = http_client.send(build_request(&token)?).await?; if !needs_llm_token_refresh(&response) && response.status() != StatusCode::UNAUTHORIZED { return Ok(response); } log::info!("LLM token rejected; refreshing and retrying request"); let token = token_provider.refresh_token(auth_context).await?; http_client.send(build_request(&token)?).await} #[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]#[serde(tag = "type", rename_all = "lowercase")]pub enum ModelMode { #[default] Default, Thinking { /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`. budget_tokens: Option<u32>, },} impl From<ModelMode> for AnthropicModelMode { fn from(value: ModelMode) -> Self { match value { ModelMode::Default => AnthropicModelMode::Default, ModelMode::Thinking { budget_tokens } => AnthropicModelMode::Thinking { budget_tokens }, } }} pub struct CloudLanguageModel<TP: CloudLlmTokenProvider> { pub id: LanguageModelId, pub model: Arc<cloud_llm_client::LanguageModel>, pub token_provider: Arc<TP>, pub http_client: Arc<HttpClientWithUrl>, pub app_version: Option<Version>, pub request_limiter: RateLimiter,} pub struct PerformLlmCompletionResponse { pub response: Response<AsyncBody>, pub includes_status_messages: bool,} impl<TP: CloudLlmTokenProvider> CloudLanguageModel<TP> { pub async fn perform_llm_completion( http_client: &HttpClientWithUrl, token_provider: &TP, auth_context: TP::AuthContext, app_version: Option<Version>, body: CompletionBody, ) -> Result<PerformLlmCompletionResponse, LanguageModelCompletionError> { let url = http_client .build_zed_llm_url("/completions", &[]) .map_err(LanguageModelCompletionError::Other)?; let body = serde_json::to_string(&body).map_err(|error| { LanguageModelCompletionError::SerializeRequest { provider: PROVIDER_NAME, error, } })?; let mut response = authenticated_llm_request(http_client, token_provider, auth_context, |token| { Ok(http_client::Request::builder() .method(Method::POST) .uri(url.as_ref()) .when_some(app_version.as_ref(), |builder, app_version| { builder.header(ZED_VERSION_HEADER_NAME, app_version.to_string()) }) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {token}")) .header(CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, "true") .header(CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME, "true") .body(body.clone().into())?) }) .await .map_err(|error| LanguageModelCompletionError::HttpSend { provider: PROVIDER_NAME, error, })?; let status = response.status(); if status.is_success() { let includes_status_messages = response .headers() .get(SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME) .is_some(); return Ok(PerformLlmCompletionResponse { response, includes_status_messages, }); } if status == StatusCode::PAYMENT_REQUIRED { return Err(LanguageModelCompletionError::PaymentRequired); } let mut body = String::new(); let headers = response.headers().clone(); response .body_mut() .read_to_string(&mut body) .await .map_err(|error| LanguageModelCompletionError::ApiReadResponseError { provider: PROVIDER_NAME, error, })?; Err(ApiError { status, body, headers, } .into()) }} fn needs_llm_token_refresh(response: &Response<AsyncBody>) -> bool { response .headers() .get(EXPIRED_LLM_TOKEN_HEADER_NAME) .is_some() || response .headers() .get(OUTDATED_LLM_TOKEN_HEADER_NAME) .is_some()} #[derive(Debug, Error)]#[error("cloud language model request failed with status {status}: {body}")]struct ApiError { status: StatusCode, body: String, headers: HeaderMap<HeaderValue>,} /// 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<StatusCode>, #[serde(default)] retry_after: Option<f64>,} fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result<Option<StatusCode>, D::Error>where D: serde::Deserializer<'de>,{ let opt: Option<u16> = Option::deserialize(deserializer)?; Ok(opt.and_then(|code| StatusCode::from_u16(code).ok()))} impl From<ApiError> for LanguageModelCompletionError { fn from(error: ApiError) -> Self { if let Ok(cloud_error) = serde_json::from_str::<CloudApiError>(&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::<u16>().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), }; } return LanguageModelCompletionError::from_http_status( PROVIDER_NAME, error.status, cloud_error.message, None, ); } let retry_after = None; LanguageModelCompletionError::from_http_status( PROVIDER_NAME, error.status, error.body, retry_after, ) }} impl<TP: CloudLlmTokenProvider + 'static> LanguageModel for CloudLanguageModel<TP> { fn id(&self) -> LanguageModelId { self.id.clone() } fn name(&self) -> LanguageModelName { LanguageModelName::from(self.model.display_name.clone()) } fn provider_id(&self) -> LanguageModelProviderId { PROVIDER_ID } fn provider_name(&self) -> LanguageModelProviderName { PROVIDER_NAME } fn upstream_provider_id(&self) -> LanguageModelProviderId { use cloud_llm_client::LanguageModelProvider::*; match self.model.provider { Anthropic => ANTHROPIC_PROVIDER_ID, OpenAi => OPEN_AI_PROVIDER_ID, Google => GOOGLE_PROVIDER_ID, XAi => X_AI_PROVIDER_ID, } } fn upstream_provider_name(&self) -> LanguageModelProviderName { use cloud_llm_client::LanguageModelProvider::*; match self.model.provider { Anthropic => ANTHROPIC_PROVIDER_NAME, OpenAi => OPEN_AI_PROVIDER_NAME, Google => GOOGLE_PROVIDER_NAME, XAi => X_AI_PROVIDER_NAME, } } fn is_latest(&self) -> bool { self.model.is_latest } fn is_disabled(&self) -> Option<DisabledReason> { if self.model.is_disabled { self.model.disabled_reason.clone().map(DisabledReason::new) } else { None } } fn requires_data_retention(&self) -> bool { // Anthropic cannot offer Fable models with Zero Data Retention self.id .0 .as_ref() .starts_with(anthropic::FABLE_MODEL_ID_PREFIX) } fn refusal_fallback_model_id(&self) -> Option<&'static str> { if self .id .0 .as_ref() .starts_with(anthropic::FABLE_MODEL_ID_PREFIX) { Some(anthropic::FABLE_FALLBACK_MODEL_ID) } else { None } } fn supports_tools(&self) -> bool { self.model.supports_tools } fn supports_images(&self) -> bool { self.model.supports_images } fn supports_thinking(&self) -> bool { self.model.supports_thinking } fn supports_disabling_thinking(&self) -> bool { self.model.supports_disabling_thinking } fn supports_fast_mode(&self) -> bool { self.model.supports_fast_mode } fn supports_server_side_compaction(&self) -> bool { self.model.supports_server_side_compaction } fn supported_effort_levels(&self) -> Vec<LanguageModelEffortLevel> { self.model .supported_effort_levels .iter() .map(|effort_level| LanguageModelEffortLevel { name: effort_level.name.clone().into(), value: effort_level.value.clone().into(), is_default: effort_level.is_default.unwrap_or(false), }) .collect() } fn supports_streaming_tools(&self) -> bool { self.model.supports_streaming_tools } fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { match choice { LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any | LanguageModelToolChoice::None => true, } } fn supports_split_token_display(&self) -> bool { use cloud_llm_client::LanguageModelProvider::*; matches!(self.model.provider, OpenAi | XAi) } fn telemetry_id(&self) -> String { format!("zed.dev/{}", self.model.id) } fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { match self.model.provider { cloud_llm_client::LanguageModelProvider::Anthropic | cloud_llm_client::LanguageModelProvider::OpenAi => { LanguageModelToolSchemaFormat::JsonSchema } cloud_llm_client::LanguageModelProvider::Google | cloud_llm_client::LanguageModelProvider::XAi => { LanguageModelToolSchemaFormat::JsonSchemaSubset } } } fn max_token_count(&self) -> u64 { self.model.max_token_count as u64 } fn max_output_tokens(&self) -> Option<u64> { Some(self.model.max_output_tokens as u64) } fn stream_completion( &self, request: LanguageModelRequest, cx: &AsyncApp, ) -> BoxFuture< 'static, Result< BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>, LanguageModelCompletionError, >, > { if self.requires_data_retention() && !self.token_provider.has_data_retention_consent(cx) { let model_name = self.model.display_name.clone(); return async move { Err(LanguageModelCompletionError::DataRetentionConsentRequired { model_name }) } .boxed(); } let thread_id = request.thread_id.clone(); let prompt_id = request.prompt_id.clone(); let app_version = self.app_version.clone(); let thinking_allowed = request.thinking_allowed; let enable_thinking = thinking_allowed && self.model.supports_thinking; let provider_name = provider_name(&self.model.provider); match self.model.provider { cloud_llm_client::LanguageModelProvider::Anthropic => { let effort = request .thinking_effort .as_ref() .and_then(|effort| anthropic::Effort::from_str(effort).ok()); let mut request = match into_anthropic( request, self.model.id.to_string(), 1.0, self.model.max_output_tokens as u64, if enable_thinking { AnthropicModelMode::Thinking { budget_tokens: Some(4_096), } } else { AnthropicModelMode::Default }, AnthropicPromptCacheMode::Automatic, ) { Ok(request) => request, Err(error) => return async move { Err(error.into()) }.boxed(), }; if enable_thinking && effort.is_some() { request.thinking = Some(anthropic::Thinking::Adaptive { display: Some(anthropic::AdaptiveThinkingDisplay::Summarized), }); request.output_config = Some(anthropic::OutputConfig { effort }); } if !self.model.supports_fast_mode { request.speed = None; } let http_client = self.http_client.clone(); let token_provider = self.token_provider.clone(); let auth_context = token_provider.auth_context(cx); let future = self.request_limiter.stream(async move { let PerformLlmCompletionResponse { response, includes_status_messages,Lines 1–500 of 1147. Load more
use anthropic::AnthropicModelMode; use anyhow::{Context as _, Result}; use cloud_llm_client::{ CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME, CLIENT_SUPPORTS_X_AI_HEADER_NAME, CompletionBody, CompletionEvent, CompletionRequestStatus, EXPIRED_LLM_TOKEN_HEADER_NAME, ListModelsResponse, OUTDATED_LLM_TOKEN_HEADER_NAME, SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, ZED_VERSION_HEADER_NAME, }; use futures::{ AsyncBufReadExt, AsyncReadExt as _, FutureExt, Stream, StreamExt, future::BoxFuture, io::BufReader, stream::{self, BoxStream}, }; use google_ai::GoogleModelMode; use gpui::{AppContext, AsyncApp, Context, Task}; use http_client::http::{HeaderMap, HeaderValue}; use http_client::{ AsyncBody, HttpClient, HttpClientWithUrl, HttpRequestExt, Method, Response, StatusCode, }; use language_model::{ ANTHROPIC_PROVIDER_ID, ANTHROPIC_PROVIDER_NAME, 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, }; use schemars::JsonSchema; use semver::Version; use serde::{Deserialize, Serialize, de::DeserializeOwned}; use std::collections::VecDeque; use std::pin::Pin; use std::str::FromStr; use std::sync::Arc; use std::task::Poll; use std::time::Duration; use thiserror::Error; use anthropic::completion::{AnthropicEventMapper, AnthropicPromptCacheMode, into_anthropic}; use google_ai::completion::{GoogleEventMapper, into_google}; use open_ai::completion::{ ChatCompletionMaxTokensParameter, OpenAiEventMapper, OpenAiResponseEventMapper, into_open_ai, into_open_ai_response, }; const PROVIDER_ID: LanguageModelProviderId = ZED_CLOUD_PROVIDER_ID; const PROVIDER_NAME: LanguageModelProviderName = ZED_CLOUD_PROVIDER_NAME; /// Trait for acquiring and refreshing LLM authentication tokens. pub trait CloudLlmTokenProvider: Send + Sync { type AuthContext: Clone + Send + 'static; fn auth_context(&self, cx: &impl AppContext) -> Self::AuthContext; fn cached_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>; fn refresh_token(&self, auth_context: Self::AuthContext) -> BoxFuture<'static, Result<String>>; /// Whether the user has consented to upstream providers retaining /// inference logs for models that require it (see /// [`LanguageModel::requires_data_retention`]). fn has_data_retention_consent(&self, cx: &impl AppContext) -> bool; } /// Sends an authenticated request to the Zed LLM service, retrying once with /// a refreshed token if the server signals that the cached LLM token is /// expired or otherwise rejected. Returns the raw response so callers can /// inspect headers and stream the body. pub async fn authenticated_llm_request<TP: CloudLlmTokenProvider>( http_client: &HttpClientWithUrl, token_provider: &TP, auth_context: TP::AuthContext, build_request: impl Fn(&str) -> Result<http_client::Request<AsyncBody>>, ) -> Result<Response<AsyncBody>> { let token = token_provider.cached_token(auth_context.clone()).await?; let response = http_client.send(build_request(&token)?).await?; if !needs_llm_token_refresh(&response) && response.status() != StatusCode::UNAUTHORIZED { return Ok(response); } log::info!("LLM token rejected; refreshing and retrying request"); let token = token_provider.refresh_token(auth_context).await?; http_client.send(build_request(&token)?).await } #[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] #[serde(tag = "type", rename_all = "lowercase")] pub enum ModelMode { #[default] Default, Thinking { /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`. budget_tokens: Option<u32>, }, } impl From<ModelMode> for AnthropicModelMode { fn from(value: ModelMode) -> Self { match value { ModelMode::Default => AnthropicModelMode::Default, ModelMode::Thinking { budget_tokens } => AnthropicModelMode::Thinking { budget_tokens }, } } } pub struct CloudLanguageModel<TP: CloudLlmTokenProvider> { pub id: LanguageModelId, pub model: Arc<cloud_llm_client::LanguageModel>, pub token_provider: Arc<TP>, pub http_client: Arc<HttpClientWithUrl>, pub app_version: Option<Version>, pub request_limiter: RateLimiter, } pub struct PerformLlmCompletionResponse { pub response: Response<AsyncBody>, pub includes_status_messages: bool, } impl<TP: CloudLlmTokenProvider> CloudLanguageModel<TP> { pub async fn perform_llm_completion( http_client: &HttpClientWithUrl, token_provider: &TP, auth_context: TP::AuthContext, app_version: Option<Version>, body: CompletionBody, ) -> Result<PerformLlmCompletionResponse, LanguageModelCompletionError> { let url = http_client .build_zed_llm_url("/completions", &[]) .map_err(LanguageModelCompletionError::Other)?; let body = serde_json::to_string(&body).map_err(|error| { LanguageModelCompletionError::SerializeRequest { provider: PROVIDER_NAME, error, } })?; let mut response = authenticated_llm_request(http_client, token_provider, auth_context, |token| { Ok(http_client::Request::builder() .method(Method::POST) .uri(url.as_ref()) .when_some(app_version.as_ref(), |builder, app_version| { builder.header(ZED_VERSION_HEADER_NAME, app_version.to_string()) }) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {token}")) .header(CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, "true") .header(CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME, "true") .body(body.clone().into())?) }) .await .map_err(|error| LanguageModelCompletionError::HttpSend { provider: PROVIDER_NAME, error, })?; let status = response.status(); if status.is_success() { let includes_status_messages = response .headers() .get(SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME) .is_some(); return Ok(PerformLlmCompletionResponse { response, includes_status_messages, }); } if status == StatusCode::PAYMENT_REQUIRED { return Err(LanguageModelCompletionError::PaymentRequired); } let mut body = String::new(); let headers = response.headers().clone(); response .body_mut() .read_to_string(&mut body) .await .map_err(|error| LanguageModelCompletionError::ApiReadResponseError { provider: PROVIDER_NAME, error, })?; Err(ApiError { status, body, headers, } .into()) } } fn needs_llm_token_refresh(response: &Response<AsyncBody>) -> bool { response .headers() .get(EXPIRED_LLM_TOKEN_HEADER_NAME) .is_some() || response .headers() .get(OUTDATED_LLM_TOKEN_HEADER_NAME) .is_some() } #[derive(Debug, Error)] #[error("cloud language model request failed with status {status}: {body}")] struct ApiError { status: StatusCode, body: String, headers: HeaderMap<HeaderValue>, } /// 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<StatusCode>, #[serde(default)] retry_after: Option<f64>, } fn deserialize_optional_status_code<'de, D>(deserializer: D) -> Result<Option<StatusCode>, D::Error> where D: serde::Deserializer<'de>, { let opt: Option<u16> = Option::deserialize(deserializer)?; Ok(opt.and_then(|code| StatusCode::from_u16(code).ok())) } impl From<ApiError> for LanguageModelCompletionError { fn from(error: ApiError) -> Self { if let Ok(cloud_error) = serde_json::from_str::<CloudApiError>(&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::<u16>().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), }; } return LanguageModelCompletionError::from_http_status( PROVIDER_NAME, error.status, cloud_error.message, None, ); } let retry_after = None; LanguageModelCompletionError::from_http_status( PROVIDER_NAME, error.status, error.body, retry_after, ) } } impl<TP: CloudLlmTokenProvider + 'static> LanguageModel for CloudLanguageModel<TP> { fn id(&self) -> LanguageModelId { self.id.clone() } fn name(&self) -> LanguageModelName { LanguageModelName::from(self.model.display_name.clone()) } fn provider_id(&self) -> LanguageModelProviderId { PROVIDER_ID } fn provider_name(&self) -> LanguageModelProviderName { PROVIDER_NAME } fn upstream_provider_id(&self) -> LanguageModelProviderId { use cloud_llm_client::LanguageModelProvider::*; match self.model.provider { Anthropic => ANTHROPIC_PROVIDER_ID, OpenAi => OPEN_AI_PROVIDER_ID, Google => GOOGLE_PROVIDER_ID, XAi => X_AI_PROVIDER_ID, } } fn upstream_provider_name(&self) -> LanguageModelProviderName { use cloud_llm_client::LanguageModelProvider::*; match self.model.provider { Anthropic => ANTHROPIC_PROVIDER_NAME, OpenAi => OPEN_AI_PROVIDER_NAME, Google => GOOGLE_PROVIDER_NAME, XAi => X_AI_PROVIDER_NAME, } } fn is_latest(&self) -> bool { self.model.is_latest } fn is_disabled(&self) -> Option<DisabledReason> { if self.model.is_disabled { self.model.disabled_reason.clone().map(DisabledReason::new) } else { None } } fn requires_data_retention(&self) -> bool { // Anthropic cannot offer Fable models with Zero Data Retention self.id .0 .as_ref() .starts_with(anthropic::FABLE_MODEL_ID_PREFIX) } fn refusal_fallback_model_id(&self) -> Option<&'static str> { if self .id .0 .as_ref() .starts_with(anthropic::FABLE_MODEL_ID_PREFIX) { Some(anthropic::FABLE_FALLBACK_MODEL_ID) } else { None } } fn supports_tools(&self) -> bool { self.model.supports_tools } fn supports_images(&self) -> bool { self.model.supports_images } fn supports_thinking(&self) -> bool { self.model.supports_thinking } fn supports_disabling_thinking(&self) -> bool { self.model.supports_disabling_thinking } fn supports_fast_mode(&self) -> bool { self.model.supports_fast_mode } fn supports_server_side_compaction(&self) -> bool { self.model.supports_server_side_compaction } fn supported_effort_levels(&self) -> Vec<LanguageModelEffortLevel> { self.model .supported_effort_levels .iter() .map(|effort_level| LanguageModelEffortLevel { name: effort_level.name.clone().into(), value: effort_level.value.clone().into(), is_default: effort_level.is_default.unwrap_or(false), }) .collect() } fn supports_streaming_tools(&self) -> bool { self.model.supports_streaming_tools } fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { match choice { LanguageModelToolChoice::Auto | LanguageModelToolChoice::Any | LanguageModelToolChoice::None => true, } } fn supports_split_token_display(&self) -> bool { use cloud_llm_client::LanguageModelProvider::*; matches!(self.model.provider, OpenAi | XAi) } fn telemetry_id(&self) -> String { format!("zed.dev/{}", self.model.id) } fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { match self.model.provider { cloud_llm_client::LanguageModelProvider::Anthropic | cloud_llm_client::LanguageModelProvider::OpenAi => { LanguageModelToolSchemaFormat::JsonSchema } cloud_llm_client::LanguageModelProvider::Google | cloud_llm_client::LanguageModelProvider::XAi => { LanguageModelToolSchemaFormat::JsonSchemaSubset } } } fn max_token_count(&self) -> u64 { self.model.max_token_count as u64 } fn max_output_tokens(&self) -> Option<u64> { Some(self.model.max_output_tokens as u64) } fn stream_completion( &self, request: LanguageModelRequest, cx: &AsyncApp, ) -> BoxFuture< 'static, Result< BoxStream<'static, Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>, LanguageModelCompletionError, >, > { if self.requires_data_retention() && !self.token_provider.has_data_retention_consent(cx) { let model_name = self.model.display_name.clone(); return async move { Err(LanguageModelCompletionError::DataRetentionConsentRequired { model_name }) } .boxed(); } let thread_id = request.thread_id.clone(); let prompt_id = request.prompt_id.clone(); let app_version = self.app_version.clone(); let thinking_allowed = request.thinking_allowed; let enable_thinking = thinking_allowed && self.model.supports_thinking; let provider_name = provider_name(&self.model.provider); match self.model.provider { cloud_llm_client::LanguageModelProvider::Anthropic => { let effort = request .thinking_effort .as_ref() .and_then(|effort| anthropic::Effort::from_str(effort).ok()); let mut request = match into_anthropic( request, self.model.id.to_string(), 1.0, self.model.max_output_tokens as u64, if enable_thinking { AnthropicModelMode::Thinking { budget_tokens: Some(4_096), } } else { AnthropicModelMode::Default }, AnthropicPromptCacheMode::Automatic, ) { Ok(request) => request, Err(error) => return async move { Err(error.into()) }.boxed(), }; if enable_thinking && effort.is_some() { request.thinking = Some(anthropic::Thinking::Adaptive { display: Some(anthropic::AdaptiveThinkingDisplay::Summarized), }); request.output_config = Some(anthropic::OutputConfig { effort }); } if !self.model.supports_fast_mode { request.speed = None; } let http_client = self.http_client.clone(); let token_provider = self.token_provider.clone(); let auth_context = token_provider.auth_context(cx); let future = self.request_limiter.stream(async move { let PerformLlmCompletionResponse { response, includes_status_messages,