mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-03 01:40:21 +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:
@@ -674,6 +674,26 @@ impl GatewayDataState {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn summarize_usage_by_provider_api_key_ids(
|
||||||
|
&self,
|
||||||
|
provider_api_key_ids: &[String],
|
||||||
|
) -> Result<
|
||||||
|
std::collections::BTreeMap<
|
||||||
|
String,
|
||||||
|
aether_data_contracts::repository::usage::StoredProviderApiKeyUsageSummary,
|
||||||
|
>,
|
||||||
|
DataLayerError,
|
||||||
|
> {
|
||||||
|
match &self.usage_reader {
|
||||||
|
Some(repository) => {
|
||||||
|
repository
|
||||||
|
.summarize_usage_by_provider_api_key_ids(provider_api_key_ids)
|
||||||
|
.await
|
||||||
|
}
|
||||||
|
None => Ok(std::collections::BTreeMap::new()),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn list_users_by_ids(
|
pub(crate) async fn list_users_by_ids(
|
||||||
&self,
|
&self,
|
||||||
user_ids: &[String],
|
user_ids: &[String],
|
||||||
|
|||||||
@@ -7,6 +7,7 @@ use super::{
|
|||||||
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
ADMIN_POOL_PROVIDER_CATALOG_READER_UNAVAILABLE_DETAIL,
|
||||||
};
|
};
|
||||||
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
use crate::handlers::admin::request::{AdminAppState, AdminRequestContext};
|
||||||
|
use crate::handlers::shared::unix_secs_to_rfc3339;
|
||||||
use crate::GatewayError;
|
use crate::GatewayError;
|
||||||
use axum::{
|
use axum::{
|
||||||
body::Body,
|
body::Body,
|
||||||
@@ -130,17 +131,38 @@ pub(super) async fn build_admin_pool_list_keys_response(
|
|||||||
}
|
}
|
||||||
_ => AdminProviderPoolRuntimeState::default(),
|
_ => AdminProviderPoolRuntimeState::default(),
|
||||||
};
|
};
|
||||||
|
let usage_summary_by_key_id = state
|
||||||
|
.app()
|
||||||
|
.summarize_usage_by_provider_api_key_ids(&key_ids)
|
||||||
|
.await?;
|
||||||
|
|
||||||
let items = keys
|
let items = keys
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|key| {
|
.map(|key| {
|
||||||
pool_payloads::build_admin_pool_key_payload(
|
let mut payload = pool_payloads::build_admin_pool_key_payload(
|
||||||
state,
|
state,
|
||||||
&provider.provider_type,
|
&provider.provider_type,
|
||||||
&key,
|
&key,
|
||||||
&runtime,
|
&runtime,
|
||||||
pool_config,
|
pool_config,
|
||||||
)
|
);
|
||||||
|
if let Some(summary) = usage_summary_by_key_id.get(&key.id) {
|
||||||
|
if let Some(object) = payload.as_object_mut() {
|
||||||
|
object.insert("request_count".to_string(), json!(summary.request_count));
|
||||||
|
object.insert("total_tokens".to_string(), json!(summary.total_tokens));
|
||||||
|
object.insert(
|
||||||
|
"total_cost_usd".to_string(),
|
||||||
|
json!(format!("{:.8}", summary.total_cost_usd)),
|
||||||
|
);
|
||||||
|
object.insert(
|
||||||
|
"last_used_at".to_string(),
|
||||||
|
json!(summary
|
||||||
|
.last_used_at_unix_secs
|
||||||
|
.and_then(unix_secs_to_rfc3339)),
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
payload
|
||||||
})
|
})
|
||||||
.collect::<Vec<_>>();
|
.collect::<Vec<_>>();
|
||||||
|
|
||||||
|
|||||||
@@ -65,6 +65,19 @@ impl AppState {
|
|||||||
.map_err(|err| GatewayError::Internal(err.to_string()))
|
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) async fn summarize_usage_by_provider_api_key_ids(
|
||||||
|
&self,
|
||||||
|
provider_api_key_ids: &[String],
|
||||||
|
) -> Result<
|
||||||
|
std::collections::BTreeMap<String, usage::StoredProviderApiKeyUsageSummary>,
|
||||||
|
GatewayError,
|
||||||
|
> {
|
||||||
|
self.data
|
||||||
|
.summarize_usage_by_provider_api_key_ids(provider_api_key_ids)
|
||||||
|
.await
|
||||||
|
.map_err(|err| GatewayError::Internal(err.to_string()))
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn list_users_by_ids(
|
pub(crate) async fn list_users_by_ids(
|
||||||
&self,
|
&self,
|
||||||
user_ids: &[String],
|
user_ids: &[String],
|
||||||
|
|||||||
@@ -1,8 +1,11 @@
|
|||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
use aether_crypto::{encrypt_python_fernet_plaintext, DEVELOPMENT_ENCRYPTION_KEY};
|
||||||
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
|
use aether_data::repository::{
|
||||||
|
provider_catalog::InMemoryProviderCatalogReadRepository, usage::InMemoryUsageReadRepository,
|
||||||
|
};
|
||||||
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository;
|
||||||
|
use aether_data_contracts::repository::usage::StoredRequestUsageAudit;
|
||||||
use axum::body::{to_bytes, Body, Bytes};
|
use axum::body::{to_bytes, Body, Bytes};
|
||||||
use axum::routing::{any, get, post};
|
use axum::routing::{any, get, post};
|
||||||
use axum::{extract::Request, Router};
|
use axum::{extract::Request, Router};
|
||||||
@@ -467,6 +470,152 @@ async fn gateway_pool_list_includes_usage_totals_and_nullable_lru_score() {
|
|||||||
assert!(keys[0]["lru_score"].is_null());
|
assert!(keys[0]["lru_score"].is_null());
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn gateway_pool_list_uses_usage_aggregates_for_stats_and_last_used() {
|
||||||
|
let provider = sample_provider("provider-openai", "openai", 10).with_transport_fields(
|
||||||
|
true,
|
||||||
|
false,
|
||||||
|
true,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(json!({
|
||||||
|
"pool_advanced": {
|
||||||
|
"enabled": true
|
||||||
|
}
|
||||||
|
})),
|
||||||
|
);
|
||||||
|
let mut key = sample_key(
|
||||||
|
"key-openai-usage-live",
|
||||||
|
"provider-openai",
|
||||||
|
"openai:chat",
|
||||||
|
"sk-usage-live",
|
||||||
|
);
|
||||||
|
key.name = "live usage key".to_string();
|
||||||
|
key.request_count = Some(0);
|
||||||
|
key.total_tokens = 0;
|
||||||
|
key.total_cost_usd = 0.0;
|
||||||
|
key.last_used_at_unix_secs = None;
|
||||||
|
|
||||||
|
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||||
|
vec![provider],
|
||||||
|
Vec::new(),
|
||||||
|
vec![key],
|
||||||
|
));
|
||||||
|
let usage_repository = Arc::new(InMemoryUsageReadRepository::seed(vec![
|
||||||
|
StoredRequestUsageAudit::new(
|
||||||
|
"usage-1".to_string(),
|
||||||
|
"request-1".to_string(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
"openai".to_string(),
|
||||||
|
"gpt-5.4".to_string(),
|
||||||
|
None,
|
||||||
|
Some("provider-openai".to_string()),
|
||||||
|
None,
|
||||||
|
Some("key-openai-usage-live".to_string()),
|
||||||
|
None,
|
||||||
|
Some("openai".to_string()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
false,
|
||||||
|
false,
|
||||||
|
20,
|
||||||
|
30,
|
||||||
|
50,
|
||||||
|
1.25,
|
||||||
|
1.25,
|
||||||
|
Some(200),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(900),
|
||||||
|
None,
|
||||||
|
"completed".to_string(),
|
||||||
|
"settled".to_string(),
|
||||||
|
1_711_000_000,
|
||||||
|
1_711_000_000,
|
||||||
|
Some(1_711_000_010),
|
||||||
|
)
|
||||||
|
.expect("usage should build"),
|
||||||
|
StoredRequestUsageAudit::new(
|
||||||
|
"usage-2".to_string(),
|
||||||
|
"request-2".to_string(),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
"openai".to_string(),
|
||||||
|
"gpt-5.4".to_string(),
|
||||||
|
None,
|
||||||
|
Some("provider-openai".to_string()),
|
||||||
|
None,
|
||||||
|
Some("key-openai-usage-live".to_string()),
|
||||||
|
None,
|
||||||
|
Some("openai".to_string()),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
false,
|
||||||
|
true,
|
||||||
|
15,
|
||||||
|
35,
|
||||||
|
50,
|
||||||
|
2.5,
|
||||||
|
2.5,
|
||||||
|
Some(200),
|
||||||
|
None,
|
||||||
|
None,
|
||||||
|
Some(1200),
|
||||||
|
None,
|
||||||
|
"completed".to_string(),
|
||||||
|
"settled".to_string(),
|
||||||
|
1_711_000_120,
|
||||||
|
1_711_000_120,
|
||||||
|
Some(1_711_000_140),
|
||||||
|
)
|
||||||
|
.expect("usage should build"),
|
||||||
|
]));
|
||||||
|
let state = AppState::new()
|
||||||
|
.expect("gateway should build")
|
||||||
|
.with_data_state_for_tests(
|
||||||
|
GatewayDataState::with_provider_catalog_and_usage_reader_for_tests(
|
||||||
|
provider_catalog_repository,
|
||||||
|
usage_repository,
|
||||||
|
),
|
||||||
|
);
|
||||||
|
|
||||||
|
let response = local_admin_pool_response(
|
||||||
|
&state,
|
||||||
|
http::Method::GET,
|
||||||
|
"/api/admin/pool/provider-openai/keys?page=1&page_size=50&status=all",
|
||||||
|
None,
|
||||||
|
)
|
||||||
|
.await;
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::OK);
|
||||||
|
let payload: serde_json::Value = serde_json::from_slice(
|
||||||
|
&to_bytes(response.into_body(), usize::MAX)
|
||||||
|
.await
|
||||||
|
.expect("body should read"),
|
||||||
|
)
|
||||||
|
.expect("json body should parse");
|
||||||
|
let keys = payload["keys"].as_array().expect("keys should be array");
|
||||||
|
assert_eq!(keys.len(), 1);
|
||||||
|
assert_eq!(keys[0]["request_count"], json!(2));
|
||||||
|
assert_eq!(keys[0]["total_tokens"], json!(100u64));
|
||||||
|
assert_eq!(keys[0]["total_cost_usd"], json!("3.75000000"));
|
||||||
|
assert_eq!(keys[0]["last_used_at"], json!("2024-03-21T05:48:40Z"));
|
||||||
|
}
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_principal() {
|
async fn gateway_handles_admin_pool_list_keys_locally_with_trusted_admin_principal() {
|
||||||
let upstream_hits = Arc::new(Mutex::new(0usize));
|
let upstream_hits = Arc::new(Mutex::new(0usize));
|
||||||
|
|||||||
@@ -1,7 +1,7 @@
|
|||||||
mod types;
|
mod types;
|
||||||
|
|
||||||
pub use types::{
|
pub use types::{
|
||||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
|
||||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
|
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
|
||||||
UsageWriteRepository,
|
UsageRepository, UsageWriteRepository,
|
||||||
};
|
};
|
||||||
|
|||||||
@@ -279,6 +279,15 @@ pub struct StoredProviderUsageSummary {
|
|||||||
pub total_cost_usd: f64,
|
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)]
|
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||||
pub struct UsageAuditListQuery {
|
pub struct UsageAuditListQuery {
|
||||||
pub created_from_unix_secs: Option<u64>,
|
pub created_from_unix_secs: Option<u64>,
|
||||||
@@ -316,6 +325,14 @@ pub trait UsageReadRepository: Send + Sync {
|
|||||||
api_key_ids: &[String],
|
api_key_ids: &[String],
|
||||||
) -> Result<std::collections::BTreeMap<String, u64>, crate::DataLayerError>;
|
) -> 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(
|
async fn summarize_provider_usage_since(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
|
|||||||
@@ -4,8 +4,9 @@ use std::sync::RwLock;
|
|||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
|
||||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
|
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
|
||||||
|
UsageWriteRepository,
|
||||||
};
|
};
|
||||||
use crate::DataLayerError;
|
use crate::DataLayerError;
|
||||||
|
|
||||||
@@ -165,6 +166,46 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
|||||||
Ok(totals)
|
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(
|
async fn summarize_provider_usage_since(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
|
|||||||
@@ -3,9 +3,9 @@ mod sql;
|
|||||||
|
|
||||||
#[allow(unused_imports)]
|
#[allow(unused_imports)]
|
||||||
pub(crate) use aether_data_contracts::repository::usage::{
|
pub(crate) use aether_data_contracts::repository::usage::{
|
||||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredProviderUsageWindow,
|
||||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
|
StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository,
|
||||||
UsageWriteRepository,
|
UsageRepository, UsageWriteRepository,
|
||||||
};
|
};
|
||||||
pub use memory::InMemoryUsageReadRepository;
|
pub use memory::InMemoryUsageReadRepository;
|
||||||
pub use sql::SqlxUsageReadRepository;
|
pub use sql::SqlxUsageReadRepository;
|
||||||
|
|||||||
@@ -4,8 +4,8 @@ use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
|||||||
use uuid::Uuid;
|
use uuid::Uuid;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery,
|
StoredProviderApiKeyUsageSummary, StoredProviderUsageSummary, StoredRequestUsageAudit,
|
||||||
UsageReadRepository, UsageWriteRepository,
|
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
|
||||||
};
|
};
|
||||||
use crate::postgres::PostgresTransactionRunner;
|
use crate::postgres::PostgresTransactionRunner;
|
||||||
use crate::{error::SqlxResultExt, DataLayerError};
|
use crate::{error::SqlxResultExt, DataLayerError};
|
||||||
@@ -154,6 +154,27 @@ GROUP BY api_key_id
|
|||||||
ORDER BY api_key_id ASC
|
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#"
|
const LIST_USAGE_AUDITS_PREFIX: &str = r#"
|
||||||
SELECT
|
SELECT
|
||||||
id,
|
id,
|
||||||
@@ -645,6 +666,76 @@ impl SqlxUsageReadRepository {
|
|||||||
Ok(totals)
|
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(
|
pub async fn upsert(
|
||||||
&self,
|
&self,
|
||||||
usage: UpsertUsageRecord,
|
usage: UpsertUsageRecord,
|
||||||
@@ -762,6 +853,14 @@ impl UsageReadRepository for SqlxUsageReadRepository {
|
|||||||
Self::summarize_total_tokens_by_api_key_ids(self, api_key_ids).await
|
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(
|
async fn summarize_provider_usage_since(
|
||||||
&self,
|
&self,
|
||||||
provider_id: &str,
|
provider_id: &str,
|
||||||
@@ -975,6 +1074,14 @@ mod tests {
|
|||||||
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("ANY($1::TEXT[])"));
|
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]
|
#[test]
|
||||||
fn usage_sql_supports_recent_usage_audits_query() {
|
fn usage_sql_supports_recent_usage_audits_query() {
|
||||||
assert!(super::LIST_RECENT_USAGE_AUDITS_PREFIX.contains("FROM \"usage\""));
|
assert!(super::LIST_RECENT_USAGE_AUDITS_PREFIX.contains("FROM \"usage\""));
|
||||||
|
|||||||
Reference in New Issue
Block a user