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

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

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2025-04-30 15:06:39 -04:00 committed by GitHub
parent 50f705e779
commit 399eced884
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 62 additions and 0 deletions

View file

@ -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,