From 340662e2f7e29463ec1716741158ae1e674b5e02 Mon Sep 17 00:00:00 2001 From: Marshall Bowers Date: Fri, 23 Aug 2024 16:41:16 -0400 Subject: [PATCH] collab: Add lifetime spending limit for LLM usage (#16780) This PR adds a lifetime spending limit on LLM usage. Exceeding this limit will prevent further use of the Zed LLM provider. Currently the cap is $1,000. Release Notes: - N/A --- crates/collab/src/llm.rs | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/crates/collab/src/llm.rs b/crates/collab/src/llm.rs index 8d54008f0dda0c..bab2e92e65975d 100644 --- a/crates/collab/src/llm.rs +++ b/crates/collab/src/llm.rs @@ -411,6 +411,11 @@ fn normalize_model_name(known_models: Vec, name: String) -> String { } } +/// The maximum lifetime spending an individual user can reach before being cut off. +/// +/// Represented in cents. +const LIFETIME_SPENDING_LIMIT_IN_CENTS: usize = 1_000 * 100; + async fn check_usage_limit( state: &Arc, provider: LanguageModelProvider, @@ -428,6 +433,13 @@ async fn check_usage_limit( ) .await?; + if usage.lifetime_spending >= LIFETIME_SPENDING_LIMIT_IN_CENTS { + return Err(Error::http( + StatusCode::FORBIDDEN, + "Maximum spending limit reached.".to_string(), + )); + } + let active_users = state.get_active_user_count(provider, model_name).await?; let users_in_recent_minutes = active_users.users_in_recent_minutes.max(1);