mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 17:30:23 +08:00
feat: 扩展 Rust gateway 全功能模块，新增 billing/crypto/wallet crate 及完整数据层
- 新增 aether-billing、aether-crypto、aether-wallet 独立 crate - aether-data 扩展 repository 层：announcements、auth_modules、billing、 candidate_selection、gemini_file_mappings、global_models、management_tokens、 oauth_providers、proxy_nodes、quota、users、wallet 等模块 - aether-gateway 新增 api/auth/billing/control/middleware/scheduler/usage/ video_tasks/hooks/maintenance/model_fetch/provider_transport 等功能模块 - 重构 executor decision 和 gateway state 为模块目录结构 - 新增 gateway router、frontdoor 路由层及对应测试 - Python 侧 API 路由重构，新增 compat/support 模块 - 前端 Logo 组件更新及 Provider 管理页面调整
This commit is contained in:
fawney19
@@ -3,12 +3,16 @@ use std::sync::RwLock;
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use super::types::{
|
||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageWriteRepository,
|
||||
};
|
||||
use crate::DataLayerError;
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryUsageReadRepository {
|
||||
by_request_id: RwLock<BTreeMap<String, StoredRequestUsageAudit>>,
|
||||
provider_usage_windows: RwLock<Vec<StoredProviderUsageWindow>>,
|
||||
}
|
||||
|
||||
impl InMemoryUsageReadRepository {
|
||||
@@ -22,12 +26,36 @@ impl InMemoryUsageReadRepository {
|
||||
}
|
||||
Self {
|
||||
by_request_id: RwLock::new(by_request_id),
|
||||
provider_usage_windows: RwLock::new(Vec::new()),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_provider_usage_windows<I>(self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredProviderUsageWindow>,
|
||||
{
|
||||
Self {
|
||||
by_request_id: self.by_request_id,
|
||||
provider_usage_windows: RwLock::new(items.into_iter().collect()),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Ok(self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.find(|item| item.id == id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
@@ -39,12 +67,250 @@ impl UsageReadRepository for InMemoryUsageReadRepository {
|
||||
.get(request_id)
|
||||
.cloned())
|
||||
}
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut items: Vec<_> = self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| {
|
||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||
if item.created_at_unix_secs < created_from_unix_secs {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||||
if item.created_at_unix_secs >= created_until_unix_secs {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(user_id) = query.user_id.as_deref() {
|
||||
if item.user_id.as_deref() != Some(user_id) {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(provider_name) = query.provider_name.as_deref() {
|
||||
if item.provider_name != provider_name {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
if let Some(model) = query.model.as_deref() {
|
||||
if item.model != model {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
true
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
items.sort_by(|left, right| {
|
||||
left.created_at_unix_secs
|
||||
.cmp(&right.created_at_unix_secs)
|
||||
.then_with(|| left.request_id.cmp(&right.request_id))
|
||||
});
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut items: Vec<_> = self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
.filter(|item| match user_id {
|
||||
Some(user_id) => item.user_id.as_deref() == Some(user_id),
|
||||
None => true,
|
||||
})
|
||||
.cloned()
|
||||
.collect();
|
||||
items.sort_by(|left, right| {
|
||||
right
|
||||
.created_at_unix_secs
|
||||
.cmp(&left.created_at_unix_secs)
|
||||
.then_with(|| left.id.cmp(&right.id))
|
||||
});
|
||||
items.truncate(limit);
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<BTreeMap<String, u64>, DataLayerError> {
|
||||
let api_key_id_set = api_key_ids.iter().map(String::as_str).collect::<Vec<_>>();
|
||||
let mut totals = BTreeMap::<String, u64>::new();
|
||||
for item in self
|
||||
.by_request_id
|
||||
.read()
|
||||
.expect("usage repository lock")
|
||||
.values()
|
||||
{
|
||||
let Some(api_key_id) = item.api_key_id.as_deref() else {
|
||||
continue;
|
||||
};
|
||||
if !api_key_id_set.contains(&api_key_id) {
|
||||
continue;
|
||||
}
|
||||
let entry = totals.entry(api_key_id.to_string()).or_insert(0);
|
||||
*entry = (*entry).saturating_add(item.total_tokens);
|
||||
}
|
||||
Ok(totals)
|
||||
}
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
let windows = self
|
||||
.provider_usage_windows
|
||||
.read()
|
||||
.expect("provider usage repository lock");
|
||||
|
||||
let mut summary = StoredProviderUsageSummary::default();
|
||||
let mut response_time_samples = 0u64;
|
||||
for window in windows.iter().filter(|window| {
|
||||
window.provider_id == provider_id && window.window_start_unix_secs >= since_unix_secs
|
||||
}) {
|
||||
summary.total_requests = summary.total_requests.saturating_add(window.total_requests);
|
||||
summary.successful_requests = summary
|
||||
.successful_requests
|
||||
.saturating_add(window.successful_requests);
|
||||
summary.failed_requests = summary
|
||||
.failed_requests
|
||||
.saturating_add(window.failed_requests);
|
||||
summary.total_cost_usd += window.total_cost_usd;
|
||||
summary.avg_response_time_ms += window.avg_response_time_ms;
|
||||
response_time_samples = response_time_samples.saturating_add(1);
|
||||
}
|
||||
|
||||
if response_time_samples > 0 {
|
||||
summary.avg_response_time_ms /= response_time_samples as f64;
|
||||
}
|
||||
|
||||
Ok(summary)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageWriteRepository for InMemoryUsageReadRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
usage.validate()?;
|
||||
let mut by_request_id = self.by_request_id.write().expect("usage repository lock");
|
||||
|
||||
let created_at_unix_secs = by_request_id
|
||||
.get(&usage.request_id)
|
||||
.map(|existing| existing.created_at_unix_secs)
|
||||
.or(usage.created_at_unix_secs)
|
||||
.unwrap_or(usage.updated_at_unix_secs);
|
||||
|
||||
let total_tokens = usage
|
||||
.total_tokens
|
||||
.or_else(|| {
|
||||
Some(
|
||||
usage.input_tokens.unwrap_or_default()
|
||||
+ usage.output_tokens.unwrap_or_default(),
|
||||
)
|
||||
})
|
||||
.unwrap_or_default();
|
||||
let existing = by_request_id.get(&usage.request_id);
|
||||
|
||||
let stored = StoredRequestUsageAudit {
|
||||
id: existing
|
||||
.map(|existing| existing.id.clone())
|
||||
.unwrap_or_else(|| format!("usage-{}", usage.request_id)),
|
||||
request_id: usage.request_id.clone(),
|
||||
user_id: usage.user_id,
|
||||
api_key_id: usage.api_key_id,
|
||||
username: usage.username,
|
||||
api_key_name: usage.api_key_name,
|
||||
provider_name: usage.provider_name,
|
||||
model: usage.model,
|
||||
target_model: usage.target_model,
|
||||
provider_id: usage.provider_id,
|
||||
provider_endpoint_id: usage.provider_endpoint_id,
|
||||
provider_api_key_id: usage.provider_api_key_id,
|
||||
request_type: usage.request_type,
|
||||
api_format: usage.api_format,
|
||||
api_family: usage.api_family,
|
||||
endpoint_kind: usage.endpoint_kind,
|
||||
endpoint_api_format: usage.endpoint_api_format,
|
||||
provider_api_family: usage.provider_api_family,
|
||||
provider_endpoint_kind: usage.provider_endpoint_kind,
|
||||
has_format_conversion: usage.has_format_conversion.unwrap_or(false),
|
||||
is_stream: usage.is_stream.unwrap_or(false),
|
||||
input_tokens: usage.input_tokens.unwrap_or_default(),
|
||||
output_tokens: usage.output_tokens.unwrap_or_default(),
|
||||
total_tokens,
|
||||
cache_creation_input_tokens: usage.cache_creation_input_tokens.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_creation_input_tokens)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
cache_read_input_tokens: usage.cache_read_input_tokens.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_read_input_tokens)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
cache_creation_cost_usd: usage.cache_creation_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_creation_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
cache_read_cost_usd: usage.cache_read_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.cache_read_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
output_price_per_1m: usage
|
||||
.output_price_per_1m
|
||||
.or_else(|| existing.and_then(|existing| existing.output_price_per_1m)),
|
||||
total_cost_usd: usage.total_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.total_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
actual_total_cost_usd: usage.actual_total_cost_usd.unwrap_or_else(|| {
|
||||
existing
|
||||
.map(|existing| existing.actual_total_cost_usd)
|
||||
.unwrap_or_default()
|
||||
}),
|
||||
status_code: usage.status_code,
|
||||
error_message: usage.error_message,
|
||||
error_category: usage.error_category,
|
||||
response_time_ms: usage.response_time_ms,
|
||||
first_byte_time_ms: usage.first_byte_time_ms,
|
||||
status: usage.status,
|
||||
billing_status: usage.billing_status,
|
||||
created_at_unix_secs,
|
||||
updated_at_unix_secs: usage.updated_at_unix_secs,
|
||||
finalized_at_unix_secs: usage.finalized_at_unix_secs,
|
||||
};
|
||||
|
||||
by_request_id.insert(stored.request_id.clone(), stored.clone());
|
||||
Ok(stored)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::InMemoryUsageReadRepository;
|
||||
use crate::repository::usage::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use crate::repository::usage::{
|
||||
StoredProviderUsageWindow, StoredRequestUsageAudit, UpsertUsageRecord, UsageReadRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
use serde_json::json;
|
||||
|
||||
fn sample_usage(request_id: &str, created_at_unix_secs: i64) -> StoredRequestUsageAudit {
|
||||
StoredRequestUsageAudit::new(
|
||||
@@ -104,4 +370,124 @@ mod tests {
|
||||
assert_eq!(usage.request_id, "req-2");
|
||||
assert_eq!(usage.total_tokens, 150);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_writes_usage_record() {
|
||||
let repository = InMemoryUsageReadRepository::default();
|
||||
let stored = repository
|
||||
.upsert(UpsertUsageRecord {
|
||||
request_id: "req-upsert-1".to_string(),
|
||||
user_id: Some("user-1".to_string()),
|
||||
api_key_id: Some("key-1".to_string()),
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "OpenAI".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: Some("gpt-5-mini".to_string()),
|
||||
provider_id: Some("provider-1".to_string()),
|
||||
provider_endpoint_id: Some("endpoint-1".to_string()),
|
||||
provider_api_key_id: Some("provider-key-1".to_string()),
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(true),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: None,
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: Some(0.25),
|
||||
actual_total_cost_usd: Some(0.15),
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(300),
|
||||
first_byte_time_ms: Some(120),
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: Some(json!({"authorization": "Bearer test"})),
|
||||
request_body: Some(json!({"model": "gpt-5"})),
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
})
|
||||
.await
|
||||
.expect("upsert should succeed");
|
||||
|
||||
assert_eq!(stored.request_id, "req-upsert-1");
|
||||
assert_eq!(stored.total_tokens, 30);
|
||||
assert_eq!(stored.total_cost_usd, 0.25);
|
||||
assert_eq!(stored.actual_total_cost_usd, 0.15);
|
||||
assert_eq!(
|
||||
repository
|
||||
.find_by_request_id("req-upsert-1")
|
||||
.await
|
||||
.expect("find should succeed")
|
||||
.expect("usage should exist")
|
||||
.model,
|
||||
"gpt-5"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn summarizes_provider_usage_windows_since_timestamp() {
|
||||
let repository = InMemoryUsageReadRepository::default().with_provider_usage_windows(vec![
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-1".to_string(),
|
||||
1_700_000_000,
|
||||
10,
|
||||
9,
|
||||
1,
|
||||
120.0,
|
||||
1.25,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-1".to_string(),
|
||||
1_700_003_600,
|
||||
6,
|
||||
5,
|
||||
1,
|
||||
180.0,
|
||||
0.75,
|
||||
)
|
||||
.expect("window should build"),
|
||||
StoredProviderUsageWindow::new(
|
||||
"provider-2".to_string(),
|
||||
1_700_003_600,
|
||||
99,
|
||||
99,
|
||||
0,
|
||||
50.0,
|
||||
5.0,
|
||||
)
|
||||
.expect("window should build"),
|
||||
]);
|
||||
|
||||
let summary = repository
|
||||
.summarize_provider_usage_since("provider-1", 1_700_000_100)
|
||||
.await
|
||||
.expect("summary should succeed");
|
||||
|
||||
assert_eq!(summary.total_requests, 6);
|
||||
assert_eq!(summary.successful_requests, 5);
|
||||
assert_eq!(summary.failed_requests, 1);
|
||||
assert_eq!(summary.avg_response_time_ms, 180.0);
|
||||
assert_eq!(summary.total_cost_usd, 0.75);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -4,4 +4,8 @@ mod types;
|
||||
|
||||
pub use memory::InMemoryUsageReadRepository;
|
||||
pub use sql::SqlxUsageReadRepository;
|
||||
pub use types::{StoredRequestUsageAudit, UsageReadRepository, UsageRepository};
|
||||
pub use types::{
|
||||
StoredProviderUsageSummary, StoredProviderUsageWindow, StoredRequestUsageAudit,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageReadRepository, UsageRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
|
||||
@@ -1,7 +1,13 @@
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{PgPool, Row};
|
||||
use futures_util::future::BoxFuture;
|
||||
use sqlx::{PgPool, Postgres, QueryBuilder, Row};
|
||||
use uuid::Uuid;
|
||||
|
||||
use super::types::{StoredRequestUsageAudit, UsageReadRepository};
|
||||
use super::types::{
|
||||
StoredProviderUsageSummary, StoredRequestUsageAudit, UpsertUsageRecord, UsageAuditListQuery,
|
||||
UsageReadRepository, UsageWriteRepository,
|
||||
};
|
||||
use crate::postgres::PostgresTransactionRunner;
|
||||
use crate::DataLayerError;
|
||||
|
||||
const FIND_BY_REQUEST_ID_SQL: &str = r#"
|
||||
@@ -30,6 +36,11 @@ SELECT
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
@@ -40,27 +51,399 @@ SELECT
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM updated_at) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE request_id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const FIND_BY_ID_SQL: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
WHERE id = $1
|
||||
LIMIT 1
|
||||
"#;
|
||||
|
||||
const SUMMARIZE_PROVIDER_USAGE_SINCE_SQL: &str = r#"
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0) AS total_requests,
|
||||
COALESCE(SUM(successful_requests), 0) AS successful_requests,
|
||||
COALESCE(SUM(failed_requests), 0) AS failed_requests,
|
||||
COALESCE(AVG(avg_response_time_ms), 0) AS avg_response_time_ms,
|
||||
COALESCE(SUM(total_cost_usd), 0) AS total_cost_usd
|
||||
FROM provider_usage_tracking
|
||||
WHERE provider_id = $1
|
||||
AND window_start >= TO_TIMESTAMP($2::double precision)
|
||||
"#;
|
||||
|
||||
const SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL: &str = r#"
|
||||
SELECT
|
||||
api_key_id,
|
||||
COALESCE(
|
||||
SUM(
|
||||
COALESCE(
|
||||
total_tokens,
|
||||
COALESCE(input_tokens, 0) + COALESCE(output_tokens, 0)
|
||||
)
|
||||
),
|
||||
0
|
||||
) AS total_tokens
|
||||
FROM "usage"
|
||||
WHERE api_key_id = ANY($1::TEXT[])
|
||||
GROUP BY api_key_id
|
||||
ORDER BY api_key_id ASC
|
||||
"#;
|
||||
|
||||
const LIST_USAGE_AUDITS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
"#;
|
||||
|
||||
const LIST_RECENT_USAGE_AUDITS_PREFIX: &str = r#"
|
||||
SELECT
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
FROM "usage"
|
||||
"#;
|
||||
|
||||
const UPSERT_SQL: &str = r#"
|
||||
INSERT INTO "usage" (
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
has_format_conversion,
|
||||
is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
cache_creation_input_tokens,
|
||||
cache_read_input_tokens,
|
||||
cache_creation_cost_usd,
|
||||
cache_read_cost_usd,
|
||||
output_price_per_1m,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
request_headers,
|
||||
request_body,
|
||||
provider_request_headers,
|
||||
provider_request_body,
|
||||
response_headers,
|
||||
response_body,
|
||||
client_response_headers,
|
||||
client_response_body,
|
||||
request_metadata,
|
||||
finalized_at,
|
||||
created_at
|
||||
) VALUES (
|
||||
$1,
|
||||
$2,
|
||||
$3,
|
||||
$4,
|
||||
$5,
|
||||
$6,
|
||||
$7,
|
||||
$8,
|
||||
$9,
|
||||
$10,
|
||||
$11,
|
||||
$12,
|
||||
$13,
|
||||
$14,
|
||||
$15,
|
||||
$16,
|
||||
$17,
|
||||
$18,
|
||||
$19,
|
||||
COALESCE($20, FALSE),
|
||||
COALESCE($21, FALSE),
|
||||
COALESCE($22, 0),
|
||||
COALESCE($23, 0),
|
||||
COALESCE($24, COALESCE($22, 0) + COALESCE($23, 0)),
|
||||
COALESCE($25, 0),
|
||||
COALESCE($26, 0),
|
||||
COALESCE($27, 0),
|
||||
COALESCE($28, 0),
|
||||
$29,
|
||||
COALESCE($30, 0),
|
||||
COALESCE($31, 0),
|
||||
$32,
|
||||
$33,
|
||||
$34,
|
||||
$35,
|
||||
$36,
|
||||
$37,
|
||||
$38,
|
||||
$39,
|
||||
$40,
|
||||
$41,
|
||||
$42,
|
||||
$43,
|
||||
$44,
|
||||
$45,
|
||||
$46,
|
||||
CASE
|
||||
WHEN $47 IS NULL THEN NULL
|
||||
ELSE TO_TIMESTAMP($47::double precision)
|
||||
END,
|
||||
COALESCE(TO_TIMESTAMP($48::double precision), NOW())
|
||||
)
|
||||
ON CONFLICT (request_id)
|
||||
DO UPDATE SET
|
||||
user_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.user_id, "usage".user_id) ELSE "usage".user_id END,
|
||||
api_key_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_key_id, "usage".api_key_id) ELSE "usage".api_key_id END,
|
||||
username = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.username, "usage".username) ELSE "usage".username END,
|
||||
api_key_name = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_key_name, "usage".api_key_name) ELSE "usage".api_key_name END,
|
||||
provider_name = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_name, "usage".provider_name) ELSE "usage".provider_name END,
|
||||
model = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.model, "usage".model) ELSE "usage".model END,
|
||||
target_model = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.target_model, "usage".target_model) ELSE "usage".target_model END,
|
||||
provider_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_id, "usage".provider_id) ELSE "usage".provider_id END,
|
||||
provider_endpoint_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_endpoint_id, "usage".provider_endpoint_id) ELSE "usage".provider_endpoint_id END,
|
||||
provider_api_key_id = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_api_key_id, "usage".provider_api_key_id) ELSE "usage".provider_api_key_id END,
|
||||
request_type = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_type, "usage".request_type) ELSE "usage".request_type END,
|
||||
api_format = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_format, "usage".api_format) ELSE "usage".api_format END,
|
||||
api_family = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.api_family, "usage".api_family) ELSE "usage".api_family END,
|
||||
endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.endpoint_kind, "usage".endpoint_kind) ELSE "usage".endpoint_kind END,
|
||||
endpoint_api_format = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.endpoint_api_format, "usage".endpoint_api_format) ELSE "usage".endpoint_api_format END,
|
||||
provider_api_family = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_api_family, "usage".provider_api_family) ELSE "usage".provider_api_family END,
|
||||
provider_endpoint_kind = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_endpoint_kind, "usage".provider_endpoint_kind) ELSE "usage".provider_endpoint_kind END,
|
||||
has_format_conversion = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.has_format_conversion, "usage".has_format_conversion) ELSE "usage".has_format_conversion END,
|
||||
is_stream = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.is_stream, "usage".is_stream) ELSE "usage".is_stream END,
|
||||
input_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.input_tokens, "usage".input_tokens) ELSE "usage".input_tokens END,
|
||||
output_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.output_tokens, "usage".output_tokens) ELSE "usage".output_tokens END,
|
||||
total_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.total_tokens, "usage".total_tokens) ELSE "usage".total_tokens END,
|
||||
cache_creation_input_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_creation_input_tokens, "usage".cache_creation_input_tokens) ELSE "usage".cache_creation_input_tokens END,
|
||||
cache_read_input_tokens = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_read_input_tokens, "usage".cache_read_input_tokens) ELSE "usage".cache_read_input_tokens END,
|
||||
cache_creation_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_creation_cost_usd, "usage".cache_creation_cost_usd) ELSE "usage".cache_creation_cost_usd END,
|
||||
cache_read_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.cache_read_cost_usd, "usage".cache_read_cost_usd) ELSE "usage".cache_read_cost_usd END,
|
||||
output_price_per_1m = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.output_price_per_1m, "usage".output_price_per_1m) ELSE "usage".output_price_per_1m END,
|
||||
total_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.total_cost_usd, "usage".total_cost_usd) ELSE "usage".total_cost_usd END,
|
||||
actual_total_cost_usd = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.actual_total_cost_usd, "usage".actual_total_cost_usd) ELSE "usage".actual_total_cost_usd END,
|
||||
status_code = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.status_code, "usage".status_code) ELSE "usage".status_code END,
|
||||
error_message = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.error_message, "usage".error_message) ELSE "usage".error_message END,
|
||||
error_category = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.error_category, "usage".error_category) ELSE "usage".error_category END,
|
||||
response_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_time_ms, "usage".response_time_ms) ELSE "usage".response_time_ms END,
|
||||
first_byte_time_ms = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.first_byte_time_ms, "usage".first_byte_time_ms) ELSE "usage".first_byte_time_ms END,
|
||||
status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.status ELSE "usage".status END,
|
||||
billing_status = CASE WHEN "usage".billing_status = 'pending' THEN EXCLUDED.billing_status ELSE "usage".billing_status END,
|
||||
request_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_headers, "usage".request_headers) ELSE "usage".request_headers END,
|
||||
request_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_body, "usage".request_body) ELSE "usage".request_body END,
|
||||
provider_request_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_request_headers, "usage".provider_request_headers) ELSE "usage".provider_request_headers END,
|
||||
provider_request_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.provider_request_body, "usage".provider_request_body) ELSE "usage".provider_request_body END,
|
||||
response_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_headers, "usage".response_headers) ELSE "usage".response_headers END,
|
||||
response_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.response_body, "usage".response_body) ELSE "usage".response_body END,
|
||||
client_response_headers = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.client_response_headers, "usage".client_response_headers) ELSE "usage".client_response_headers END,
|
||||
client_response_body = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.client_response_body, "usage".client_response_body) ELSE "usage".client_response_body END,
|
||||
request_metadata = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.request_metadata, "usage".request_metadata) ELSE "usage".request_metadata END,
|
||||
finalized_at = CASE WHEN "usage".billing_status = 'pending' THEN COALESCE(EXCLUDED.finalized_at, "usage".finalized_at) ELSE "usage".finalized_at END
|
||||
RETURNING
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
provider_name,
|
||||
model,
|
||||
target_model,
|
||||
provider_id,
|
||||
provider_endpoint_id,
|
||||
provider_api_key_id,
|
||||
request_type,
|
||||
api_format,
|
||||
api_family,
|
||||
endpoint_kind,
|
||||
endpoint_api_format,
|
||||
provider_api_family,
|
||||
provider_endpoint_kind,
|
||||
COALESCE(has_format_conversion, FALSE) AS has_format_conversion,
|
||||
COALESCE(is_stream, FALSE) AS is_stream,
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
total_tokens,
|
||||
COALESCE(cache_creation_input_tokens, 0) AS cache_creation_input_tokens,
|
||||
COALESCE(cache_read_input_tokens, 0) AS cache_read_input_tokens,
|
||||
COALESCE(CAST(cache_creation_cost_usd AS DOUBLE PRECISION), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(CAST(cache_read_cost_usd AS DOUBLE PRECISION), 0) AS cache_read_cost_usd,
|
||||
CAST(output_price_per_1m AS DOUBLE PRECISION) AS output_price_per_1m,
|
||||
COALESCE(CAST(total_cost_usd AS DOUBLE PRECISION), 0) AS total_cost_usd,
|
||||
COALESCE(CAST(actual_total_cost_usd AS DOUBLE PRECISION), 0) AS actual_total_cost_usd,
|
||||
status_code,
|
||||
error_message,
|
||||
error_category,
|
||||
response_time_ms,
|
||||
first_byte_time_ms,
|
||||
status,
|
||||
billing_status,
|
||||
CAST(EXTRACT(EPOCH FROM created_at) AS BIGINT) AS created_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM COALESCE(finalized_at, created_at)) AS BIGINT) AS updated_at_unix_secs,
|
||||
CAST(EXTRACT(EPOCH FROM finalized_at) AS BIGINT) AS finalized_at_unix_secs
|
||||
"#;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct SqlxUsageReadRepository {
|
||||
pool: PgPool,
|
||||
tx_runner: PostgresTransactionRunner,
|
||||
}
|
||||
|
||||
impl SqlxUsageReadRepository {
|
||||
pub fn new(pool: PgPool) -> Self {
|
||||
Self { pool }
|
||||
let tx_runner = PostgresTransactionRunner::new(pool.clone());
|
||||
Self { pool, tx_runner }
|
||||
}
|
||||
|
||||
pub fn pool(&self) -> &PgPool {
|
||||
&self.pool
|
||||
}
|
||||
|
||||
pub fn transaction_runner(&self) -> &PostgresTransactionRunner {
|
||||
&self.tx_runner
|
||||
}
|
||||
|
||||
pub async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
@@ -71,20 +454,262 @@ impl SqlxUsageReadRepository {
|
||||
.await?;
|
||||
row.as_ref().map(map_usage_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let row = sqlx::query(FIND_BY_ID_SQL)
|
||||
.bind(id)
|
||||
.fetch_optional(&self.pool)
|
||||
.await?;
|
||||
row.as_ref().map(map_usage_row).transpose()
|
||||
}
|
||||
|
||||
pub async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
let row = sqlx::query(SUMMARIZE_PROVIDER_USAGE_SINCE_SQL)
|
||||
.bind(provider_id)
|
||||
.bind(since_unix_secs as f64)
|
||||
.fetch_one(&self.pool)
|
||||
.await?;
|
||||
|
||||
Ok(StoredProviderUsageSummary {
|
||||
total_requests: row.try_get::<i64, _>("total_requests")?.max(0) as u64,
|
||||
successful_requests: row.try_get::<i64, _>("successful_requests")?.max(0) as u64,
|
||||
failed_requests: row.try_get::<i64, _>("failed_requests")?.max(0) as u64,
|
||||
avg_response_time_ms: row.try_get::<f64, _>("avg_response_time_ms")?,
|
||||
total_cost_usd: row.try_get::<f64, _>("total_cost_usd")?,
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(LIST_USAGE_AUDITS_PREFIX);
|
||||
let mut has_where = false;
|
||||
|
||||
if let Some(created_from_unix_secs) = query.created_from_unix_secs {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder
|
||||
.push("created_at >= TO_TIMESTAMP(")
|
||||
.push_bind(created_from_unix_secs as f64)
|
||||
.push("::double precision)");
|
||||
}
|
||||
if let Some(created_until_unix_secs) = query.created_until_unix_secs {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder
|
||||
.push("created_at < TO_TIMESTAMP(")
|
||||
.push_bind(created_until_unix_secs as f64)
|
||||
.push("::double precision)");
|
||||
}
|
||||
if let Some(user_id) = query.user_id.as_deref() {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder.push("user_id = ").push_bind(user_id.to_string());
|
||||
}
|
||||
if let Some(provider_name) = query.provider_name.as_deref() {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
has_where = true;
|
||||
builder
|
||||
.push("provider_name = ")
|
||||
.push_bind(provider_name.to_string());
|
||||
}
|
||||
if let Some(model) = query.model.as_deref() {
|
||||
builder.push(if has_where { " AND " } else { " WHERE " });
|
||||
builder.push("model = ").push_bind(model.to_string());
|
||||
}
|
||||
|
||||
builder.push(" ORDER BY created_at ASC, request_id ASC");
|
||||
let rows = builder.build().fetch_all(&self.pool).await?;
|
||||
rows.iter().map(map_usage_row).collect()
|
||||
}
|
||||
|
||||
pub async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Postgres>::new(LIST_RECENT_USAGE_AUDITS_PREFIX);
|
||||
if let Some(user_id) = user_id {
|
||||
builder
|
||||
.push(" WHERE user_id = ")
|
||||
.push_bind(user_id.to_string());
|
||||
}
|
||||
builder
|
||||
.push(" ORDER BY created_at DESC, id ASC LIMIT ")
|
||||
.push_bind(i64::try_from(limit).map_err(|_| {
|
||||
DataLayerError::InvalidInput(format!("invalid recent usage limit: {limit}"))
|
||||
})?);
|
||||
let rows = builder.build().fetch_all(&self.pool).await?;
|
||||
rows.iter().map(map_usage_row).collect()
|
||||
}
|
||||
|
||||
pub async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
|
||||
if api_key_ids.is_empty() {
|
||||
return Ok(std::collections::BTreeMap::new());
|
||||
}
|
||||
|
||||
let rows = sqlx::query(SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL)
|
||||
.bind(api_key_ids)
|
||||
.fetch_all(&self.pool)
|
||||
.await?;
|
||||
|
||||
let mut totals = std::collections::BTreeMap::new();
|
||||
for row in rows {
|
||||
let api_key_id: String = row.try_get("api_key_id")?;
|
||||
let total_tokens = row.try_get::<i64, _>("total_tokens")?.max(0) as u64;
|
||||
totals.insert(api_key_id, total_tokens);
|
||||
}
|
||||
Ok(totals)
|
||||
}
|
||||
|
||||
pub async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
usage.validate()?;
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let row = sqlx::query(UPSERT_SQL)
|
||||
.bind(Uuid::new_v4().to_string())
|
||||
.bind(&usage.request_id)
|
||||
.bind(&usage.user_id)
|
||||
.bind(&usage.api_key_id)
|
||||
.bind(&usage.username)
|
||||
.bind(&usage.api_key_name)
|
||||
.bind(&usage.provider_name)
|
||||
.bind(&usage.model)
|
||||
.bind(&usage.target_model)
|
||||
.bind(&usage.provider_id)
|
||||
.bind(&usage.provider_endpoint_id)
|
||||
.bind(&usage.provider_api_key_id)
|
||||
.bind(&usage.request_type)
|
||||
.bind(&usage.api_format)
|
||||
.bind(&usage.api_family)
|
||||
.bind(&usage.endpoint_kind)
|
||||
.bind(&usage.endpoint_api_format)
|
||||
.bind(&usage.provider_api_family)
|
||||
.bind(&usage.provider_endpoint_kind)
|
||||
.bind(usage.has_format_conversion)
|
||||
.bind(usage.is_stream)
|
||||
.bind(usage.input_tokens.map(to_i32).transpose()?)
|
||||
.bind(usage.output_tokens.map(to_i32).transpose()?)
|
||||
.bind(
|
||||
usage
|
||||
.total_tokens
|
||||
.or_else(|| {
|
||||
Some(
|
||||
usage.input_tokens.unwrap_or_default()
|
||||
+ usage.output_tokens.unwrap_or_default(),
|
||||
)
|
||||
})
|
||||
.map(to_i32)
|
||||
.transpose()?,
|
||||
)
|
||||
.bind(usage.cache_creation_input_tokens.map(to_i32).transpose()?)
|
||||
.bind(usage.cache_read_input_tokens.map(to_i32).transpose()?)
|
||||
.bind(usage.cache_creation_cost_usd)
|
||||
.bind(usage.cache_read_cost_usd)
|
||||
.bind(usage.output_price_per_1m)
|
||||
.bind(usage.total_cost_usd)
|
||||
.bind(usage.actual_total_cost_usd)
|
||||
.bind(usage.status_code.map(i32::from))
|
||||
.bind(&usage.error_message)
|
||||
.bind(&usage.error_category)
|
||||
.bind(usage.response_time_ms.map(to_i32).transpose()?)
|
||||
.bind(usage.first_byte_time_ms.map(to_i32).transpose()?)
|
||||
.bind(&usage.status)
|
||||
.bind(&usage.billing_status)
|
||||
.bind(&usage.request_headers)
|
||||
.bind(&usage.request_body)
|
||||
.bind(&usage.provider_request_headers)
|
||||
.bind(&usage.provider_request_body)
|
||||
.bind(&usage.response_headers)
|
||||
.bind(&usage.response_body)
|
||||
.bind(&usage.client_response_headers)
|
||||
.bind(&usage.client_response_body)
|
||||
.bind(&usage.request_metadata)
|
||||
.bind(usage.finalized_at_unix_secs.map(|value| value as f64))
|
||||
.bind(usage.created_at_unix_secs.map(|value| value as f64))
|
||||
.fetch_one(&mut **tx)
|
||||
.await?;
|
||||
map_usage_row(&row)
|
||||
}) as BoxFuture<'_, Result<StoredRequestUsageAudit, DataLayerError>>
|
||||
})
|
||||
.await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageReadRepository for SqlxUsageReadRepository {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::find_by_id(self, id).await
|
||||
}
|
||||
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::find_by_request_id(self, request_id).await
|
||||
}
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::list_usage_audits(self, query).await
|
||||
}
|
||||
|
||||
async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, DataLayerError> {
|
||||
Self::list_recent_usage_audits(self, user_id, limit).await
|
||||
}
|
||||
|
||||
async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, DataLayerError> {
|
||||
Self::summarize_total_tokens_by_api_key_ids(self, api_key_ids).await
|
||||
}
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, DataLayerError> {
|
||||
Self::summarize_provider_usage_since(self, provider_id, since_unix_secs).await
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl UsageWriteRepository for SqlxUsageReadRepository {
|
||||
async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
Self::upsert(self, usage).await
|
||||
}
|
||||
}
|
||||
|
||||
fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit, DataLayerError> {
|
||||
StoredRequestUsageAudit::new(
|
||||
let mut usage = StoredRequestUsageAudit::new(
|
||||
row.try_get("id")?,
|
||||
row.try_get("request_id")?,
|
||||
row.try_get("user_id")?,
|
||||
@@ -121,13 +746,39 @@ fn map_usage_row(row: &sqlx::postgres::PgRow) -> Result<StoredRequestUsageAudit,
|
||||
row.try_get("created_at_unix_secs")?,
|
||||
row.try_get("updated_at_unix_secs")?,
|
||||
row.try_get("finalized_at_unix_secs")?,
|
||||
)
|
||||
)?;
|
||||
usage.cache_creation_input_tokens = row
|
||||
.try_get::<Option<i32>, _>("cache_creation_input_tokens")?
|
||||
.map(|value| to_u64(value, "usage.cache_creation_input_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
usage.cache_read_input_tokens = row
|
||||
.try_get::<Option<i32>, _>("cache_read_input_tokens")?
|
||||
.map(|value| to_u64(value, "usage.cache_read_input_tokens"))
|
||||
.transpose()?
|
||||
.unwrap_or_default();
|
||||
usage.cache_creation_cost_usd = row.try_get::<f64, _>("cache_creation_cost_usd")?;
|
||||
usage.cache_read_cost_usd = row.try_get::<f64, _>("cache_read_cost_usd")?;
|
||||
usage.output_price_per_1m = row.try_get("output_price_per_1m")?;
|
||||
Ok(usage)
|
||||
}
|
||||
|
||||
fn to_i32(value: u64) -> Result<i32, DataLayerError> {
|
||||
i32::try_from(value).map_err(|_| {
|
||||
DataLayerError::UnexpectedValue(format!("invalid usage integer value: {value}"))
|
||||
})
|
||||
}
|
||||
|
||||
fn to_u64(value: i32, field_name: &str) -> Result<u64, DataLayerError> {
|
||||
u64::try_from(value)
|
||||
.map_err(|_| DataLayerError::UnexpectedValue(format!("invalid {field_name}: {value}")))
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::SqlxUsageReadRepository;
|
||||
use crate::postgres::{PostgresPoolConfig, PostgresPoolFactory};
|
||||
use crate::repository::usage::UpsertUsageRecord;
|
||||
|
||||
#[tokio::test]
|
||||
async fn repository_constructs_from_lazy_pool() {
|
||||
@@ -146,5 +797,98 @@ mod tests {
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxUsageReadRepository::new(pool);
|
||||
let _ = repository.pool();
|
||||
let _ = repository.transaction_runner();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn validates_upsert_before_hitting_database() {
|
||||
let factory = PostgresPoolFactory::new(PostgresPoolConfig {
|
||||
database_url: "postgres://localhost/aether".to_string(),
|
||||
min_connections: 1,
|
||||
max_connections: 4,
|
||||
acquire_timeout_ms: 1_000,
|
||||
idle_timeout_ms: 5_000,
|
||||
max_lifetime_ms: 30_000,
|
||||
statement_cache_capacity: 64,
|
||||
require_ssl: false,
|
||||
})
|
||||
.expect("factory should build");
|
||||
|
||||
let pool = factory.connect_lazy().expect("pool should build");
|
||||
let repository = SqlxUsageReadRepository::new(pool);
|
||||
let result = repository
|
||||
.upsert(UpsertUsageRecord {
|
||||
request_id: "".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "openai".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: None,
|
||||
provider_id: None,
|
||||
provider_endpoint_id: None,
|
||||
provider_api_key_id: None,
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(false),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: None,
|
||||
actual_total_cost_usd: None,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(100),
|
||||
first_byte_time_ms: None,
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: None,
|
||||
request_body: None,
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
})
|
||||
.await;
|
||||
|
||||
assert!(result.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_does_not_require_updated_at_column() {
|
||||
assert!(!super::FIND_BY_REQUEST_ID_SQL.contains("COALESCE(updated_at, created_at)"));
|
||||
assert!(!super::LIST_USAGE_AUDITS_PREFIX.contains("COALESCE(updated_at, created_at)"));
|
||||
assert!(!super::UPSERT_SQL.contains("\n updated_at\n"));
|
||||
assert!(!super::UPSERT_SQL.contains("updated_at = CASE"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_summarizes_tokens_by_api_key_ids_in_database() {
|
||||
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_API_KEY_IDS_SQL.contains("GROUP BY api_key_id"));
|
||||
assert!(super::SUMMARIZE_TOTAL_TOKENS_BY_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\""));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
use async_trait::async_trait;
|
||||
use serde_json::Value;
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredRequestUsageAudit {
|
||||
@@ -26,6 +27,11 @@ pub struct StoredRequestUsageAudit {
|
||||
pub input_tokens: u64,
|
||||
pub output_tokens: u64,
|
||||
pub total_tokens: u64,
|
||||
pub cache_creation_input_tokens: u64,
|
||||
pub cache_read_input_tokens: u64,
|
||||
pub cache_creation_cost_usd: f64,
|
||||
pub cache_read_cost_usd: f64,
|
||||
pub output_price_per_1m: Option<f64>,
|
||||
pub total_cost_usd: f64,
|
||||
pub actual_total_cost_usd: f64,
|
||||
pub status_code: Option<u16>,
|
||||
@@ -141,6 +147,11 @@ impl StoredRequestUsageAudit {
|
||||
input_tokens: parse_u64(input_tokens, "usage.input_tokens")?,
|
||||
output_tokens: parse_u64(output_tokens, "usage.output_tokens")?,
|
||||
total_tokens: parse_u64(total_tokens, "usage.total_tokens")?,
|
||||
cache_creation_input_tokens: 0,
|
||||
cache_read_input_tokens: 0,
|
||||
cache_creation_cost_usd: 0.0,
|
||||
cache_read_cost_usd: 0.0,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd,
|
||||
actual_total_cost_usd,
|
||||
status_code: parse_u16(status_code, "usage.status_code")?,
|
||||
@@ -163,19 +174,258 @@ impl StoredRequestUsageAudit {
|
||||
.transpose()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn with_cache_input_tokens(
|
||||
mut self,
|
||||
cache_creation_input_tokens: u64,
|
||||
cache_read_input_tokens: u64,
|
||||
) -> Self {
|
||||
self.cache_creation_input_tokens = cache_creation_input_tokens;
|
||||
self.cache_read_input_tokens = cache_read_input_tokens;
|
||||
self
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderUsageWindow {
|
||||
pub provider_id: String,
|
||||
pub window_start_unix_secs: u64,
|
||||
pub total_requests: u64,
|
||||
pub successful_requests: u64,
|
||||
pub failed_requests: u64,
|
||||
pub avg_response_time_ms: f64,
|
||||
pub total_cost_usd: f64,
|
||||
}
|
||||
|
||||
impl StoredProviderUsageWindow {
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
pub fn new(
|
||||
provider_id: String,
|
||||
window_start_unix_secs: i64,
|
||||
total_requests: i64,
|
||||
successful_requests: i64,
|
||||
failed_requests: i64,
|
||||
avg_response_time_ms: f64,
|
||||
total_cost_usd: f64,
|
||||
) -> Result<Self, crate::DataLayerError> {
|
||||
if provider_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider usage window provider_id is empty".to_string(),
|
||||
));
|
||||
}
|
||||
if !avg_response_time_ms.is_finite() || !total_cost_usd.is_finite() {
|
||||
return Err(crate::DataLayerError::UnexpectedValue(
|
||||
"provider usage window value is not finite".to_string(),
|
||||
));
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
provider_id,
|
||||
window_start_unix_secs: parse_timestamp(
|
||||
window_start_unix_secs,
|
||||
"provider_usage_tracking.window_start_unix_secs",
|
||||
)?,
|
||||
total_requests: parse_timestamp(
|
||||
total_requests,
|
||||
"provider_usage_tracking.total_requests",
|
||||
)?,
|
||||
successful_requests: parse_timestamp(
|
||||
successful_requests,
|
||||
"provider_usage_tracking.successful_requests",
|
||||
)?,
|
||||
failed_requests: parse_timestamp(
|
||||
failed_requests,
|
||||
"provider_usage_tracking.failed_requests",
|
||||
)?,
|
||||
avg_response_time_ms,
|
||||
total_cost_usd,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct StoredProviderUsageSummary {
|
||||
pub total_requests: u64,
|
||||
pub successful_requests: u64,
|
||||
pub failed_requests: u64,
|
||||
pub avg_response_time_ms: f64,
|
||||
pub total_cost_usd: f64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UsageAuditListQuery {
|
||||
pub created_from_unix_secs: Option<u64>,
|
||||
pub created_until_unix_secs: Option<u64>,
|
||||
pub user_id: Option<String>,
|
||||
pub provider_name: Option<String>,
|
||||
pub model: Option<String>,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageReadRepository: Send + Sync {
|
||||
async fn find_by_id(
|
||||
&self,
|
||||
id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn find_by_request_id(
|
||||
&self,
|
||||
request_id: &str,
|
||||
) -> Result<Option<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn list_usage_audits(
|
||||
&self,
|
||||
query: &UsageAuditListQuery,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn list_recent_usage_audits(
|
||||
&self,
|
||||
user_id: Option<&str>,
|
||||
limit: usize,
|
||||
) -> Result<Vec<StoredRequestUsageAudit>, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_total_tokens_by_api_key_ids(
|
||||
&self,
|
||||
api_key_ids: &[String],
|
||||
) -> Result<std::collections::BTreeMap<String, u64>, crate::DataLayerError>;
|
||||
|
||||
async fn summarize_provider_usage_since(
|
||||
&self,
|
||||
provider_id: &str,
|
||||
since_unix_secs: u64,
|
||||
) -> Result<StoredProviderUsageSummary, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait UsageRepository: UsageReadRepository + Send + Sync {}
|
||||
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
|
||||
pub struct UpsertUsageRecord {
|
||||
pub request_id: String,
|
||||
pub user_id: Option<String>,
|
||||
pub api_key_id: Option<String>,
|
||||
pub username: Option<String>,
|
||||
pub api_key_name: Option<String>,
|
||||
pub provider_name: String,
|
||||
pub model: String,
|
||||
pub target_model: Option<String>,
|
||||
pub provider_id: Option<String>,
|
||||
pub provider_endpoint_id: Option<String>,
|
||||
pub provider_api_key_id: Option<String>,
|
||||
pub request_type: Option<String>,
|
||||
pub api_format: Option<String>,
|
||||
pub api_family: Option<String>,
|
||||
pub endpoint_kind: Option<String>,
|
||||
pub endpoint_api_format: Option<String>,
|
||||
pub provider_api_family: Option<String>,
|
||||
pub provider_endpoint_kind: Option<String>,
|
||||
pub has_format_conversion: Option<bool>,
|
||||
pub is_stream: Option<bool>,
|
||||
pub input_tokens: Option<u64>,
|
||||
pub output_tokens: Option<u64>,
|
||||
pub total_tokens: Option<u64>,
|
||||
pub cache_creation_input_tokens: Option<u64>,
|
||||
pub cache_read_input_tokens: Option<u64>,
|
||||
pub cache_creation_cost_usd: Option<f64>,
|
||||
pub cache_read_cost_usd: Option<f64>,
|
||||
pub output_price_per_1m: Option<f64>,
|
||||
pub total_cost_usd: Option<f64>,
|
||||
pub actual_total_cost_usd: Option<f64>,
|
||||
pub status_code: Option<u16>,
|
||||
pub error_message: Option<String>,
|
||||
pub error_category: Option<String>,
|
||||
pub response_time_ms: Option<u64>,
|
||||
pub first_byte_time_ms: Option<u64>,
|
||||
pub status: String,
|
||||
pub billing_status: String,
|
||||
pub request_headers: Option<Value>,
|
||||
pub request_body: Option<Value>,
|
||||
pub provider_request_headers: Option<Value>,
|
||||
pub provider_request_body: Option<Value>,
|
||||
pub response_headers: Option<Value>,
|
||||
pub response_body: Option<Value>,
|
||||
pub client_response_headers: Option<Value>,
|
||||
pub client_response_body: Option<Value>,
|
||||
pub request_metadata: Option<Value>,
|
||||
pub finalized_at_unix_secs: Option<u64>,
|
||||
pub created_at_unix_secs: Option<u64>,
|
||||
pub updated_at_unix_secs: u64,
|
||||
}
|
||||
|
||||
impl<T> UsageRepository for T where T: UsageReadRepository + Send + Sync {}
|
||||
impl UpsertUsageRecord {
|
||||
pub fn validate(&self) -> Result<(), crate::DataLayerError> {
|
||||
if self.request_id.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert request_id cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.provider_name.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert provider_name cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.model.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert model cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert status cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if self.billing_status.trim().is_empty() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert billing_status cannot be empty".to_string(),
|
||||
));
|
||||
}
|
||||
if let Some(value) = self.total_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert total_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.cache_creation_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert cache_creation_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.cache_read_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert cache_read_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.output_price_per_1m {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert output_price_per_1m must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
if let Some(value) = self.actual_total_cost_usd {
|
||||
if !value.is_finite() {
|
||||
return Err(crate::DataLayerError::InvalidInput(
|
||||
"usage upsert actual_total_cost_usd must be finite".to_string(),
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait UsageWriteRepository: Send + Sync {
|
||||
async fn upsert(
|
||||
&self,
|
||||
usage: UpsertUsageRecord,
|
||||
) -> Result<StoredRequestUsageAudit, crate::DataLayerError>;
|
||||
}
|
||||
|
||||
pub trait UsageRepository: UsageReadRepository + UsageWriteRepository + Send + Sync {}
|
||||
|
||||
impl<T> UsageRepository for T where T: UsageReadRepository + UsageWriteRepository + Send + Sync {}
|
||||
|
||||
fn parse_u64(value: i32, field_name: &str) -> Result<u64, crate::DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| {
|
||||
@@ -214,7 +464,8 @@ fn parse_timestamp(value: i64, field_name: &str) -> Result<u64, crate::DataLayer
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::StoredRequestUsageAudit;
|
||||
use super::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||
use serde_json::json;
|
||||
|
||||
#[test]
|
||||
fn rejects_empty_request_id() {
|
||||
@@ -301,4 +552,61 @@ mod tests {
|
||||
)
|
||||
.is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn rejects_invalid_upsert_payload() {
|
||||
let record = UpsertUsageRecord {
|
||||
request_id: "".to_string(),
|
||||
user_id: None,
|
||||
api_key_id: None,
|
||||
username: None,
|
||||
api_key_name: None,
|
||||
provider_name: "openai".to_string(),
|
||||
model: "gpt-5".to_string(),
|
||||
target_model: None,
|
||||
provider_id: None,
|
||||
provider_endpoint_id: None,
|
||||
provider_api_key_id: None,
|
||||
request_type: Some("chat".to_string()),
|
||||
api_format: Some("openai:chat".to_string()),
|
||||
api_family: Some("openai".to_string()),
|
||||
endpoint_kind: Some("chat".to_string()),
|
||||
endpoint_api_format: Some("openai:chat".to_string()),
|
||||
provider_api_family: Some("openai".to_string()),
|
||||
provider_endpoint_kind: Some("chat".to_string()),
|
||||
has_format_conversion: Some(false),
|
||||
is_stream: Some(false),
|
||||
input_tokens: Some(10),
|
||||
output_tokens: Some(20),
|
||||
total_tokens: Some(30),
|
||||
cache_creation_input_tokens: None,
|
||||
cache_read_input_tokens: None,
|
||||
cache_creation_cost_usd: None,
|
||||
cache_read_cost_usd: None,
|
||||
output_price_per_1m: None,
|
||||
total_cost_usd: None,
|
||||
actual_total_cost_usd: None,
|
||||
status_code: Some(200),
|
||||
error_message: None,
|
||||
error_category: None,
|
||||
response_time_ms: Some(120),
|
||||
first_byte_time_ms: None,
|
||||
status: "completed".to_string(),
|
||||
billing_status: "pending".to_string(),
|
||||
request_headers: Some(json!({"authorization": "Bearer test"})),
|
||||
request_body: Some(json!({"model": "gpt-5"})),
|
||||
provider_request_headers: None,
|
||||
provider_request_body: None,
|
||||
response_headers: None,
|
||||
response_body: None,
|
||||
client_response_headers: None,
|
||||
client_response_body: None,
|
||||
request_metadata: None,
|
||||
finalized_at_unix_secs: None,
|
||||
created_at_unix_secs: Some(100),
|
||||
updated_at_unix_secs: 101,
|
||||
};
|
||||
|
||||
assert!(record.validate().is_err());
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user