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

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(),