mod authorization;pub mod db;mod token; use crate::api::events::SnowflakeRow;use crate::api::CloudflareIpCountryHeader;use crate::build_kinesis_client;use crate::{db::UserId, executor::Executor, Cents, Config, Error, Result};use anyhow::{anyhow, Context as _};use authorization::authorize_access_to_language_model;use axum::routing::get;use axum::{    body::Body,    http::{self, HeaderName, HeaderValue, Request, StatusCode},    middleware::{self, Next},    response::{IntoResponse, Response},    routing::post,    Extension, Json, Router, TypedHeader,};use chrono::{DateTime, Duration, Utc};use collections::HashMap;use db::TokenUsage;use db::{usage_measure::UsageMeasure, ActiveUserCount, LlmDatabase};use futures::{Stream, StreamExt as _};use reqwest_client::ReqwestClient;use rpc::{    proto::Plan, LanguageModelProvider, PerformCompletionParams, EXPIRED_LLM_TOKEN_HEADER_NAME,};use rpc::{    ListModelsResponse, PredictEditsParams, PredictEditsResponse,    MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME,};use serde_json::json;use std::{    pin::Pin,    sync::Arc,    task::{Context, Poll},};use strum::IntoEnumIterator;use tokio::sync::RwLock;use util::ResultExt; pub use token::*; pub struct LlmState {    pub config: Config,    pub executor: Executor,    pub db: Arc<LlmDatabase>,    pub http_client: ReqwestClient,    pub kinesis_client: Option<aws_sdk_kinesis::Client>,    active_user_count_by_model:        RwLock<HashMap<(LanguageModelProvider, String), (DateTime<Utc>, ActiveUserCount)>>,} const ACTIVE_USER_COUNT_CACHE_DURATION: Duration = Duration::seconds(30); impl LlmState {    pub async fn new(config: Config, executor: Executor) -> Result<Arc<Self>> {        let database_url = config            .llm_database_url            .as_ref()            .ok_or_else(|| anyhow!("missing LLM_DATABASE_URL"))?;        let max_connections = config            .llm_database_max_connections            .ok_or_else(|| anyhow!("missing LLM_DATABASE_MAX_CONNECTIONS"))?;         let mut db_options = db::ConnectOptions::new(database_url);        db_options.max_connections(max_connections);        let mut db = LlmDatabase::new(db_options, executor.clone()).await?;        db.initialize().await?;         let db = Arc::new(db);         let user_agent = format!("Zed Server/{}", env!("CARGO_PKG_VERSION"));        let http_client =            ReqwestClient::user_agent(&user_agent).context("failed to construct http client")?;         let this = Self {            executor,            db,            http_client,            kinesis_client: if config.kinesis_access_key.is_some() {                build_kinesis_client(&config).await.log_err()            } else {                None            },            active_user_count_by_model: RwLock::new(HashMap::default()),            config,        };         Ok(Arc::new(this))    }     pub async fn get_active_user_count(        &self,        provider: LanguageModelProvider,        model: &str,    ) -> Result<ActiveUserCount> {        let now = Utc::now();         {            let active_user_count_by_model = self.active_user_count_by_model.read().await;            if let Some((last_updated, count)) =                active_user_count_by_model.get(&(provider, model.to_string()))            {                if now - *last_updated < ACTIVE_USER_COUNT_CACHE_DURATION {                    return Ok(*count);                }            }        }         let mut cache = self.active_user_count_by_model.write().await;        let new_count = self.db.get_active_user_count(provider, model, now).await?;        cache.insert((provider, model.to_string()), (now, new_count));        Ok(new_count)    }} pub fn routes() -> Router<(), Body> {    Router::new()        .route("/models", get(list_models))        .route("/completion", post(perform_completion))        .route("/predict_edits", post(predict_edits))        .layer(middleware::from_fn(validate_api_token))} async fn validate_api_token<B>(mut req: Request<B>, next: Next<B>) -> impl IntoResponse {    let token = req        .headers()        .get(http::header::AUTHORIZATION)        .and_then(|header| header.to_str().ok())        .ok_or_else(|| {            Error::http(                StatusCode::BAD_REQUEST,                "missing authorization header".to_string(),            )        })?        .strip_prefix("Bearer ")        .ok_or_else(|| {            Error::http(                StatusCode::BAD_REQUEST,                "invalid authorization header".to_string(),            )        })?;     let state = req.extensions().get::<Arc<LlmState>>().unwrap();    match LlmTokenClaims::validate(token, &state.config) {        Ok(claims) => {            if state.db.is_access_token_revoked(&claims.jti).await? {                return Err(Error::http(                    StatusCode::UNAUTHORIZED,                    "unauthorized".to_string(),                ));            }             tracing::Span::current()                .record("user_id", claims.user_id)                .record("login", claims.github_user_login.clone())                .record("authn.jti", &claims.jti)                .record("is_staff", claims.is_staff);             req.extensions_mut().insert(claims);            Ok::<_, Error>(next.run(req).await.into_response())        }        Err(ValidateLlmTokenError::Expired) => Err(Error::Http(            StatusCode::UNAUTHORIZED,            "unauthorized".to_string(),            [(                HeaderName::from_static(EXPIRED_LLM_TOKEN_HEADER_NAME),                HeaderValue::from_static("true"),            )]            .into_iter()            .collect(),        )),        Err(_err) => Err(Error::http(            StatusCode::UNAUTHORIZED,            "unauthorized".to_string(),        )),    }} async fn list_models(    Extension(state): Extension<Arc<LlmState>>,    Extension(claims): Extension<LlmTokenClaims>,    country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,) -> Result<Json<ListModelsResponse>> {    let country_code = country_code_header.map(|header| header.to_string());     let mut accessible_models = Vec::new();     for (provider, model) in state.db.all_models() {        let authorize_result = authorize_access_to_language_model(            &state.config,            &claims,            country_code.as_deref(),            provider,            &model.name,        );         if authorize_result.is_ok() {            accessible_models.push(rpc::LanguageModel {                provider,                name: model.name,            });        }    }     Ok(Json(ListModelsResponse {        models: accessible_models,    }))} async fn perform_completion(    Extension(state): Extension<Arc<LlmState>>,    Extension(claims): Extension<LlmTokenClaims>,    country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,    Json(params): Json<PerformCompletionParams>,) -> Result<impl IntoResponse> {    let model = normalize_model_name(        state.db.model_names_for_provider(params.provider),        params.model,    );     authorize_access_to_language_model(        &state.config,        &claims,        country_code_header            .map(|header| header.to_string())            .as_deref(),        params.provider,        &model,    )?;     check_usage_limit(&state, params.provider, &model, &claims).await?;     let stream = match params.provider {        LanguageModelProvider::Anthropic => {            let api_key = if claims.is_staff {                state                    .config                    .anthropic_staff_api_key                    .as_ref()                    .context("no Anthropic AI staff API key configured on the server")?            } else {                state                    .config                    .anthropic_api_key                    .as_ref()                    .context("no Anthropic AI API key configured on the server")?            };             let mut request: anthropic::Request =                serde_json::from_str(params.provider_request.get())?;             // Override the model on the request with the latest version of the model that is            // known to the server.            //            // Right now, we use the version that's defined in `model.id()`, but we will likely            // want to change this code once a new version of an Anthropic model is released,            // so that users can use the new version, without having to update Zed.            request.model = match model.as_str() {                "claude-3-5-sonnet" => anthropic::Model::Claude3_5Sonnet.id().to_string(),                "claude-3-opus" => anthropic::Model::Claude3Opus.id().to_string(),                "claude-3-haiku" => anthropic::Model::Claude3Haiku.id().to_string(),                "claude-3-sonnet" => anthropic::Model::Claude3Sonnet.id().to_string(),                _ => request.model,            };             let (chunks, rate_limit_info) = anthropic::stream_completion_with_rate_limit_info(                &state.http_client,                anthropic::ANTHROPIC_API_URL,                api_key,                request,            )            .await            .map_err(|err| match err {                anthropic::AnthropicError::ApiError(ref api_error) => match api_error.code() {                    Some(anthropic::ApiErrorCode::RateLimitError) => {                        tracing::info!(                            target: "upstream rate limit exceeded",                            user_id = claims.user_id,                            login = claims.github_user_login,                            authn.jti = claims.jti,                            is_staff = claims.is_staff,                            provider = params.provider.to_string(),                            model = model                        );                         Error::http(                            StatusCode::TOO_MANY_REQUESTS,                            "Upstream Anthropic rate limit exceeded.".to_string(),                        )                    }                    Some(anthropic::ApiErrorCode::InvalidRequestError) => {                        Error::http(StatusCode::BAD_REQUEST, api_error.message.clone())                    }                    Some(anthropic::ApiErrorCode::OverloadedError) => {                        Error::http(StatusCode::SERVICE_UNAVAILABLE, api_error.message.clone())                    }                    Some(_) => {                        Error::http(StatusCode::INTERNAL_SERVER_ERROR, api_error.message.clone())                    }                    None => Error::Internal(anyhow!(err)),                },                anthropic::AnthropicError::Other(err) => Error::Internal(err),            })?;             if let Some(rate_limit_info) = rate_limit_info {                tracing::info!(                    target: "upstream rate limit",                    is_staff = claims.is_staff,                    provider = params.provider.to_string(),                    model = model,                    tokens_remaining = rate_limit_info.tokens_remaining,                    requests_remaining = rate_limit_info.requests_remaining,                    requests_reset = ?rate_limit_info.requests_reset,                    tokens_reset = ?rate_limit_info.tokens_reset,                );            }             chunks                .map(move |event| {                    let chunk = event?;                    let (                        input_tokens,                        output_tokens,                        cache_creation_input_tokens,                        cache_read_input_tokens,                    ) = match &chunk {                        anthropic::Event::MessageStart {                            message: anthropic::Response { usage, .. },                        }                        | anthropic::Event::MessageDelta { usage, .. } => (                            usage.input_tokens.unwrap_or(0) as usize,                            usage.output_tokens.unwrap_or(0) as usize,                            usage.cache_creation_input_tokens.unwrap_or(0) as usize,                            usage.cache_read_input_tokens.unwrap_or(0) as usize,                        ),                        _ => (0, 0, 0, 0),                    };                     anyhow::Ok(CompletionChunk {                        bytes: serde_json::to_vec(&chunk).unwrap(),                        input_tokens,                        output_tokens,                        cache_creation_input_tokens,                        cache_read_input_tokens,                    })                })                .boxed()        }        LanguageModelProvider::OpenAi => {            let api_key = state                .config                .openai_api_key                .as_ref()                .context("no OpenAI API key configured on the server")?;            let chunks = open_ai::stream_completion(                &state.http_client,                open_ai::OPEN_AI_API_URL,                api_key,                serde_json::from_str(params.provider_request.get())?,            )            .await?;             chunks                .map(|event| {                    event.map(|chunk| {                        let input_tokens =                            chunk.usage.as_ref().map_or(0, |u| u.prompt_tokens) as usize;                        let output_tokens =                            chunk.usage.as_ref().map_or(0, |u| u.completion_tokens) as usize;                        CompletionChunk {                            bytes: serde_json::to_vec(&chunk).unwrap(),                            input_tokens,                            output_tokens,                            cache_creation_input_tokens: 0,                            cache_read_input_tokens: 0,                        }                    })                })                .boxed()        }        LanguageModelProvider::Google => {            let api_key = state                .config                .google_ai_api_key                .as_ref()                .context("no Google AI API key configured on the server")?;            let chunks = google_ai::stream_generate_content(                &state.http_client,                google_ai::API_URL,                api_key,                serde_json::from_str(params.provider_request.get())?,            )            .await?;             chunks                .map(|event| {                    event.map(|chunk| {                        // TODO - implement token counting for Google AI                        CompletionChunk {                            bytes: serde_json::to_vec(&chunk).unwrap(),                            input_tokens: 0,                            output_tokens: 0,                            cache_creation_input_tokens: 0,                            cache_read_input_tokens: 0,                        }                    })                })                .boxed()        }    };     Ok(Response::new(Body::wrap_stream(TokenCountingStream {        state,        claims,        provider: params.provider,        model,        tokens: TokenUsage::default(),        inner_stream: stream,    })))} fn normalize_model_name(known_models: Vec<String>, name: String) -> String {    if let Some(known_model_name) = known_models        .iter()        .filter(|known_model_name| name.starts_with(known_model_name.as_str()))        .max_by_key(|known_model_name| known_model_name.len())    {        known_model_name.to_string()    } else {        name    }} async fn predict_edits(    Extension(state): Extension<Arc<LlmState>>,    Extension(claims): Extension<LlmTokenClaims>,    _country_code_header: Option<TypedHeader<CloudflareIpCountryHeader>>,    Json(params): Json<PredictEditsParams>,) -> Result<impl IntoResponse> {    if !claims.is_staff {        return Err(anyhow!("not found"))?;    }     let api_url = state        .config        .prediction_api_url        .as_ref()        .context("no PREDICTION_API_URL configured on the server")?;    let api_key = state        .config        .prediction_api_key        .as_ref()        .context("no PREDICTION_API_KEY configured on the server")?;    let model = state        .config        .prediction_model        .as_ref()        .context("no PREDICTION_MODEL configured on the server")?;    let prompt = include_str!("./llm/prediction_prompt.md")        .replace("<events>", &params.input_events)        .replace("<excerpt>", &params.input_excerpt);    let mut response = open_ai::complete_text(        &state.http_client,        api_url,        api_key,        open_ai::CompletionRequest {            model: model.to_string(),            prompt: prompt.clone(),            max_tokens: 1024,            temperature: 0.,            prediction: Some(open_ai::Prediction::Content {                content: params.input_excerpt,            }),            rewrite_speculation: Some(true),        },    )    .await?;    let choice = response        .choices        .pop()        .context("no output from completion response")?;    Ok(Json(PredictEditsResponse {        output_excerpt: choice.text,    }))} /// The maximum monthly spending an individual user can reach on the free tier/// before they have to pay.pub const FREE_TIER_MONTHLY_SPENDING_LIMIT: Cents = Cents::from_dollars(10); /// The default value to use for maximum spend per month if the user did not/// explicitly set a maximum spend.////// Used to prevent surprise bills.pub const DEFAULT_MAX_MONTHLY_SPEND: Cents = Cents::from_dollars(10); async fn check_usage_limit(

Lines 1–500 of 780. Load more