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

collab: Update billing code for LLM usage billing (#18879)

Browse source 
This PR reworks our existing billing code in preparation for charging
based on LLM usage.

We aren't yet exercising the new billing-related code outside of
development.

There are some noteworthy changes for our existing LLM usage tracking:

- A new `monthly_usages` table has been added for tracking usage
per-user, per-model, per-month
- The per-month usage measures have been removed, in favor of the
`monthly_usages` table
- All of the per-month metrics in the Clickhouse rows have been changed
from a rolling 30-day window to a calendar month

Release Notes:

- N/A

---------

Co-authored-by: Antonio Scandurra <me@as-cii.com>
Co-authored-by: Richard <richard@zed.dev>
Co-authored-by: Max <max@zed.dev>
This commit is contained in:
Marshall Bowers 2024-10-08 18:29:38 -04:00 • committed by GitHub
parent a95fb8f1f9
commit f861479890
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 390 additions and 132 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -191,16 +191,26 @@ impl Session {
}
}
pub async fn current_plan(&self, db: MutexGuard<'_, DbHandle>) -> anyhow::Result<proto::Plan> {
pub async fn has_llm_subscription(
&self,
db: &MutexGuard<'_, DbHandle>,
) -> anyhow::Result<bool> {
if self.is_staff() {
return Ok(proto::Plan::ZedPro);
return Ok(true);
}
let Some(user_id) = self.user_id() else {
return Ok(proto::Plan::Free);
return Ok(false);
};
if db.has_active_billing_subscription(user_id).await? {
Ok(db.has_active_billing_subscription(user_id).await?)
}
pub async fn current_plan(
&self,
_db: &MutexGuard<'_, DbHandle>,
) -> anyhow::Result<proto::Plan> {
if self.is_staff() {
Ok(proto::Plan::ZedPro)
} else {
Ok(proto::Plan::Free)
@ -3471,7 +3481,7 @@ fn should_auto_subscribe_to_channels(version: ZedVersion) -> bool {
}
async fn update_user_plan(_user_id: UserId, session: &Session) -> Result<()> {
let plan = session.current_plan(session.db().await).await?;
let plan = session.current_plan(&session.db().await).await?;
session
.peer
@ -4471,7 +4481,7 @@ async fn count_language_model_tokens(
};
authorize_access_to_legacy_llm_endpoints(&session).await?;
let rate_limit: Box<dyn RateLimit> = match session.current_plan(session.db().await).await? {
let rate_limit: Box<dyn RateLimit> = match session.current_plan(&session.db().await).await? {
proto::Plan::ZedPro => Box::new(ZedProCountLanguageModelTokensRateLimit),
proto::Plan::Free => Box::new(FreeCountLanguageModelTokensRateLimit),
};
@ -4592,7 +4602,7 @@ async fn compute_embeddings(
let api_key = api_key.context("no OpenAI API key configured on the server")?;
authorize_access_to_legacy_llm_endpoints(&session).await?;
let rate_limit: Box<dyn RateLimit> = match session.current_plan(session.db().await).await? {
let rate_limit: Box<dyn RateLimit> = match session.current_plan(&session.db().await).await? {
proto::Plan::ZedPro => Box::new(ZedProComputeEmbeddingsRateLimit),
proto::Plan::Free => Box::new(FreeComputeEmbeddingsRateLimit),
};
@ -4915,7 +4925,8 @@ async fn get_llm_api_token(
user.github_login.clone(),
session.is_staff(),
has_llm_closed_beta_feature_flag,
session.current_plan(db).await?,
session.has_llm_subscription(&db).await?,
session.current_plan(&db).await?,
&session.app_state.config,
)?;
response.send(proto::GetLlmTokenResponse { token })?;