Merge pull request #429 from AAEE86/rust

fix(usage): 统一用户排行榜 Token 统计口径
This commit is contained in:
fawney19
2026-05-11 22:46:56 +08:00
committed by GitHub
8 changed files with 133 additions and 54 deletions

View File

@@ -931,6 +931,13 @@ fn usage_effective_input_tokens(item: &StoredRequestUsageAudit) -> u64 {
normalize_usage_input_tokens(api_format, input_tokens, cache_read_tokens) as u64
}
fn usage_total_tokens(item: &StoredRequestUsageAudit) -> u64 {
usage_effective_input_tokens(item)
.saturating_add(item.output_tokens)
.saturating_add(usage_cache_creation_tokens(item))
.saturating_add(item.cache_read_input_tokens)
}
fn usage_is_success(item: &StoredRequestUsageAudit) -> bool {
matches!(
item.status.as_str(),
@@ -1426,7 +1433,9 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
.effective_input_tokens
.saturating_add(usage_effective_input_tokens(item));
summary.output_tokens = summary.output_tokens.saturating_add(item.output_tokens);
summary.total_tokens = summary.total_tokens.saturating_add(item.total_tokens);
summary.total_tokens = summary
.total_tokens
.saturating_add(usage_total_tokens(item));
summary.cache_creation_tokens = summary
.cache_creation_tokens
.saturating_add(usage_cache_creation_tokens(item));
@@ -2187,12 +2196,7 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
entry.legacy_name = legacy_name;
}
entry.request_count = entry.request_count.saturating_add(1);
entry.total_tokens = entry.total_tokens.saturating_add(
item.input_tokens
.saturating_add(item.output_tokens)
.saturating_add(item.cache_creation_input_tokens)
.saturating_add(item.cache_read_input_tokens),
);
entry.total_tokens = entry.total_tokens.saturating_add(usage_total_tokens(item));
entry.total_cost_usd += item.total_cost_usd;
}
Ok(grouped.into_values().collect())
@@ -2989,7 +2993,8 @@ mod tests {
};
use aether_data_contracts::repository::usage::{
usage_body_ref, ProviderApiKeyWindowUsageRequest, UsageAuditAggregationGroupBy,
UsageAuditAggregationQuery, UsageBodyField, UsageProviderPerformanceQuery,
UsageAuditAggregationQuery, UsageBodyField, UsageDashboardSummaryQuery,
UsageLeaderboardGroupBy, UsageLeaderboardQuery, UsageProviderPerformanceQuery,
UsageTimeSeriesGranularity,
};
use serde_json::json;
@@ -4647,6 +4652,46 @@ mod tests {
assert_eq!(items[0].request_id, "req-2");
}
#[tokio::test]
async fn dashboard_and_leaderboard_total_tokens_use_effective_cache_aware_tokens() {
let mut item = sample_usage("req-cache-aware-total", 1_711_000_000);
item.input_tokens = 100;
item.output_tokens = 20;
item.total_tokens = 999;
item.cache_creation_input_tokens = 0;
item.cache_creation_ephemeral_5m_input_tokens = 12;
item.cache_creation_ephemeral_1h_input_tokens = 8;
item.cache_read_input_tokens = 80;
let repository = InMemoryUsageReadRepository::seed(vec![item]);
let dashboard = repository
.summarize_dashboard_usage(&UsageDashboardSummaryQuery {
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
user_id: None,
})
.await
.expect("dashboard should summarize");
assert_eq!(dashboard.effective_input_tokens, 20);
assert_eq!(dashboard.cache_creation_tokens, 20);
assert_eq!(dashboard.total_tokens, 140);
let leaderboard = repository
.summarize_usage_leaderboard(&UsageLeaderboardQuery {
created_from_unix_secs: 1_711_000_000,
created_until_unix_secs: 1_711_000_001,
group_by: UsageLeaderboardGroupBy::User,
user_id: None,
provider_name: None,
model: None,
})
.await
.expect("leaderboard should summarize");
assert_eq!(leaderboard.len(), 1);
assert_eq!(leaderboard[0].total_tokens, 140);
}
#[tokio::test]
async fn summarizes_provider_api_key_last_used_at_in_seconds() {
let repository = InMemoryUsageReadRepository::seed(vec![

View File

@@ -1104,7 +1104,7 @@ mod tests {
.expect("dashboard summary should load");
assert_eq!(summary.total_requests, 2);
assert_eq!(summary.error_requests, 1);
assert_eq!(summary.total_tokens, 14);
assert_eq!(summary.total_tokens, 10);
}
async fn seed_stats_targets(

View File

@@ -6019,9 +6019,30 @@ SELECT
MAX({legacy_name_expr}) AS legacy_name,
COUNT(*)::BIGINT AS request_count,
COALESCE(SUM(
GREATEST(COALESCE("usage".input_tokens, 0), 0)
CASE
WHEN GREATEST(COALESCE("usage".input_tokens, 0), 0) <= 0 THEN 0
WHEN GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0) <= 0
THEN GREATEST(COALESCE("usage".input_tokens, 0), 0)
WHEN split_part(lower(COALESCE(COALESCE("usage".endpoint_api_format, "usage".api_format), '')), ':', 1)
IN ('openai', 'gemini', 'google')
THEN GREATEST(
GREATEST(COALESCE("usage".input_tokens, 0), 0)
- GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0),
0
)
ELSE GREATEST(COALESCE("usage".input_tokens, 0), 0)
END
+ GREATEST(COALESCE("usage".output_tokens, 0), 0)
+ GREATEST(COALESCE("usage".cache_creation_input_tokens, 0), 0)
+ CASE
WHEN COALESCE("usage".cache_creation_input_tokens, 0) = 0
AND (
COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
) > 0
THEN COALESCE("usage".cache_creation_input_tokens_5m, 0)
+ COALESCE("usage".cache_creation_input_tokens_1h, 0)
ELSE COALESCE("usage".cache_creation_input_tokens, 0)
END
+ GREATEST(COALESCE("usage".cache_read_input_tokens, 0), 0)
), 0)::BIGINT AS total_tokens,
COALESCE(SUM(COALESCE(CAST("usage".total_cost_usd AS DOUBLE PRECISION), 0)), 0)
@@ -6169,42 +6190,20 @@ WHERE date >=
fetch_usage_leaderboard_query(builder.build(), &self.pool).await?
}
UsageLeaderboardGroupBy::User => {
let mut builder = if query.provider_name.is_some() && query.model.is_some() {
if query.provider_name.is_some() && query.model.is_some() {
return Ok(None);
}
let mut builder = if let Some(provider_name) = query.provider_name.as_deref() {
let mut builder = QueryBuilder::<Postgres>::new(
r#"
SELECT
user_id AS group_key,
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd
FROM stats_user_daily_model_provider
WHERE date >=
"#,
);
builder
.push_bind(start_day_utc)
.push(" AND date < ")
.push_bind(end_day_utc)
.push(" AND provider_name = ")
.push_bind(query.provider_name.as_deref().unwrap().to_string())
.push(" AND model = ")
.push_bind(query.model.as_deref().unwrap().to_string());
if let Some(user_id) = query.user_id.as_deref() {
builder
.push(" AND user_id = ")
.push_bind(user_id.to_string());
}
builder.push(" GROUP BY user_id ORDER BY user_id ASC");
builder
} else if let Some(provider_name) = query.provider_name.as_deref() {
let mut builder = QueryBuilder::<Postgres>::new(
r#"
SELECT
user_id AS group_key,
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
COALESCE(
SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd
FROM stats_user_daily_provider
WHERE date >=
@@ -6230,7 +6229,10 @@ SELECT
user_id AS group_key,
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(SUM(total_tokens), 0)::BIGINT AS total_tokens,
COALESCE(
SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd
FROM stats_user_daily_model
WHERE date >=
@@ -6257,7 +6259,7 @@ SELECT
MAX(NULLIF(BTRIM(username), '')) AS legacy_name,
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
COALESCE(
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
0
)::BIGINT AS total_tokens,
CAST(COALESCE(SUM(total_cost), 0) AS DOUBLE PRECISION) AS total_cost_usd

View File

@@ -1170,7 +1170,7 @@ ORDER BY request_id
.expect("dashboard summary should load");
assert_eq!(summary.total_requests, 2);
assert_eq!(summary.error_requests, 1);
assert_eq!(summary.total_tokens, 14);
assert_eq!(summary.total_tokens, 10);
}
async fn seed_stats_targets(pool: &sqlx::SqlitePool) {