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:
parent
50f705e779
commit
399eced884
@ -1090,9 +1090,17 @@ struct UsageCounts {
|
|||||||
pub remaining: Option<i32>,
|
pub remaining: Option<i32>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Serialize)]
|
||||||
|
struct ModelRequestUsage {
|
||||||
|
pub model: String,
|
||||||
|
pub mode: CompletionMode,
|
||||||
|
pub requests: i32,
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Serialize)]
|
#[derive(Debug, Serialize)]
|
||||||
struct GetCurrentUsageResponse {
|
struct GetCurrentUsageResponse {
|
||||||
pub model_requests: UsageCounts,
|
pub model_requests: UsageCounts,
|
||||||
|
pub model_request_usage: Vec<ModelRequestUsage>,
|
||||||
pub edit_predictions: UsageCounts,
|
pub edit_predictions: UsageCounts,
|
||||||
}
|
}
|
||||||
|
|
||||||
@ -1119,6 +1127,7 @@ async fn get_current_usage(
|
|||||||
limit: Some(0),
|
limit: Some(0),
|
||||||
remaining: Some(0),
|
remaining: Some(0),
|
||||||
},
|
},
|
||||||
|
model_request_usage: Vec::new(),
|
||||||
edit_predictions: UsageCounts {
|
edit_predictions: UsageCounts {
|
||||||
used: 0,
|
used: 0,
|
||||||
limit: Some(0),
|
limit: Some(0),
|
||||||
@ -1163,12 +1172,30 @@ async fn get_current_usage(
|
|||||||
zed_llm_client::UsageLimit::Unlimited => None,
|
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 {
|
Ok(Json(GetCurrentUsageResponse {
|
||||||
model_requests: UsageCounts {
|
model_requests: UsageCounts {
|
||||||
used: usage.model_requests,
|
used: usage.model_requests,
|
||||||
limit: model_requests_limit,
|
limit: model_requests_limit,
|
||||||
remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)),
|
remaining: model_requests_limit.map(|limit| (limit - usage.model_requests).max(0)),
|
||||||
},
|
},
|
||||||
|
model_request_usage,
|
||||||
edit_predictions: UsageCounts {
|
edit_predictions: UsageCounts {
|
||||||
used: usage.edit_predictions,
|
used: usage.edit_predictions,
|
||||||
limit: edit_prediction_limit,
|
limit: edit_prediction_limit,
|
||||||
|
@ -1,3 +1,4 @@
|
|||||||
|
use crate::db::UserId;
|
||||||
use crate::llm::db::queries::subscription_usages::convert_chrono_to_time;
|
use crate::llm::db::queries::subscription_usages::convert_chrono_to_time;
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
@ -34,4 +35,38 @@ impl LlmDatabase {
|
|||||||
})
|
})
|
||||||
.await
|
.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
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
Loading…
Reference in New Issue
Block a user