i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-09-02 01:10:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

Exclude unknown providers from usage aggregations

Browse Source
This commit is contained in:
fawney19
2026-05-06 13:40:18 +08:00
parent 94a6f315ca
commit 6fbb867f5f
7 changed files with 115 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
crates/aether-data-contracts/src/repository/usage/types.rs
View File

@@ -673,6 +673,7 @@ pub struct UsageAuditAggregationQuery {
pub created_until_unix_secs: u64,
pub group_by: UsageAuditAggregationGroupBy,
pub limit: usize,
pub exclude_reserved_provider_labels: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]

53
crates/aether-data/src/repository/usage/memory.rs
View File

@@ -1155,6 +1155,8 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
if item.created_at_unix_ms < query.created_from_unix_secs
|| item.created_at_unix_ms >= query.created_until_unix_secs
|| matches!(item.status.as_str(), "pending" | "streaming")
|| (query.exclude_reserved_provider_labels
&& usage_provider_display_name(item).is_none())
{
continue;
}
@@ -3090,6 +3092,7 @@ mod tests {
created_until_unix_secs: 1_000,
group_by: UsageAuditAggregationGroupBy::Provider,
limit: 10,
exclude_reserved_provider_labels: false,
})
.await
.expect("aggregation should succeed");
@@ -3099,6 +3102,56 @@ mod tests {
assert_eq!(rows[0].display_name.as_deref(), Some("OpenAI"));
}
#[tokio::test]
async fn aggregation_can_skip_unknown_provider_records_for_model_and_api_format() {
let mut unknown = sample_usage("req-unknown-provider", 100);
unknown.provider_id = None;
unknown.provider_name = "unknown".to_string();
let mut typo_unknown = sample_usage("req-unknow-provider", 200);
typo_unknown.provider_id = Some("unknow".to_string());
typo_unknown.provider_name = "unknow".to_string();
let mut pending_provider = sample_usage("req-pending-provider", 300);
pending_provider.provider_id = None;
pending_provider.provider_name = "pending".to_string();
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-valid-provider", 400),
unknown,
typo_unknown,
pending_provider,
]);
let model_rows = repository
.aggregate_usage_audits(&UsageAuditAggregationQuery {
created_from_unix_secs: 0,
created_until_unix_secs: 1_000,
group_by: UsageAuditAggregationGroupBy::Model,
limit: 10,
exclude_reserved_provider_labels: true,
})
.await
.expect("model aggregation should succeed");
assert_eq!(model_rows.len(), 1);
assert_eq!(model_rows[0].group_key, "gpt-4.1");
assert_eq!(model_rows[0].request_count, 1);
let api_format_rows = repository
.aggregate_usage_audits(&UsageAuditAggregationQuery {
created_from_unix_secs: 0,
created_until_unix_secs: 1_000,
group_by: UsageAuditAggregationGroupBy::ApiFormat,
limit: 10,
exclude_reserved_provider_labels: true,
})
.await
.expect("api format aggregation should succeed");
assert_eq!(api_format_rows.len(), 1);
assert_eq!(api_format_rows[0].group_key, "openai:chat");
assert_eq!(api_format_rows[0].request_count, 1);
}
#[tokio::test]
async fn stale_pending_update_does_not_regress_finalized_usage() {
let repository = InMemoryUsageReadRepository::default();

28
crates/aether-data/src/repository/usage/postgres/mod.rs
View File

@@ -1336,6 +1336,8 @@ const REBUILD_PROVIDER_API_KEY_USAGE_STATS_SQL: &str =
include_str!("queries/rebuild_provider_api_key_usage_stats_sql.sql");
const LIST_USAGE_AUDITS_PREFIX: &str = include_str!("queries/list_usage_audits_prefix.sql");
const USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL: &str =
" AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')";
struct UsageAuditAggregationSqlFragments {
filtered_extra_where: &'static str,
@@ -1363,7 +1365,7 @@ fn usage_audit_aggregation_sql_fragments(
success_count_expr: "NULL::BIGINT",
},
UsageAuditAggregationGroupBy::Provider => UsageAuditAggregationSqlFragments {
filtered_extra_where: " AND BTRIM(COALESCE(\"usage\".provider_name, '')) <> '' AND lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')",
filtered_extra_where: "",
group_key_expr: "provider_group_key",
display_name_expr: "provider_display_name",
secondary_name_expr: "NULL::varchar",
@@ -6417,6 +6419,12 @@ WHERE stats_daily_api_key.date >=
}
};
let provider_extra_where = if matches!(group_by, UsageAuditAggregationGroupBy::Provider) {
" AND BTRIM(COALESCE(provider_name, '')) <> '' AND lower(BTRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')"
} else {
""
};
let sql = format!(
r#"
SELECT
@@ -6441,6 +6449,7 @@ SELECT
FROM {table_name}
WHERE date >= $1
AND date < $2
{provider_extra_where}
GROUP BY {group_column}
ORDER BY request_count DESC, group_key ASC
"#,
@@ -6449,6 +6458,7 @@ ORDER BY request_count DESC, group_key ASC
avg_response_time_expr = avg_response_time_expr,
success_count_expr = success_count_expr,
table_name = table_name,
provider_extra_where = provider_extra_where,
);
let mut rows = sqlx::query(&sql)
@@ -6467,6 +6477,13 @@ ORDER BY request_count DESC, group_key ASC
query: &UsageAuditAggregationQuery,
) -> Result<Vec<StoredUsageAuditAggregation>, DataLayerError> {
let fragments = usage_audit_aggregation_sql_fragments(query.group_by);
let provider_extra_where = if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
|| query.exclude_reserved_provider_labels
{
USAGE_RESERVED_PROVIDER_LABELS_FILTER_SQL
} else {
""
};
let sql = format!(
r#"
WITH filtered_usage AS (
@@ -6518,6 +6535,7 @@ WITH filtered_usage AS (
WHERE "usage".created_at >= TO_TIMESTAMP($1::double precision)
AND "usage".created_at < TO_TIMESTAMP($2::double precision)
AND "usage".status NOT IN ('pending', 'streaming')
{provider_extra_where}
{filtered_extra_where}
),
normalized_usage AS (
@@ -6611,6 +6629,7 @@ FROM aggregated_usage
ORDER BY request_count DESC, group_key ASC
LIMIT $3
"#,
provider_extra_where = provider_extra_where,
filtered_extra_where = fragments.filtered_extra_where,
group_key_expr = fragments.group_key_expr,
display_name_expr = fragments.display_name_expr,
@@ -6646,6 +6665,11 @@ LIMIT $3
if matches!(query.group_by, UsageAuditAggregationGroupBy::User) {
return self.aggregate_usage_audits_raw(query).await;
}
if query.exclude_reserved_provider_labels
&& !matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
{
return self.aggregate_usage_audits_raw(query).await;
}
let Some(cutoff_utc) = self.read_stats_daily_cutoff_date().await? else {
return self.aggregate_usage_audits_raw(query).await;
@@ -6666,6 +6690,7 @@ LIMIT $3
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
group_by: query.group_by,
limit: raw_merge_limit,
exclude_reserved_provider_labels: query.exclude_reserved_provider_labels,
})
.await?;
absorb_usage_audit_aggregation_rows(&mut grouped, raw);
@@ -6687,6 +6712,7 @@ LIMIT $3
created_until_unix_secs: dashboard_utc_to_unix_secs(raw_end),
group_by: query.group_by,
limit: raw_merge_limit,
exclude_reserved_provider_labels: query.exclude_reserved_provider_labels,
})
.await?;
absorb_usage_audit_aggregation_rows(&mut grouped, raw);

4
crates/aether-data/src/repository/usage/postgres/tests.rs
View File

@@ -420,6 +420,10 @@ fn usage_sql_provider_aggregation_excludes_unknown_provider_labels() {
assert!(source.contains(
r#"lower(BTRIM(COALESCE(\"usage\".provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')"#
));
assert!(source.contains("exclude_reserved_provider_labels"));
assert!(source.contains(
r#"lower(BTRIM(COALESCE(provider_name, ''))) NOT IN ('unknown', 'unknow', 'pending')"#
));
assert!(source.contains("MAX(display_name)"));
assert!(!source.contains("COALESCE(MAX(NULLIF(display_name, 'Unknown')), 'Unknown')"));
}