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

3
crates/collab/migrations.sqlite/20221109000000_test_schema.sql
View file

@ -11,7 +11,8 @@ CREATE TABLE "users" (
"metrics_id" TEXT,
"github_user_id" INTEGER NOT NULL,
"accepted_tos_at" TIMESTAMP WITHOUT TIME ZONE,
"github_user_created_at" TIMESTAMP WITHOUT TIME ZONE
"github_user_created_at" TIMESTAMP WITHOUT TIME ZONE,
"custom_llm_monthly_allowance_in_cents" INTEGER
);
CREATE UNIQUE INDEX "index_users_github_login" ON "users" ("github_login");
CREATE UNIQUE INDEX "index_invite_code_users" ON "users" ("invite_code");

1
crates/collab/migrations/20241021202606_add_custom_llm_monthly_allowance_in_cents_to_users.sql Normal file
View file

@ -0,0 +1 @@
alter table users add column custom_llm_monthly_allowance_in_cents integer;

9
crates/collab/src/api/billing.rs
View file

@ -34,7 +34,7 @@ use crate::{
db::{billing_subscription::StripeSubscriptionStatus, UserId},
llm::db::LlmDatabase,
};
use crate::{AppState, Error, Result};
use crate::{AppState, Cents, Error, Result};
pub fn router() -> Router {
Router::new()
@ -700,10 +700,15 @@ async fn get_monthly_spend(
));
};
let free_tier = user
.custom_llm_monthly_allowance_in_cents
.map(|allowance| Cents(allowance as u32))
.unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT);
let monthly_spend = llm_db
.get_user_spending_for_month(user.id, Utc::now())
.await?
.saturating_sub(FREE_TIER_MONTHLY_SPENDING_LIMIT);
.saturating_sub(free_tier);
Ok(Json(GetMonthlySpendResponse {
monthly_spend_in_cents: monthly_spend.0 as i32,

1
crates/collab/src/db/tables/user.rs
View file

@ -21,6 +21,7 @@ pub struct Model {
pub metrics_id: Uuid,
pub created_at: NaiveDateTime,
pub accepted_tos_at: Option<NaiveDateTime>,
pub custom_llm_monthly_allowance_in_cents: Option<i32>,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]

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

7
crates/collab/src/llm/db/queries/usages.rs
View file

@ -1,5 +1,5 @@
use crate::db::UserId;
use crate::llm::Cents;
use crate::{db::UserId, llm::FREE_TIER_MONTHLY_SPENDING_LIMIT};
use chrono::{Datelike, Duration};
use futures::StreamExt as _;
use rpc::LanguageModelProvider;
@ -299,6 +299,7 @@ impl LlmDatabase {
tokens: TokenUsage,
has_llm_subscription: bool,
max_monthly_spend: Cents,
free_tier_monthly_spending_limit: Cents,
now: DateTimeUtc,
) -> Result<Usage> {
self.transaction(|tx| async move {
@ -410,9 +411,9 @@ impl LlmDatabase {
);
if !is_staff
&& spending_this_month > FREE_TIER_MONTHLY_SPENDING_LIMIT
&& spending_this_month > free_tier_monthly_spending_limit
&& has_llm_subscription
&& (spending_this_month - FREE_TIER_MONTHLY_SPENDING_LIMIT) <= max_monthly_spend
&& (spending_this_month - free_tier_monthly_spending_limit) <= max_monthly_spend
{
billing_event::ActiveModel {
id: ActiveValue::not_set(),

3
crates/collab/src/llm/db/tests/billing_tests.rs
View file

@ -66,6 +66,7 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
usage,
true,
max_monthly_spend,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -103,6 +104,7 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
usage_2,
true,
max_monthly_spend,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -132,6 +134,7 @@ async fn test_billing_limit_exceeded(db: &mut LlmDatabase) {
model,
usage_exceeding,
true,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
max_monthly_spend,
now,
)

7
crates/collab/src/llm/db/tests/usage_tests.rs
View file

@ -1,3 +1,4 @@
use crate::llm::FREE_TIER_MONTHLY_SPENDING_LIMIT;
use crate::{
db::UserId,
llm::db::{
@ -49,6 +50,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
},
false,
Cents::ZERO,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -68,6 +70,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
},
false,
Cents::ZERO,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -124,6 +127,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
},
false,
Cents::ZERO,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -180,6 +184,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
},
false,
Cents::ZERO,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -222,6 +227,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
},
false,
Cents::ZERO,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await
@ -259,6 +265,7 @@ async fn test_tracking_usage(db: &mut LlmDatabase) {
},
false,
Cents::ZERO,
FREE_TIER_MONTHLY_SPENDING_LIMIT,
now,
)
.await

26
crates/collab/src/llm/token.rs
View file

@ -1,8 +1,7 @@
use crate::llm::DEFAULT_MAX_MONTHLY_SPEND;
use crate::{
db::{billing_preference, UserId},
Config,
};
use crate::db::user;
use crate::llm::{DEFAULT_MAX_MONTHLY_SPEND, FREE_TIER_MONTHLY_SPENDING_LIMIT};
use crate::Cents;
use crate::{db::billing_preference, Config};
use anyhow::{anyhow, Result};
use chrono::Utc;
use jsonwebtoken::{DecodingKey, EncodingKey, Header, Validation};
@ -22,6 +21,7 @@ pub struct LlmTokenClaims {
pub has_llm_closed_beta_feature_flag: bool,
pub has_llm_subscription: bool,
pub max_monthly_spend_in_cents: u32,
pub custom_llm_monthly_allowance_in_cents: Option<u32>,
pub plan: rpc::proto::Plan,
}
@ -30,8 +30,7 @@ const LLM_TOKEN_LIFETIME: Duration = Duration::from_secs(60 * 60);
impl LlmTokenClaims {
#[allow(clippy::too_many_arguments)]
pub fn create(
user_id: UserId,
github_user_login: String,
user: &user::Model,
is_staff: bool,
billing_preferences: Option<billing_preference::Model>,
has_llm_closed_beta_feature_flag: bool,
@ -49,8 +48,8 @@ impl LlmTokenClaims {
iat: now.timestamp() as u64,
exp: (now + LLM_TOKEN_LIFETIME).timestamp() as u64,
jti: uuid::Uuid::new_v4().to_string(),
user_id: user_id.to_proto(),
github_user_login,
user_id: user.id.to_proto(),
github_user_login: user.github_login.clone(),
is_staff,
has_llm_closed_beta_feature_flag,
has_llm_subscription,
@ -58,6 +57,9 @@ impl LlmTokenClaims {
.map_or(DEFAULT_MAX_MONTHLY_SPEND.0, |preferences| {
preferences.max_monthly_llm_usage_spending_in_cents as u32
}),
custom_llm_monthly_allowance_in_cents: user
.custom_llm_monthly_allowance_in_cents
.map(|allowance| allowance as u32),
plan,
};
@ -89,6 +91,12 @@ impl LlmTokenClaims {
}
}
}
pub fn free_tier_monthly_spending_limit(&self) -> Cents {
self.custom_llm_monthly_allowance_in_cents
.map(Cents)
.unwrap_or(FREE_TIER_MONTHLY_SPENDING_LIMIT)
}
}
#[derive(Error, Debug)]

3
crates/collab/src/rpc.rs
View file

@ -4930,8 +4930,7 @@ async fn get_llm_api_token(
let billing_preferences = db.get_billing_preferences(user.id).await?;
let token = LlmTokenClaims::create(
user.id,
user.github_login.clone(),
&user,
session.is_staff(),
billing_preferences,
has_llm_closed_beta_feature_flag,