collab: Return current usage by model from `GET /billing/usage` (#29693)

Marshall Bowers created 7 months ago

This PR updates the `GET /billing/usage` endpoint to return the number
of requests made to each model and mode.

Release Notes:

- N/A

Change summary

crates/collab/src/api/billing.rs                              | 27 +++
crates/collab/src/llm/db/queries/subscription_usage_meters.rs | 35 +++++
2 files changed, 62 insertions(+)

Detailed changes

crates/collab/src/api/billing.rs

@@ -1090,9 +1090,17 @@ struct UsageCounts {
     pub remaining: Option<i32>,
 }
 
+#[derive(Debug, Serialize)]
+struct ModelRequestUsage {
+    pub model: String,
+    pub mode: CompletionMode,
+    pub requests: i32,
+}
+
 #[derive(Debug, Serialize)]
 struct GetCurrentUsageResponse {
     pub model_requests: UsageCounts,
+    pub model_request_usage: Vec<ModelRequestUsage>,
     pub edit_predictions: UsageCounts,
 }
 
@@ -1119,6 +1127,7 @@ async fn get_current_usage(
             limit: Some(0),
             remaining: Some(0),
         },
+        model_request_usage: Vec::new(),
         edit_predictions: UsageCounts {
             used: 0,
             limit: Some(0),
@@ -1163,12 +1172,30 @@ async fn get_current_usage(
         zed_llm_client::UsageLimit::Unlimited => None,
     };
 
+    let subscription_usage_meters = llm_db
+        .get_current_subscription_usage_meters_for_user(user.id, Utc::now())
+        .await?;
+
+    let model_request_usage = subscription_usage_meters
+        .into_iter()
+        .filter_map(|(usage_meter, _usage)| {
+            let model = llm_db.model_by_id(usage_meter.model_id).ok()?;
+
+            Some(ModelRequestUsage {
+                model: model.name.clone(),
+                mode: usage_meter.mode,
+                requests: usage_meter.requests,
+            })
+        })
+        .collect::<Vec<_>>();
+
     Ok(Json(GetCurrentUsageResponse {
         model_requests: UsageCounts {
             used: usage.model_requests,
             limit: model_requests_limit,
             remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)),
         },
+        model_request_usage,
         edit_predictions: UsageCounts {
             used: usage.edit_predictions,
             limit: edit_prediction_limit,

crates/collab/src/llm/db/queries/subscription_usage_meters.rs

@@ -1,3 +1,4 @@
+use crate::db::UserId;
 use crate::llm::db::queries::subscription_usages::convert_chrono_to_time;
 
 use super::*;
@@ -34,4 +35,38 @@ impl LlmDatabase {
         })
         .await
     }
+
+    /// Returns all current subscription usage meters for the given user as of the given timestamp.
+    pub async fn get_current_subscription_usage_meters_for_user(
+        &self,
+        user_id: UserId,
+        now: DateTimeUtc,
+    ) -> Result<Vec<(subscription_usage_meter::Model, subscription_usage::Model)>> {
+        let now = convert_chrono_to_time(now)?;
+
+        self.transaction(|tx| async move {
+            let result = subscription_usage_meter::Entity::find()
+                .inner_join(subscription_usage::Entity)
+                .filter(subscription_usage::Column::UserId.eq(user_id))
+                .filter(
+                    subscription_usage::Column::PeriodStartAt
+                        .lte(now)
+                        .and(subscription_usage::Column::PeriodEndAt.gte(now)),
+                )
+                .select_also(subscription_usage::Entity)
+                .all(&*tx)
+                .await?;
+
+            let result = result
+                .into_iter()
+                .filter_map(|(meter, usage)| {
+                    let usage = usage?;
+                    Some((meter, usage))
+                })
+                .collect();
+
+            Ok(result)
+        })
+        .await
+    }
 }