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
This commit is contained in:
Marshall Bowers 2024-08-23 16:41:16 -04:00 committed by GitHub
parent 77bb60f1d1
commit 340662e2f7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -411,6 +411,11 @@ fn normalize_model_name(known_models: Vec<String>, 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<LlmState>,
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);