mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +08:00
perf(usage): heatmap 改为数据库端按天聚合查询
将 admin 和 user heatmap 从逐条加载 usage audit 记录后在应用层聚合, 改为通过 SQL GROUP BY DATE 在数据库端直接按天汇总, 大幅减少数据传输量。 新增 UsageDailyHeatmapQuery / StoredUsageDailySummary 类型, 在 trait、SQL、内存实现中均补齐 summarize_usage_daily_heatmap 方法。 同时优化 docker-compose: postgres 增加空闲事务超时与 keepalive 参数, gateway 增加健康检查配置。
This commit is contained in:
@@ -1,15 +1,17 @@
|
||||
use crate::handlers::admin::request::AdminAppState;
|
||||
use crate::GatewayError;
|
||||
use aether_admin::observability::stats::round_to;
|
||||
use aether_admin::observability::usage::{
|
||||
admin_usage_data_unavailable_response, admin_usage_heatmap_json,
|
||||
ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
|
||||
admin_usage_data_unavailable_response, ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
|
||||
};
|
||||
use aether_data_contracts::repository::usage::UsageAuditListQuery;
|
||||
use aether_data_contracts::repository::usage::UsageDailyHeatmapQuery;
|
||||
use axum::{
|
||||
body::Body,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use std::collections::BTreeMap;
|
||||
|
||||
pub(super) async fn build_admin_usage_heatmap_response(
|
||||
state: &AdminAppState<'_>,
|
||||
@@ -19,14 +21,64 @@ pub(super) async fn build_admin_usage_heatmap_response(
|
||||
ADMIN_USAGE_DATA_UNAVAILABLE_DETAIL,
|
||||
));
|
||||
}
|
||||
let now_unix_secs = u64::try_from(chrono::Utc::now().timestamp()).unwrap_or_default();
|
||||
let created_from_unix_secs = now_unix_secs.saturating_sub(365 * 24 * 3600);
|
||||
let mut usage = state
|
||||
.list_usage_audits(&UsageAuditListQuery {
|
||||
created_from_unix_secs: Some(created_from_unix_secs),
|
||||
..Default::default()
|
||||
let today = chrono::Utc::now().date_naive();
|
||||
let start_date = today
|
||||
.checked_sub_signed(chrono::Duration::days(364))
|
||||
.unwrap_or(today);
|
||||
let created_from_unix_secs = u64::try_from(
|
||||
chrono::DateTime::<chrono::Utc>::from_naive_utc_and_offset(
|
||||
start_date.and_hms_opt(0, 0, 0).unwrap_or_default(),
|
||||
chrono::Utc,
|
||||
)
|
||||
.timestamp(),
|
||||
)
|
||||
.unwrap_or_default();
|
||||
|
||||
let summaries = state
|
||||
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
|
||||
created_from_unix_secs,
|
||||
user_id: None,
|
||||
admin_mode: true,
|
||||
})
|
||||
.await?;
|
||||
usage.retain(|item| item.status != "pending" && item.status != "streaming");
|
||||
Ok(Json(admin_usage_heatmap_json(&usage)).into_response())
|
||||
|
||||
let grouped: BTreeMap<String, _> = summaries.into_iter().map(|s| (s.date.clone(), s)).collect();
|
||||
|
||||
let mut max_requests = 0_u64;
|
||||
let mut cursor = start_date;
|
||||
let mut days = Vec::new();
|
||||
while cursor <= today {
|
||||
let date_str = cursor.to_string();
|
||||
let (requests, total_tokens, total_cost, actual_total_cost) =
|
||||
if let Some(s) = grouped.get(&date_str) {
|
||||
(
|
||||
s.requests,
|
||||
s.total_tokens,
|
||||
s.total_cost_usd,
|
||||
s.actual_total_cost_usd,
|
||||
)
|
||||
} else {
|
||||
(0, 0, 0.0, 0.0)
|
||||
};
|
||||
max_requests = max_requests.max(requests);
|
||||
days.push(json!({
|
||||
"date": date_str,
|
||||
"requests": requests,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": round_to(total_cost, 6),
|
||||
"actual_total_cost": round_to(actual_total_cost, 6),
|
||||
}));
|
||||
cursor = cursor
|
||||
.checked_add_signed(chrono::Duration::days(1))
|
||||
.unwrap_or(today + chrono::Duration::days(1));
|
||||
}
|
||||
|
||||
Ok(Json(json!({
|
||||
"start_date": start_date.to_string(),
|
||||
"end_date": today.to_string(),
|
||||
"total_days": days.len(),
|
||||
"max_requests": max_requests,
|
||||
"days": days,
|
||||
}))
|
||||
.into_response())
|
||||
}
|
||||
|
||||
@@ -40,6 +40,14 @@ impl<'a> AdminAppState<'a> {
|
||||
self.app.list_usage_audits(query).await
|
||||
}
|
||||
|
||||
pub(crate) async fn summarize_usage_daily_heatmap(
|
||||
&self,
|
||||
query: &aether_data_contracts::repository::usage::UsageDailyHeatmapQuery,
|
||||
) -> Result<Vec<aether_data_contracts::repository::usage::StoredUsageDailySummary>, GatewayError>
|
||||
{
|
||||
self.app.summarize_usage_daily_heatmap(query).await
|
||||
}
|
||||
|
||||
pub(crate) async fn find_request_usage_by_id(
|
||||
&self,
|
||||
usage_id: &str,
|
||||
|
||||
@@ -1026,14 +1026,14 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
|
||||
)
|
||||
.unwrap_or_default();
|
||||
|
||||
let items = match state
|
||||
.list_usage_audits(&UsageAuditListQuery {
|
||||
created_from_unix_secs: Some(created_from_unix_secs),
|
||||
created_until_unix_secs: None,
|
||||
user_id: Some(auth.user.id.clone()),
|
||||
provider_name: None,
|
||||
model: None,
|
||||
})
|
||||
let summaries = match state
|
||||
.summarize_usage_daily_heatmap(
|
||||
&aether_data_contracts::repository::usage::UsageDailyHeatmapQuery {
|
||||
created_from_unix_secs,
|
||||
user_id: Some(auth.user.id.clone()),
|
||||
admin_mode: false,
|
||||
},
|
||||
)
|
||||
.await
|
||||
{
|
||||
Ok(value) => value,
|
||||
@@ -1047,40 +1047,28 @@ pub(super) async fn handle_users_me_usage_heatmap_get(
|
||||
};
|
||||
|
||||
let include_actual_cost = auth.user.role.eq_ignore_ascii_case("admin");
|
||||
let mut daily = BTreeMap::<chrono::NaiveDate, (u64, u64, f64, f64)>::new();
|
||||
for item in items {
|
||||
if item.billing_status != "settled" || item.total_cost_usd <= 0.0 {
|
||||
continue;
|
||||
}
|
||||
let effective = users_me_usage_effective_unix_secs(&item);
|
||||
let Some(timestamp) = chrono::DateTime::<chrono::Utc>::from_timestamp(
|
||||
i64::try_from(effective).unwrap_or_default(),
|
||||
0,
|
||||
) else {
|
||||
continue;
|
||||
};
|
||||
let entry = daily
|
||||
.entry(timestamp.date_naive())
|
||||
.or_insert((0, 0, 0.0, 0.0));
|
||||
entry.0 = entry.0.saturating_add(1);
|
||||
entry.1 = entry
|
||||
.1
|
||||
.saturating_add(item.total_tokens)
|
||||
.saturating_add(item.cache_creation_input_tokens)
|
||||
.saturating_add(item.cache_read_input_tokens);
|
||||
entry.2 += item.total_cost_usd;
|
||||
entry.3 += item.actual_total_cost_usd;
|
||||
}
|
||||
let grouped: std::collections::HashMap<String, _> =
|
||||
summaries.into_iter().map(|s| (s.date.clone(), s)).collect();
|
||||
|
||||
let mut max_requests = 0_u64;
|
||||
let mut cursor = start_date;
|
||||
let mut days = Vec::new();
|
||||
while cursor <= today {
|
||||
let date_str = cursor.to_string();
|
||||
let (requests, total_tokens, total_cost, actual_total_cost) =
|
||||
daily.get(&cursor).copied().unwrap_or((0, 0, 0.0, 0.0));
|
||||
if let Some(s) = grouped.get(&date_str) {
|
||||
(
|
||||
s.requests,
|
||||
s.total_tokens,
|
||||
s.total_cost_usd,
|
||||
s.actual_total_cost_usd,
|
||||
)
|
||||
} else {
|
||||
(0, 0, 0.0, 0.0)
|
||||
};
|
||||
max_requests = max_requests.max(requests);
|
||||
let mut day = json!({
|
||||
"date": cursor.to_string(),
|
||||
"date": date_str,
|
||||
"requests": requests,
|
||||
"total_tokens": total_tokens,
|
||||
"total_cost": round_to(total_cost, 6),
|
||||
|
||||
Reference in New Issue
Block a user