Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

collab: Add support for a custom monthly allowance for LLM usage (#19525)

Browse source 
This PR adds support for setting a monthly LLM usage allowance for
certain users.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-10-21 17:12:33 -04:00 • committed by GitHub
parent 89f6b65ee6
commit 1a4b253ee5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 47 additions and 21 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/collab/src/llm.rs
View file

@ -459,8 +459,9 @@ async fn check_usage_limit(
Utc::now(),
)
.await?;
let free_tier = claims.free_tier_monthly_spending_limit();
if usage.spending_this_month >= FREE_TIER_MONTHLY_SPENDING_LIMIT {
if usage.spending_this_month >= free_tier {
if !claims.has_llm_subscription {
return Err(Error::http(
StatusCode::PAYMENT_REQUIRED,
@ -468,9 +469,7 @@ async fn check_usage_limit(
));
}
if (usage.spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT)
>= Cents(claims.max_monthly_spend_in_cents)
{
if (usage.spending_this_month - free_tier) >= Cents(claims.max_monthly_spend_in_cents) {
return Err(Error::Http(
StatusCode::FORBIDDEN,
"Maximum spending limit reached for this month.".to_string(),
@ -640,6 +639,7 @@ impl<S> Drop for TokenCountingStream<S> {
tokens,
claims.has_llm_subscription,
Cents(claims.max_monthly_spend_in_cents),
claims.free_tier_monthly_spending_limit(),
Utc::now(),
)
.await