1529 lines · 57.2 KB · Rust
RawDownload
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, CompactionResult, DisabledReason,    GOOGLE_PROVIDER_ID, GOOGLE_PROVIDER_NAME, LanguageModel, LanguageModelCompletionError,    LanguageModelCompletionEvent, LanguageModelEffortLevel, LanguageModelId, LanguageModelName,    LanguageModelProviderId, LanguageModelProviderName, LanguageModelRequest,    LanguageModelToolChoice, LanguageModelToolSchemaFormat, OPEN_AI_PROVIDER_ID,    OPEN_AI_PROVIDER_NAME, RateLimiter, X_AI_PROVIDER_ID, X_AI_PROVIDER_NAME,    ZED_CLOUD_PROVIDER_ID, ZED_CLOUD_PROVIDER_NAME,}; 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, token_usage_from_response_usage,}; 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> {        Self::perform_llm_request(            "/completions",            true,            http_client,            token_provider,            auth_context,            app_version,            body,        )        .await    }     async fn perform_llm_compaction(        http_client: &HttpClientWithUrl,        token_provider: &TP,        auth_context: TP::AuthContext,        app_version: Option<Version>,        body: CompletionBody,    ) -> Result<PerformLlmCompletionResponse, LanguageModelCompletionError> {        Self::perform_llm_request(            "/completions/compact",            false,            http_client,            token_provider,            auth_context,            app_version,            body,        )        .await    }     async fn perform_llm_request(        path: &str,        request_status_messages: bool,        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(path, &[])            .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| {                let mut request = 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}"));                if request_status_messages {                    request = request                        .header(CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME, "true")                        .header(CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME, "true");                }                Ok(request.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 = request_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 supports_explicit_compaction(&self) -> bool {        self.model.provider == cloud_llm_client::LanguageModelProvider::OpenAi            && self.model.supports_server_side_compaction    }     fn compact(        &self,        request: LanguageModelRequest,        cx: &AsyncApp,    ) -> BoxFuture<'static, Result<CompactionResult, LanguageModelCompletionError>> {        if !self.supports_explicit_compaction() {            return async {                Err(LanguageModelCompletionError::Other(anyhow::anyhow!(                    "this cloud model does not support explicit compaction"                )))            }            .boxed();        }         let thread_id = request.thread_id.clone();        let prompt_id = request.prompt_id.clone();        let app_version = self.app_version.clone();        let model_provider = self.model.provider;        let provider_name = provider_name(&self.model.provider);        let supports_none_reasoning_effort =            self.model.supported_effort_levels.iter().any(|effort| {                open_ai::ReasoningEffort::from_str(&effort.value)                    .is_ok_and(|effort| effort == open_ai::ReasoningEffort::None)            });        // Cloud proxies to OpenAI's own infrastructure, so the resulting        // compaction state is owned by (and interchangeable with) OpenAI        // proper, not by the cloud transport.        let request = match into_open_ai_response(            request,            &self.model.id.0,            self.model.supports_parallel_tool_calls,            true,            None,            None,            supports_none_reasoning_effort,            &OPEN_AI_PROVIDER_ID,        ) {            Ok(request) => request,            Err(error) => return async move { Err(error.into()) }.boxed(),        };        let compact_request = request.into_compact_request();        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.run(async move {            let PerformLlmCompletionResponse {                response,                includes_status_messages,            } = Self::perform_llm_compaction(                &http_client,                &*token_provider,                auth_context,                app_version,                CompletionBody {                    thread_id,                    prompt_id,                    provider: model_provider,                    model: compact_request.model.clone(),                    provider_request: serde_json::to_value(compact_request).map_err(|error| {                        LanguageModelCompletionError::SerializeRequest {                            provider: provider_name.clone(),                            error,                        }                    })?,                },            )            .await?;             let events = response_lines::<open_ai::responses::CompactedResponse>(                response,                includes_status_messages,            );            futures::pin_mut!(events);

Lines 1–500 of 1529.