989 lines · 36.8 KB · Rust
RawDownload
use anthropic::AnthropicModelMode;use anyhow::{Context as _, Result, anyhow};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, 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,    PaymentRequiredError, 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::{    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>>;} /// 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> {        let url = http_client.build_zed_llm_url("/completions", &[])?;        let body = serde_json::to_string(&body)?;        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?;         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(anyhow!(PaymentRequiredError));        }         let mut body = String::new();        let headers = response.headers().clone();        response.body_mut().read_to_string(&mut body).await?;        Err(anyhow!(ApiError {            status,            body,            headers        }))    }} 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 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_fast_mode(&self) -> bool {        self.model.supports_fast_mode    }     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,        >,    > {        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 = 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,                );                 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,                    } = Self::perform_llm_completion(                        &http_client,                        &*token_provider,                        auth_context,                        app_version,                        CompletionBody {                            thread_id,                            prompt_id,                            provider: cloud_llm_client::LanguageModelProvider::Anthropic,                            model: request.model.clone(),                            provider_request: serde_json::to_value(&request)                                .map_err(|e| anyhow!(e))?,                        },                    )                    .await                    .map_err(|err| match err.downcast::<ApiError>() {                        Ok(api_err) => anyhow!(LanguageModelCompletionError::from(api_err)),                        Err(err) => anyhow!(err),                    })?;                     let mut mapper = AnthropicEventMapper::new();                    Ok(map_cloud_completion_events(                        Box::pin(response_lines(response, includes_status_messages)),                        &provider_name,                        move |event| mapper.map_event(event),                    ))                });                async move { Ok(future.await?.boxed()) }.boxed()            }            cloud_llm_client::LanguageModelProvider::OpenAi => {                let http_client = self.http_client.clone();                let token_provider = self.token_provider.clone();                let effort = request                    .thinking_effort                    .as_ref()                    .and_then(|effort| open_ai::ReasoningEffort::from_str(effort).ok())                    .filter(|effort| *effort != open_ai::ReasoningEffort::None);                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)                    });                 let mut request = into_open_ai_response(                    request,                    &self.model.id.0,                    self.model.supports_parallel_tool_calls,                    true,                    None,                    None,                    supports_none_reasoning_effort,                );                 if enable_thinking && let Some(effort) = effort {                    request.reasoning = Some(open_ai::responses::ReasoningConfig {                        effort,                        summary: Some(open_ai::responses::ReasoningSummaryMode::Auto),                    });                }                 let auth_context = token_provider.auth_context(cx);                let future = self.request_limiter.stream(async move {                    let PerformLlmCompletionResponse {                        response,                        includes_status_messages,                    } = Self::perform_llm_completion(                        &http_client,                        &*token_provider,                        auth_context,                        app_version,                        CompletionBody {                            thread_id,                            prompt_id,

Lines 1–500 of 989.