Richard FeldmanThought signatures (#42915)c0fadae12 months agoHistory
1008 lines · 34.4 KB · Rust
RawDownload
mod model;mod rate_limiter;mod registry;mod request;mod role;mod telemetry;pub mod tool_schema; #[cfg(any(test, feature = "test-support"))]pub mod fake_provider; use anthropic::{AnthropicError, parse_prompt_too_long};use anyhow::{Result, anyhow};use client::Client;use cloud_llm_client::{CompletionMode, CompletionRequestStatus};use futures::FutureExt;use futures::{StreamExt, future::BoxFuture, stream::BoxStream};use gpui::{AnyView, App, AsyncApp, SharedString, Task, Window};use http_client::{StatusCode, http};use icons::IconName;use open_router::OpenRouterError;use parking_lot::Mutex;use serde::{Deserialize, Serialize};pub use settings::LanguageModelCacheConfiguration;use std::ops::{Add, Sub};use std::str::FromStr;use std::sync::Arc;use std::time::Duration;use std::{fmt, io};use thiserror::Error;use util::serde::is_default; pub use crate::model::*;pub use crate::rate_limiter::*;pub use crate::registry::*;pub use crate::request::*;pub use crate::role::*;pub use crate::telemetry::*;pub use crate::tool_schema::LanguageModelToolSchemaFormat; pub const ANTHROPIC_PROVIDER_ID: LanguageModelProviderId =    LanguageModelProviderId::new("anthropic");pub const ANTHROPIC_PROVIDER_NAME: LanguageModelProviderName =    LanguageModelProviderName::new("Anthropic"); pub const GOOGLE_PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("google");pub const GOOGLE_PROVIDER_NAME: LanguageModelProviderName =    LanguageModelProviderName::new("Google AI"); pub const OPEN_AI_PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openai");pub const OPEN_AI_PROVIDER_NAME: LanguageModelProviderName =    LanguageModelProviderName::new("OpenAI"); pub const X_AI_PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("x_ai");pub const X_AI_PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("xAI"); pub const ZED_CLOUD_PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("zed.dev");pub const ZED_CLOUD_PROVIDER_NAME: LanguageModelProviderName =    LanguageModelProviderName::new("Zed"); pub fn init(client: Arc<Client>, cx: &mut App) {    init_settings(cx);    RefreshLlmTokenListener::register(client, cx);} pub fn init_settings(cx: &mut App) {    registry::init(cx);} /// A completion event from a language model.#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]pub enum LanguageModelCompletionEvent {    StatusUpdate(CompletionRequestStatus),    Stop(StopReason),    Text(String),    Thinking {        text: String,        signature: Option<String>,    },    RedactedThinking {        data: String,    },    ToolUse(LanguageModelToolUse),    ToolUseJsonParseError {        id: LanguageModelToolUseId,        tool_name: Arc<str>,        raw_input: Arc<str>,        json_parse_error: String,    },    StartMessage {        message_id: String,    },    UsageUpdate(TokenUsage),} #[derive(Error, Debug)]pub enum LanguageModelCompletionError {    #[error("prompt too large for context window")]    PromptTooLarge { tokens: Option<u64> },    #[error("missing {provider} API key")]    NoApiKey { provider: LanguageModelProviderName },    #[error("{provider}'s API rate limit exceeded")]    RateLimitExceeded {        provider: LanguageModelProviderName,        retry_after: Option<Duration>,    },    #[error("{provider}'s API servers are overloaded right now")]    ServerOverloaded {        provider: LanguageModelProviderName,        retry_after: Option<Duration>,    },    #[error("{provider}'s API server reported an internal server error: {message}")]    ApiInternalServerError {        provider: LanguageModelProviderName,        message: String,    },    #[error("{message}")]    UpstreamProviderError {        message: String,        status: StatusCode,        retry_after: Option<Duration>,    },    #[error("HTTP response error from {provider}'s API: status {status_code} - {message:?}")]    HttpResponseError {        provider: LanguageModelProviderName,        status_code: StatusCode,        message: String,    },     // Client errors    #[error("invalid request format to {provider}'s API: {message}")]    BadRequestFormat {        provider: LanguageModelProviderName,        message: String,    },    #[error("authentication error with {provider}'s API: {message}")]    AuthenticationError {        provider: LanguageModelProviderName,        message: String,    },    #[error("Permission error with {provider}'s API: {message}")]    PermissionError {        provider: LanguageModelProviderName,        message: String,    },    #[error("language model provider API endpoint not found")]    ApiEndpointNotFound { provider: LanguageModelProviderName },    #[error("I/O error reading response from {provider}'s API")]    ApiReadResponseError {        provider: LanguageModelProviderName,        #[source]        error: io::Error,    },    #[error("error serializing request to {provider} API")]    SerializeRequest {        provider: LanguageModelProviderName,        #[source]        error: serde_json::Error,    },    #[error("error building request body to {provider} API")]    BuildRequestBody {        provider: LanguageModelProviderName,        #[source]        error: http::Error,    },    #[error("error sending HTTP request to {provider} API")]    HttpSend {        provider: LanguageModelProviderName,        #[source]        error: anyhow::Error,    },    #[error("error deserializing {provider} API response")]    DeserializeResponse {        provider: LanguageModelProviderName,        #[source]        error: serde_json::Error,    },     // TODO: Ideally this would be removed in favor of having a comprehensive list of errors.    #[error(transparent)]    Other(#[from] anyhow::Error),} impl LanguageModelCompletionError {    fn parse_upstream_error_json(message: &str) -> Option<(StatusCode, String)> {        let error_json = serde_json::from_str::<serde_json::Value>(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,        message: String,        retry_after: Option<Duration>,    ) -> Self {        if let Some(tokens) = parse_prompt_too_long(&message) {            // TODO: currently Anthropic PAYLOAD_TOO_LARGE response may cause INTERNAL_SERVER_ERROR            // to be reported. This is a temporary workaround to handle this in the case where the            // token limit has been exceeded.            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())        {            Self::from_http_status(upstream_provider, status_code, message, retry_after)        } else if let Some(status_code) = code            .strip_prefix("http_")            .and_then(|code| StatusCode::from_str(code).ok())        {            Self::from_http_status(ZED_CLOUD_PROVIDER_NAME, status_code, message, retry_after)        } else {            anyhow!("completion request failed, code: {code}, message: {message}").into()        }    }     pub fn from_http_status(        provider: LanguageModelProviderName,        status_code: StatusCode,        message: String,        retry_after: Option<Duration>,    ) -> Self {        match status_code {            StatusCode::BAD_REQUEST => Self::BadRequestFormat { provider, message },            StatusCode::UNAUTHORIZED => Self::AuthenticationError { provider, message },            StatusCode::FORBIDDEN => Self::PermissionError { provider, message },            StatusCode::NOT_FOUND => Self::ApiEndpointNotFound { provider },            StatusCode::PAYLOAD_TOO_LARGE => Self::PromptTooLarge {                tokens: parse_prompt_too_long(&message),            },            StatusCode::TOO_MANY_REQUESTS => Self::RateLimitExceeded {                provider,                retry_after,            },            StatusCode::INTERNAL_SERVER_ERROR => Self::ApiInternalServerError { provider, message },            StatusCode::SERVICE_UNAVAILABLE => Self::ServerOverloaded {                provider,                retry_after,            },            _ if status_code.as_u16() == 529 => Self::ServerOverloaded {                provider,                retry_after,            },            _ => Self::HttpResponseError {                provider,                status_code,                message,            },        }    }} impl From<AnthropicError> for LanguageModelCompletionError {    fn from(error: AnthropicError) -> Self {        let provider = ANTHROPIC_PROVIDER_NAME;        match error {            AnthropicError::SerializeRequest(error) => Self::SerializeRequest { provider, error },            AnthropicError::BuildRequestBody(error) => Self::BuildRequestBody { provider, error },            AnthropicError::HttpSend(error) => Self::HttpSend { provider, error },            AnthropicError::DeserializeResponse(error) => {                Self::DeserializeResponse { provider, error }            }            AnthropicError::ReadResponse(error) => Self::ApiReadResponseError { provider, error },            AnthropicError::HttpResponseError {                status_code,                message,            } => Self::HttpResponseError {                provider,                status_code,                message,            },            AnthropicError::RateLimit { retry_after } => Self::RateLimitExceeded {                provider,                retry_after: Some(retry_after),            },            AnthropicError::ServerOverloaded { retry_after } => Self::ServerOverloaded {                provider,                retry_after,            },            AnthropicError::ApiError(api_error) => api_error.into(),        }    }} impl From<anthropic::ApiError> for LanguageModelCompletionError {    fn from(error: anthropic::ApiError) -> Self {        use anthropic::ApiErrorCode::*;        let provider = ANTHROPIC_PROVIDER_NAME;        match error.code() {            Some(code) => match code {                InvalidRequestError => Self::BadRequestFormat {                    provider,                    message: error.message,                },                AuthenticationError => Self::AuthenticationError {                    provider,                    message: error.message,                },                PermissionError => Self::PermissionError {                    provider,                    message: error.message,                },                NotFoundError => Self::ApiEndpointNotFound { provider },                RequestTooLarge => Self::PromptTooLarge {                    tokens: parse_prompt_too_long(&error.message),                },                RateLimitError => Self::RateLimitExceeded {                    provider,                    retry_after: None,                },                ApiError => Self::ApiInternalServerError {                    provider,                    message: error.message,                },                OverloadedError => Self::ServerOverloaded {                    provider,                    retry_after: None,                },            },            None => Self::Other(error.into()),        }    }} impl From<open_ai::RequestError> for LanguageModelCompletionError {    fn from(error: open_ai::RequestError) -> Self {        match error {            open_ai::RequestError::HttpResponseError {                provider,                status_code,                body,                headers,            } => {                let retry_after = headers                    .get(http::header::RETRY_AFTER)                    .and_then(|val| val.to_str().ok()?.parse::<u64>().ok())                    .map(Duration::from_secs);                 Self::from_http_status(provider.into(), status_code, body, retry_after)            }            open_ai::RequestError::Other(e) => Self::Other(e),        }    }} impl From<OpenRouterError> for LanguageModelCompletionError {    fn from(error: OpenRouterError) -> Self {        let provider = LanguageModelProviderName::new("OpenRouter");        match error {            OpenRouterError::SerializeRequest(error) => Self::SerializeRequest { provider, error },            OpenRouterError::BuildRequestBody(error) => Self::BuildRequestBody { provider, error },            OpenRouterError::HttpSend(error) => Self::HttpSend { provider, error },            OpenRouterError::DeserializeResponse(error) => {                Self::DeserializeResponse { provider, error }            }            OpenRouterError::ReadResponse(error) => Self::ApiReadResponseError { provider, error },            OpenRouterError::RateLimit { retry_after } => Self::RateLimitExceeded {                provider,                retry_after: Some(retry_after),            },            OpenRouterError::ServerOverloaded { retry_after } => Self::ServerOverloaded {                provider,                retry_after,            },            OpenRouterError::ApiError(api_error) => api_error.into(),        }    }} impl From<open_router::ApiError> for LanguageModelCompletionError {    fn from(error: open_router::ApiError) -> Self {        use open_router::ApiErrorCode::*;        let provider = LanguageModelProviderName::new("OpenRouter");        match error.code {            InvalidRequestError => Self::BadRequestFormat {                provider,                message: error.message,            },            AuthenticationError => Self::AuthenticationError {                provider,                message: error.message,            },            PaymentRequiredError => Self::AuthenticationError {                provider,                message: format!("Payment required: {}", error.message),            },            PermissionError => Self::PermissionError {                provider,                message: error.message,            },            RequestTimedOut => Self::HttpResponseError {                provider,                status_code: StatusCode::REQUEST_TIMEOUT,                message: error.message,            },            RateLimitError => Self::RateLimitExceeded {                provider,                retry_after: None,            },            ApiError => Self::ApiInternalServerError {                provider,                message: error.message,            },            OverloadedError => Self::ServerOverloaded {                provider,                retry_after: None,            },        }    }} #[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]#[serde(rename_all = "snake_case")]pub enum StopReason {    EndTurn,    MaxTokens,    ToolUse,    Refusal,} #[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize, Default)]pub struct TokenUsage {    #[serde(default, skip_serializing_if = "is_default")]    pub input_tokens: u64,    #[serde(default, skip_serializing_if = "is_default")]    pub output_tokens: u64,    #[serde(default, skip_serializing_if = "is_default")]    pub cache_creation_input_tokens: u64,    #[serde(default, skip_serializing_if = "is_default")]    pub cache_read_input_tokens: u64,} impl TokenUsage {    pub fn total_tokens(&self) -> u64 {        self.input_tokens            + self.output_tokens            + self.cache_read_input_tokens            + self.cache_creation_input_tokens    }} impl Add<TokenUsage> for TokenUsage {    type Output = Self;     fn add(self, other: Self) -> Self {        Self {            input_tokens: self.input_tokens + other.input_tokens,            output_tokens: self.output_tokens + other.output_tokens,            cache_creation_input_tokens: self.cache_creation_input_tokens                + other.cache_creation_input_tokens,            cache_read_input_tokens: self.cache_read_input_tokens + other.cache_read_input_tokens,        }    }} impl Sub<TokenUsage> for TokenUsage {    type Output = Self;     fn sub(self, other: Self) -> Self {        Self {            input_tokens: self.input_tokens - other.input_tokens,            output_tokens: self.output_tokens - other.output_tokens,            cache_creation_input_tokens: self.cache_creation_input_tokens                - other.cache_creation_input_tokens,            cache_read_input_tokens: self.cache_read_input_tokens - other.cache_read_input_tokens,        }    }} #[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]pub struct LanguageModelToolUseId(Arc<str>); impl fmt::Display for LanguageModelToolUseId {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        write!(f, "{}", self.0)    }}

Lines 1–500 of 1008.