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

Add usage client family read model field

Browse Source
This commit is contained in:
RWDai
2026-05-18 19:28:36 +08:00
parent 0a26accca4
commit 6d994917a0
6 changed files with 54 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -29,11 +29,12 @@ use serde_json::Value;
use super::{
api_key_usage_contribution, provider_api_key_usage_contribution,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
ApiKeyUsageContribution, ApiKeyUsageDelta, ProviderApiKeyUsageContribution,
ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery,
UsageDailyHeatmapQuery, UsageReadRepository, UsageWriteRepository,
usage_request_metadata_client_family, ApiKeyUsageContribution, ApiKeyUsageDelta,
ProviderApiKeyUsageContribution, ProviderApiKeyUsageDelta, ProviderApiKeyWindowUsageRequest,
StoredProviderApiKeyUsageSummary, StoredProviderApiKeyWindowUsageSummary,
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
StoredUsageDailySummary, UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
UsageReadRepository, UsageWriteRepository,
};
use crate::repository::auth::InMemoryAuthApiKeySnapshotRepository;
use crate::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
@@ -56,6 +57,7 @@ impl InMemoryUsageReadRepository {
let mut by_request_id = BTreeMap::new();
for mut item in items {
hydrate_legacy_body_refs(&mut item);
hydrate_client_family(&mut item);
by_request_id.insert(item.request_id.clone(), item);
}
Self {
@@ -75,6 +77,7 @@ impl InMemoryUsageReadRepository {
let mut detached_bodies = BTreeMap::new();
for mut item in items {
hydrate_legacy_body_refs(&mut item);
hydrate_client_family(&mut item);
let request_id = item.request_id.clone();
if let Some(body_ref) = detach_usage_body(
&request_id,
@@ -2546,6 +2549,13 @@ fn hydrate_legacy_body_refs(item: &mut StoredRequestUsageAudit) {
}
}
fn hydrate_client_family(item: &mut StoredRequestUsageAudit) {
if item.client_family.is_none() {
item.client_family = usage_request_metadata_client_family(item.request_metadata.as_ref())
.map(ToOwned::to_owned);
}
}
fn persisted_usage_body_ref(
incoming_ref: Option<&str>,
incoming_body: Option<&Value>,
@@ -2851,6 +2861,13 @@ impl UsageWriteRepository for InMemoryUsageReadRepository {
})
},
),
client_family: usage_request_metadata_client_family(request_metadata.as_ref())
.map(ToOwned::to_owned)
.or_else(|| {
existing
.as_ref()
.and_then(|existing| existing.client_family.clone())
}),
request_metadata,
created_at_unix_ms,
updated_at_unix_secs: usage.updated_at_unix_secs,
@@ -4399,6 +4416,7 @@ mod tests {
provider_endpoint_kind: Some("chat".to_string()),
has_format_conversion: false,
is_stream: false,
client_family: None,
input_tokens: 10,
output_tokens: 20,
total_tokens: 30,

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

@@ -363,7 +363,8 @@ mod sqlite;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
usage_request_metadata_client_family, PendingUsageCleanupSummary,
ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyUsageSummary,
StoredProviderApiKeyWindowUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, StoredUsageAuditAggregation, StoredUsageAuditSummary,
StoredUsageBreakdownSummaryRow, StoredUsageCacheAffinityHitSummary,

6
crates/aether-data/src/repository/usage/mysql.rs
View File

@@ -6,8 +6,8 @@ use sqlx::{mysql::MySqlRow, Row};
use super::{
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
InMemoryUsageReadRepository, PendingUsageCleanupSummary, StoredRequestUsageAudit,
UpsertUsageRecord, UsageWriteRepository,
usage_request_metadata_client_family, InMemoryUsageReadRepository, PendingUsageCleanupSummary,
StoredRequestUsageAudit, UpsertUsageRecord, UsageWriteRepository,
};
use crate::driver::mysql::MysqlPool;
use crate::error::SqlResultExt;
@@ -850,6 +850,8 @@ fn map_usage_row(row: &MySqlRow) -> Result<StoredRequestUsageAudit, DataLayerErr
.map(|raw| serde_json::from_str(&raw))
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref())
.map(ToOwned::to_owned);
let upstream_is_stream = row
.try_get::<Option<bool>, _>("upstream_is_stream")
.map_sql_err()?;

6
crates/aether-data/src/repository/usage/sqlite.rs
View File

@@ -6,8 +6,8 @@ use sqlx::{sqlite::SqliteRow, Row};
use super::{
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
InMemoryUsageReadRepository, PendingUsageCleanupSummary, StoredRequestUsageAudit,
UpsertUsageRecord, UsageWriteRepository,
usage_request_metadata_client_family, InMemoryUsageReadRepository, PendingUsageCleanupSummary,
StoredRequestUsageAudit, UpsertUsageRecord, UsageWriteRepository,
};
use crate::driver::sqlite::{sqlite_optional_real, sqlite_real, SqlitePool};
use crate::error::SqlResultExt;
@@ -847,6 +847,8 @@ fn map_usage_row(row: &SqliteRow) -> Result<StoredRequestUsageAudit, DataLayerEr
.map(|raw| serde_json::from_str(&raw))
.transpose()
.map_err(|err| DataLayerError::UnexpectedValue(err.to_string()))?;
audit.client_family = usage_request_metadata_client_family(audit.request_metadata.as_ref())
.map(ToOwned::to_owned);
let upstream_is_stream = row
.try_get::<Option<i64>, _>("upstream_is_stream")
.map_sql_err()?