mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix: 号池列表统计与最后使用时间显示
- 为 usage 仓储新增按 provider_api_key_id 汇总请求数、Token、费用和最后使用时间的能力
- 在 /api/admin/pool/{provider_id}/keys 中优先使用 usage 汇总结果覆盖 request_count、total_tokens、total_cost_usd、last_used_at
- 补充数据层与网关侧回归测试,避免号池管理页统计全为 0 且最后使用为空
This commit is contained in:
@@ -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(0)
|
||||
.max(item.created_at_unix_secs),
|
||||
);
|
||||
}
|
||||
Ok(summaries)
|
||||
}
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -4,8 +4,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};
|
||||
@@ -154,6 +154,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,
|
||||
@@ -645,6 +666,76 @@ 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()?
|
||||
.try_into()
|
||||
.map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"usage.request_count aggregate is negative".to_string(),
|
||||
)
|
||||
})?;
|
||||
let total_tokens = row
|
||||
.try_get::<i64, _>("total_tokens")
|
||||
.map_postgres_err()?
|
||||
.try_into()
|
||||
.map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"usage.total_tokens aggregate is negative".to_string(),
|
||||
)
|
||||
})?;
|
||||
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.try_into().map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(
|
||||
"usage.last_used_at_unix_secs aggregate is negative".to_string(),
|
||||
)
|
||||
})
|
||||
})
|
||||
.transpose()?;
|
||||
|
||||
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,
|
||||
@@ -762,6 +853,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,
|
||||
@@ -975,6 +1074,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\""));
|
||||
|
||||
Reference in New Issue
Block a user