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::{FutureExt, 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 && !claims.has_predict_edits_feature_flag { return Err(Error::http( StatusCode::FORBIDDEN, "no access to Zed's edit prediction feature".to_string(), )); } 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 outline_prefix = params .outline .as_ref() .map(|outline| format!("### Outline for current file:\n{}\n", outline)) .unwrap_or_default(); let prompt = include_str!("./llm/prediction_prompt.md") .replace("<outline>", &outline_prefix) .replace("<events>", ¶ms.input_events) .replace("<excerpt>", ¶ms.input_excerpt); let request_start = std::time::Instant::now(); let timeout = state .executor .sleep(std::time::Duration::from_secs(2)) .fuse(); let response = fireworks::complete( &state.http_client, api_url, api_key, fireworks::CompletionRequest { model: model.to_string(), prompt: prompt.clone(), max_tokens: 2048, temperature: 0., prediction: Some(fireworks::Prediction::Content { content: params.input_excerpt, }), rewrite_speculation: Some(true), }, ) .fuse(); futures::pin_mut!(timeout); futures::pin_mut!(response); Lines 1–500 of 855. Load more
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::{FutureExt, 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 && !claims.has_predict_edits_feature_flag { return Err(Error::http( StatusCode::FORBIDDEN, "no access to Zed's edit prediction feature".to_string(), )); } 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 outline_prefix = params .outline .as_ref() .map(|outline| format!("### Outline for current file:\n{}\n", outline)) .unwrap_or_default(); let prompt = include_str!("./llm/prediction_prompt.md") .replace("<outline>", &outline_prefix) .replace("<events>", ¶ms.input_events) .replace("<excerpt>", ¶ms.input_excerpt); let request_start = std::time::Instant::now(); let timeout = state .executor .sleep(std::time::Duration::from_secs(2)) .fuse(); let response = fireworks::complete( &state.http_client, api_url, api_key, fireworks::CompletionRequest { model: model.to_string(), prompt: prompt.clone(), max_tokens: 2048, temperature: 0., prediction: Some(fireworks::Prediction::Content { content: params.input_excerpt, }), rewrite_speculation: Some(true), }, ) .fuse(); futures::pin_mut!(timeout); futures::pin_mut!(response);