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

Preserve usage provider identity

Browse Source
This commit is contained in:
fawney19
2026-05-20 20:56:36 +08:00
parent c972bbd397
commit 57910f906d
21 changed files with 340 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/aether-data/src/repository/usage/memory.rs
+80 -23
View File
@@ -983,6 +983,30 @@ fn usage_provider_display_name(item: &StoredRequestUsageAudit) -> Option<String>
}
}
fn usage_provider_id(item: &StoredRequestUsageAudit) -> Option<String> {
let provider_id = item.provider_id.as_deref()?.trim();
if provider_id.is_empty() || usage_reserved_provider_label(provider_id) {
None
} else {
Some(provider_id.to_string())
}
}
fn usage_provider_aggregation_identity(
item: &StoredRequestUsageAudit,
) -> Option<(String, Option<String>, String)> {
let display_name = usage_provider_display_name(item);
if let Some(provider_id) = usage_provider_id(item) {
return Some((provider_id, display_name, "provider_id".to_string()));
}
let display_name = display_name?;
Some((
display_name.clone(),
Some(display_name),
"legacy_name".to_string(),
))
}
#[async_trait]
impl UsageReadRepository for InMemoryUsageReadRepository {
async fn find_by_id(
@@ -1167,14 +1191,14 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|| 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())
&& usage_provider_aggregation_identity(item).is_none())
{
continue;
}
let provider_display_name =
let provider_identity =
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider) {
match usage_provider_display_name(item) {
match usage_provider_aggregation_identity(item) {
Some(value) => Some(value),
None => continue,
}
@@ -1184,19 +1208,11 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
let group_key = match query.group_by {
UsageAuditAggregationGroupBy::Model => item.model.clone(),
UsageAuditAggregationGroupBy::Provider => {
let display_name = provider_display_name
.as_deref()
.expect("provider display name is set for provider aggregation");
item.provider_id
.as_deref()
.map(str::trim)
.filter(|provider_id| {
!provider_id.is_empty() && !usage_reserved_provider_label(provider_id)
})
.unwrap_or(display_name)
.to_string()
}
UsageAuditAggregationGroupBy::Provider => provider_identity
.as_ref()
.expect("provider identity is set for provider aggregation")
.0
.clone(),
UsageAuditAggregationGroupBy::ApiFormat => item
.api_format
.clone()
@@ -1211,7 +1227,17 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
&& (bucket.display_name.is_none()
|| bucket.display_name.as_deref() == Some("Unknown"))
{
bucket.display_name = provider_display_name;
bucket.display_name = provider_identity
.as_ref()
.and_then(|(_, display_name, _)| display_name.clone());
}
if matches!(query.group_by, UsageAuditAggregationGroupBy::Provider)
&& (bucket.secondary_name.is_none()
|| bucket.secondary_name.as_deref() == Some("legacy_name"))
{
bucket.secondary_name = provider_identity
.as_ref()
.map(|(_, _, identity_source)| identity_source.clone());
}
bucket.request_count = bucket.request_count.saturating_add(1);
bucket.total_tokens = bucket.total_tokens.saturating_add(item.total_tokens);
@@ -3149,6 +3175,12 @@ mod tests {
#[tokio::test]
async fn provider_aggregation_skips_unknown_provider_labels() {
let valid_provider = sample_usage("req-valid-provider", 300);
let mut legacy_provider = sample_usage("req-legacy-provider", 250);
legacy_provider.provider_id = None;
legacy_provider.provider_name = "Legacy Provider".to_string();
let mut unknown = sample_usage("req-unknown-provider", 100);
unknown.provider_id = None;
unknown.provider_name = "unknown".to_string();
@@ -3158,7 +3190,8 @@ mod tests {
typo_unknown.provider_name = "unknow".to_string();
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-valid-provider", 300),
valid_provider,
legacy_provider,
unknown,
typo_unknown,
]);
@@ -3174,9 +3207,29 @@ mod tests {
.await
.expect("aggregation should succeed");
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].group_key, "provider-1");
assert_eq!(rows[0].display_name.as_deref(), Some("OpenAI"));
assert_eq!(rows.len(), 2);
let provider_id_row = rows
.iter()
.find(|row| row.group_key == "provider-1")
.expect("provider_id row should be present");
assert_eq!(provider_id_row.display_name.as_deref(), Some("OpenAI"));
assert_eq!(
provider_id_row.secondary_name.as_deref(),
Some("provider_id")
);
let legacy_name_row = rows
.iter()
.find(|row| row.group_key == "Legacy Provider")
.expect("legacy provider name row should be present");
assert_eq!(
legacy_name_row.display_name.as_deref(),
Some("Legacy Provider")
);
assert_eq!(
legacy_name_row.secondary_name.as_deref(),
Some("legacy_name")
);
}
#[tokio::test]
@@ -3193,11 +3246,15 @@ mod tests {
pending_provider.provider_id = None;
pending_provider.provider_name = "pending".to_string();
let mut id_only_provider = sample_usage("req-id-only-provider", 350);
id_only_provider.provider_name = "unknown".to_string();
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-valid-provider", 400),
unknown,
typo_unknown,
pending_provider,
id_only_provider,
]);
let model_rows = repository
@@ -3212,7 +3269,7 @@ mod tests {
.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);
assert_eq!(model_rows[0].request_count, 2);
let api_format_rows = repository
.aggregate_usage_audits(&UsageAuditAggregationQuery {
@@ -3226,7 +3283,7 @@ mod tests {
.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);
assert_eq!(api_format_rows[0].request_count, 2);
}
#[tokio::test]