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

fix(admin): restore pool key usage aggregates

Browse Source
This commit is contained in:
fawney19
2026-04-11 02:54:54 +08:00
parent cf02a10050
commit a54ac76688
9 changed files with 389 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,7 +1,7 @@
mod types;
pub use types::{
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
UsageWriteRepository,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
UsageRepository, UsageWriteRepository,
};

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

@@ -280,6 +280,15 @@ pub struct StoredProviderUsageSummary {
pub total_cost_usd: f64,
}
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
pub struct StoredProviderApiKeyUsageSummary {
pub provider_api_key_id: String,
pub request_count: u64,
pub total_tokens: u64,
pub total_cost_usd: f64,
pub last_used_at_unix_secs: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
pub struct UsageAuditListQuery {
pub created_from_unix_secs: Option<u64>,
@@ -317,6 +326,14 @@ pub trait UsageReadRepository: Send + Sync {
api_key_ids: &[String],
) -> Result<std::collections::BTreeMap<String, u64>, crate::DataLayerError>;
async fn summarize_usage_by_provider_api_key_ids(
&self,
provider_api_key_ids: &[String],
) -> Result<
std::collections::BTreeMap<String, StoredProviderApiKeyUsageSummary>,
crate::DataLayerError,
>;
async fn summarize_provider_usage_since(
&self,
provider_id: &str,

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

@@ -4,8 +4,9 @@ use std::sync::RwLock;
use async_trait::async_trait;
use super::{
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
UsageWriteRepository,
};
use crate::DataLayerError;
@@ -165,6 +166,46 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
Ok(totals)
}
async fn summarize_usage_by_provider_api_key_ids(
&self,
provider_api_key_ids: &[String],
) -> Result<BTreeMap<String, StoredProviderApiKeyUsageSummary>, DataLayerError> {
let provider_api_key_id_set = provider_api_key_ids
.iter()
.map(String::as_str)
.collect::<Vec<_>>();
let mut summaries = BTreeMap::<String, StoredProviderApiKeyUsageSummary>::new();
for item in self
.by_request_id
.read()
.expect("usage repository lock")
.values()
{
let Some(provider_api_key_id) = item.provider_api_key_id.as_deref() else {
continue;
};
if !provider_api_key_id_set.contains(&provider_api_key_id) {
continue;
}
let entry = summaries
.entry(provider_api_key_id.to_string())
.or_insert_with(|| StoredProviderApiKeyUsageSummary {
provider_api_key_id: provider_api_key_id.to_string(),
..StoredProviderApiKeyUsageSummary::default()
});
entry.request_count = entry.request_count.saturating_add(1);
entry.total_tokens = entry.total_tokens.saturating_add(item.total_tokens);
entry.total_cost_usd += item.total_cost_usd;
entry.last_used_at_unix_secs = Some(
entry
.last_used_at_unix_secs
.unwrap_or_default()
.max(item.created_at_unix_ms),
);
}
Ok(summaries)
}
async fn summarize_provider_usage_since(
&self,
provider_id: &str,
@@ -533,4 +574,25 @@ mod tests {
assert_eq!(summary.avg_response_time_ms, 180.0);
assert_eq!(summary.total_cost_usd, 0.75);
}
#[tokio::test]
async fn summarizes_usage_by_provider_api_key_ids() {
let repository = InMemoryUsageReadRepository::seed(vec![
sample_usage("req-1", 1_711_000_000),
sample_usage("req-2", 1_711_000_250),
]);
let usage = repository
.summarize_usage_by_provider_api_key_ids(&["provider-key-1".to_string()])
.await
.expect("summary should succeed");
let item = usage
.get("provider-key-1")
.expect("provider key summary should exist");
assert_eq!(item.request_count, 2);
assert_eq!(item.total_tokens, 300);
assert_eq!(item.total_cost_usd, 0.24);
assert_eq!(item.last_used_at_unix_secs, Some(1_711_000_250));
}
}

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

@@ -3,9 +3,9 @@ mod sql;
#[allow(unused_imports)]
pub(crate) use aether_data_contracts::repository::usage::{
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
UsageWriteRepository,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
UsageRepository, UsageWriteRepository,
};
pub use memory::InMemoryUsageReadRepository;
pub use sql::SqlxUsageReadRepository;

94
crates/aether-data/src/repository/usage/sql.rs
View File

@@ -5,8 +5,8 @@ use sqlx::{PgPool, Postgres, QueryBuilder, Row};
use uuid::Uuid;
use super::{
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery,
UsageReadRepository, UsageWriteRepository,
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
};
use crate::postgres::PostgresTransactionRunner;
use crate::{error::SqlxResultExt, DataLayerError};
@@ -159,6 +159,27 @@ GROUP BY api_key_id
ORDER BY api_key_id ASC
"#;
const SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL: &str = r#"
SELECT
provider_api_key_id,
COUNT(*)::BIGINT AS request_count,
COALESCE(
SUM(
COALESCE(
total_tokens,
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
)
),
0
) AS total_tokens,
COALESCE(CAST(SUM(total_cost_usd) AS DOUBLE PRECISION), 0) AS total_cost_usd,
CAST(EXTRACT(EPOCH FROM MAX(created_at)) AS BIGINT) AS last_used_at_unix_secs
FROM "usage"
WHERE provider_api_key_id = ANY($1::TEXT[])
GROUP BY provider_api_key_id
ORDER BY provider_api_key_id ASC
"#;
const LIST_USAGE_AUDITS_PREFIX: &str = r#"
SELECT
id,
@@ -671,6 +692,59 @@ impl SqlxUsageReadRepository {
Ok(totals)
}
pub async fn summarize_usage_by_provider_api_key_ids(
&self,
provider_api_key_ids: &[String],
) -> Result<std::collections::BTreeMap<String, StoredProviderApiKeyUsageSummary>, DataLayerError>
{
if provider_api_key_ids.is_empty() {
return Ok(std::collections::BTreeMap::new());
}
let rows = sqlx::query(SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL)
.bind(provider_api_key_ids)
.fetch_all(&self.pool)
.await
.map_postgres_err()?;
let mut summaries = std::collections::BTreeMap::new();
for row in rows {
let provider_api_key_id: String =
row.try_get("provider_api_key_id").map_postgres_err()?;
let request_count = row
.try_get::<i64, _>("request_count")
.map_postgres_err()?
.max(0) as u64;
let total_tokens = row
.try_get::<i64, _>("total_tokens")
.map_postgres_err()?
.max(0) as u64;
let total_cost_usd: f64 = row.try_get("total_cost_usd").map_postgres_err()?;
if !total_cost_usd.is_finite() {
return Err(DataLayerError::UnexpectedValue(
"usage.total_cost_usd aggregate is not finite".to_string(),
));
}
let last_used_at_unix_secs = row
.try_get::<Option<i64>, _>("last_used_at_unix_secs")
.map_postgres_err()?
.map(|value| value.max(0) as u64);
summaries.insert(
provider_api_key_id.clone(),
StoredProviderApiKeyUsageSummary {
provider_api_key_id,
request_count,
total_tokens,
total_cost_usd,
last_used_at_unix_secs,
},
);
}
Ok(summaries)
}
pub async fn upsert(
&self,
usage: UpsertUsageRecord,
@@ -813,6 +887,14 @@ impl UsageReadRepository for SqlxUsageReadRepository {
Self::summarize_total_tokens_by_api_key_ids(self, api_key_ids).await
}
async fn summarize_usage_by_provider_api_key_ids(
&self,
provider_api_key_ids: &[String],
) -> Result<std::collections::BTreeMap<String, StoredProviderApiKeyUsageSummary>, DataLayerError>
{
Self::summarize_usage_by_provider_api_key_ids(self, provider_api_key_ids).await
}
async fn summarize_provider_usage_since(
&self,
provider_id: &str,
@@ -1050,6 +1132,14 @@ mod tests {
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
}
#[test]
fn usage_sql_summarizes_usage_by_provider_api_key_ids_in_database() {
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL
.contains("GROUP BY provider_api_key_id"));
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL.contains("MAX(created_at)"));
assert!(super::SUMMARIZE_USAGE_BY_PROVIDER_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
}
#[test]
fn usage_sql_supports_recent_usage_audits_query() {
assert!(super::LIST_RECENT_USAGE_AUDITS_PREFIX.contains("FROM \"usage\""));