mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-10 13:10:21 +08:00
fix: read imported usage aggregates in dashboards
This commit is contained in:
fawney19
@@ -108,8 +108,7 @@ macro_rules! impl_materialized_usage_read_repository {
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<$crate::repository::usage::StoredUsageUserTotals>, $crate::DataLayerError>
|
||||
{
|
||||
let repository = self.materialize_read_model().await?;
|
||||
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_totals_by_user_ids(&repository, user_ids).await
|
||||
<$repository>::summarize_usage_totals_by_user_ids(self, user_ids).await
|
||||
}
|
||||
|
||||
async fn summarize_usage_cache_hit_summary(
|
||||
@@ -163,8 +162,7 @@ macro_rules! impl_materialized_usage_read_repository {
|
||||
$crate::repository::usage::StoredUsageDashboardSummary,
|
||||
$crate::DataLayerError,
|
||||
> {
|
||||
let repository = self.materialize_read_model().await?;
|
||||
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_dashboard_usage(&repository, query).await
|
||||
<$repository>::summarize_dashboard_usage(self, query).await
|
||||
}
|
||||
|
||||
async fn list_dashboard_daily_breakdown(
|
||||
@@ -174,8 +172,7 @@ macro_rules! impl_materialized_usage_read_repository {
|
||||
Vec<$crate::repository::usage::StoredUsageDashboardDailyBreakdownRow>,
|
||||
$crate::DataLayerError,
|
||||
> {
|
||||
let repository = self.materialize_read_model().await?;
|
||||
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::list_dashboard_daily_breakdown(&repository, query).await
|
||||
<$repository>::list_dashboard_daily_breakdown(self, query).await
|
||||
}
|
||||
|
||||
async fn summarize_dashboard_provider_counts(
|
||||
@@ -349,8 +346,7 @@ macro_rules! impl_materialized_usage_read_repository {
|
||||
Vec<$crate::repository::usage::StoredUsageDailySummary>,
|
||||
$crate::DataLayerError,
|
||||
> {
|
||||
let repository = self.materialize_read_model().await?;
|
||||
<$crate::repository::usage::InMemoryUsageReadRepository as $crate::repository::usage::UsageReadRepository>::summarize_usage_daily_heatmap(&repository, query).await
|
||||
<$repository>::summarize_usage_daily_heatmap(self, query).await
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
@@ -1,14 +1,18 @@
|
||||
use std::collections::{BTreeMap, HashSet};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
|
||||
use async_trait::async_trait;
|
||||
use sqlx::{mysql::MySqlRow, Row};
|
||||
use sqlx::{mysql::MySqlRow, MySql, QueryBuilder, Row};
|
||||
|
||||
use super::{
|
||||
provider_api_key_usage_is_error, provider_api_key_usage_is_success,
|
||||
strip_deprecated_usage_display_fields, usage_can_recover_terminal_failure,
|
||||
usage_request_metadata_client_family, InMemoryUsageReadRepository, PendingUsageCleanupSummary,
|
||||
StoredRequestUsageAudit, UpsertUsageRecord, UsageWriteRepository,
|
||||
StoredRequestUsageAudit, StoredUsageDailySummary, StoredUsageDashboardDailyBreakdownRow,
|
||||
StoredUsageDashboardSummary, StoredUsageUserTotals, UpsertUsageRecord, UsageDailyHeatmapQuery,
|
||||
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageReadRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
use crate::driver::mysql::MysqlPool;
|
||||
use crate::error::SqlResultExt;
|
||||
@@ -229,6 +233,467 @@ impl MysqlUsageReadRepository {
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
Ok(InMemoryUsageReadRepository::seed(items))
|
||||
}
|
||||
|
||||
async fn summarize_usage_daily_heatmap_raw_from_range(
|
||||
&self,
|
||||
created_from_unix_secs: u64,
|
||||
created_until_unix_secs: u64,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||||
let mut sql = String::from(
|
||||
r#"
|
||||
SELECT
|
||||
DATE_FORMAT(FROM_UNIXTIME(created_at_unix_ms), '%Y-%m-%d') AS date,
|
||||
COUNT(*) AS requests,
|
||||
COALESCE(SUM(
|
||||
GREATEST(COALESCE(input_tokens, 0), 0)
|
||||
+ GREATEST(COALESCE(output_tokens, 0), 0)
|
||||
+ CASE
|
||||
WHEN COALESCE(cache_creation_input_tokens, 0) = 0
|
||||
AND (COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0)) > 0
|
||||
THEN COALESCE(cache_creation_ephemeral_5m_input_tokens, 0) + COALESCE(cache_creation_ephemeral_1h_input_tokens, 0)
|
||||
ELSE GREATEST(COALESCE(cache_creation_input_tokens, 0), 0)
|
||||
END
|
||||
+ GREATEST(COALESCE(cache_read_input_tokens, 0), 0)
|
||||
), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(total_cost_usd, 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(actual_total_cost_usd, 0)), 0) AS actual_total_cost_usd
|
||||
FROM `usage`
|
||||
WHERE created_at_unix_ms >= ?
|
||||
AND created_at_unix_ms < ?
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
"#,
|
||||
);
|
||||
if user_id.is_some() {
|
||||
sql.push_str(" AND user_id = ?\n");
|
||||
}
|
||||
sql.push_str("GROUP BY date ORDER BY date ASC");
|
||||
|
||||
let mut query = sqlx::query(&sql)
|
||||
.bind(to_i64(created_from_unix_secs, "usage.created_at_unix_ms")?)
|
||||
.bind(to_i64(created_until_unix_secs, "usage.created_at_unix_ms")?);
|
||||
if let Some(user_id) = user_id {
|
||||
query = query.bind(user_id.to_string());
|
||||
}
|
||||
let rows = query.fetch_all(&self.pool).await.map_sql_err()?;
|
||||
rows.iter().map(map_mysql_usage_daily_summary).collect()
|
||||
}
|
||||
|
||||
async fn summarize_usage_daily_heatmap_from_daily_aggregates(
|
||||
&self,
|
||||
created_from_unix_secs: u64,
|
||||
created_until_unix_secs: u64,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||||
let rows = if let Some(user_id) = user_id {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
||||
total_requests AS requests,
|
||||
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
|
||||
total_cost AS total_cost_usd,
|
||||
total_cost AS actual_total_cost_usd
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ?
|
||||
AND `date` >= ?
|
||||
AND `date` < ?
|
||||
AND total_requests > 0
|
||||
ORDER BY `date` ASC
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(to_i64(created_from_unix_secs, "stats_user_daily.date")?)
|
||||
.bind(to_i64(created_until_unix_secs, "stats_user_daily.date")?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
||||
total_requests AS requests,
|
||||
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
|
||||
total_cost AS total_cost_usd,
|
||||
actual_total_cost AS actual_total_cost_usd
|
||||
FROM stats_daily
|
||||
WHERE `date` >= ?
|
||||
AND `date` < ?
|
||||
AND total_requests > 0
|
||||
ORDER BY `date` ASC
|
||||
"#,
|
||||
)
|
||||
.bind(to_i64(created_from_unix_secs, "stats_daily.date")?)
|
||||
.bind(to_i64(created_until_unix_secs, "stats_daily.date")?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
};
|
||||
|
||||
rows.iter().map(map_mysql_usage_daily_summary).collect()
|
||||
}
|
||||
|
||||
async fn summarize_usage_daily_heatmap(
|
||||
&self,
|
||||
query: &UsageDailyHeatmapQuery,
|
||||
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||||
let created_until_unix_secs = usage_current_unix_secs().saturating_add(1);
|
||||
let user_id = query.user_id.as_deref();
|
||||
let mut summaries = BTreeMap::<String, StoredUsageDailySummary>::new();
|
||||
|
||||
for item in self
|
||||
.summarize_usage_daily_heatmap_from_daily_aggregates(
|
||||
query.created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
user_id,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
summaries.insert(item.date.clone(), item);
|
||||
}
|
||||
for item in self
|
||||
.summarize_usage_daily_heatmap_raw_from_range(
|
||||
query.created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
user_id,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
summaries.entry(item.date.clone()).or_insert(item);
|
||||
}
|
||||
|
||||
Ok(summaries.into_values().collect())
|
||||
}
|
||||
|
||||
async fn summarize_dashboard_usage_from_daily_aggregates(
|
||||
&self,
|
||||
query: &UsageDashboardSummaryQuery,
|
||||
) -> Result<Option<StoredUsageDashboardSummary>, DataLayerError> {
|
||||
let row = if let Some(user_id) = query.user_id.as_deref() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
|
||||
0.0 AS cache_creation_cost_usd,
|
||||
0.0 AS cache_read_cost_usd,
|
||||
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS actual_total_cost_usd,
|
||||
COALESCE(SUM(error_requests), 0) AS error_requests,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ?
|
||||
AND `date` >= ?
|
||||
AND `date` < ?
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(to_i64(
|
||||
query.created_from_unix_secs,
|
||||
"stats_user_daily.date",
|
||||
)?)
|
||||
.bind(to_i64(
|
||||
query.created_until_unix_secs,
|
||||
"stats_user_daily.date",
|
||||
)?)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
|
||||
COALESCE(SUM(COALESCE(cache_creation_cost, 0)), 0) AS cache_creation_cost_usd,
|
||||
COALESCE(SUM(COALESCE(cache_read_cost, 0)), 0) AS cache_read_cost_usd,
|
||||
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(actual_total_cost, 0)), 0) AS actual_total_cost_usd,
|
||||
COALESCE(SUM(error_requests), 0) AS error_requests,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_daily
|
||||
WHERE `date` >= ?
|
||||
AND `date` < ?
|
||||
"#,
|
||||
)
|
||||
.bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?)
|
||||
.bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
};
|
||||
|
||||
let total_requests = row_u64(&row, "total_requests")?;
|
||||
if total_requests == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(StoredUsageDashboardSummary {
|
||||
total_requests,
|
||||
input_tokens: row_u64(&row, "input_tokens")?,
|
||||
effective_input_tokens: row_u64(&row, "effective_input_tokens")?,
|
||||
output_tokens: row_u64(&row, "output_tokens")?,
|
||||
total_tokens: row_u64(&row, "total_tokens")?,
|
||||
cache_creation_tokens: row_u64(&row, "cache_creation_tokens")?,
|
||||
cache_read_tokens: row_u64(&row, "cache_read_tokens")?,
|
||||
total_input_context: row_u64(&row, "total_input_context")?,
|
||||
cache_creation_cost_usd: row.try_get("cache_creation_cost_usd").map_sql_err()?,
|
||||
cache_read_cost_usd: row.try_get("cache_read_cost_usd").map_sql_err()?,
|
||||
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
|
||||
actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?,
|
||||
error_requests: row_u64(&row, "error_requests")?,
|
||||
response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?,
|
||||
response_time_samples: row_u64(&row, "response_time_samples")?,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn list_dashboard_daily_breakdown_from_daily_aggregates(
|
||||
&self,
|
||||
query: &UsageDashboardDailyBreakdownQuery,
|
||||
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||||
let rows = if let Some(user_id) = query.user_id.as_deref() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
||||
'aggregate' AS model,
|
||||
'aggregate' AS provider,
|
||||
COALESCE(SUM(total_requests), 0) AS requests,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ?
|
||||
AND `date` >= ?
|
||||
AND `date` < ?
|
||||
AND total_requests > 0
|
||||
GROUP BY `date`
|
||||
ORDER BY `date` ASC
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(to_i64(
|
||||
query.created_from_unix_secs,
|
||||
"stats_user_daily.date",
|
||||
)?)
|
||||
.bind(to_i64(
|
||||
query.created_until_unix_secs,
|
||||
"stats_user_daily.date",
|
||||
)?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
DATE_FORMAT(FROM_UNIXTIME(`date`), '%Y-%m-%d') AS date,
|
||||
'aggregate' AS model,
|
||||
'aggregate' AS provider,
|
||||
COALESCE(SUM(total_requests), 0) AS requests,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(total_cost, 0)), 0) AS total_cost_usd,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_daily
|
||||
WHERE `date` >= ?
|
||||
AND `date` < ?
|
||||
AND total_requests > 0
|
||||
GROUP BY `date`
|
||||
ORDER BY `date` ASC
|
||||
"#,
|
||||
)
|
||||
.bind(to_i64(query.created_from_unix_secs, "stats_daily.date")?)
|
||||
.bind(to_i64(query.created_until_unix_secs, "stats_daily.date")?)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
};
|
||||
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
Ok(StoredUsageDashboardDailyBreakdownRow {
|
||||
date: row.try_get("date").map_sql_err()?,
|
||||
model: row.try_get("model").map_sql_err()?,
|
||||
provider: row.try_get("provider").map_sql_err()?,
|
||||
requests: row_u64(row, "requests")?,
|
||||
total_tokens: row_u64(row, "total_tokens")?,
|
||||
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
|
||||
response_time_sum_ms: row.try_get("response_time_sum_ms").map_sql_err()?,
|
||||
response_time_samples: row_u64(row, "response_time_samples")?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
async fn summarize_dashboard_usage(
|
||||
&self,
|
||||
query: &UsageDashboardSummaryQuery,
|
||||
) -> Result<StoredUsageDashboardSummary, DataLayerError> {
|
||||
if let Some(summary) = self
|
||||
.summarize_dashboard_usage_from_daily_aggregates(query)
|
||||
.await?
|
||||
{
|
||||
return Ok(summary);
|
||||
}
|
||||
|
||||
let repository = self.materialize_read_model().await?;
|
||||
<InMemoryUsageReadRepository as UsageReadRepository>::summarize_dashboard_usage(
|
||||
&repository,
|
||||
query,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn list_dashboard_daily_breakdown(
|
||||
&self,
|
||||
query: &UsageDashboardDailyBreakdownQuery,
|
||||
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||||
let aggregate_rows = self
|
||||
.list_dashboard_daily_breakdown_from_daily_aggregates(query)
|
||||
.await?;
|
||||
if !aggregate_rows.is_empty() {
|
||||
return Ok(aggregate_rows);
|
||||
}
|
||||
|
||||
let repository = self.materialize_read_model().await?;
|
||||
<InMemoryUsageReadRepository as UsageReadRepository>::list_dashboard_daily_breakdown(
|
||||
&repository,
|
||||
query,
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
async fn summarize_usage_totals_by_user_ids(
|
||||
&self,
|
||||
user_ids: &[String],
|
||||
) -> Result<Vec<StoredUsageUserTotals>, DataLayerError> {
|
||||
if user_ids.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let unique_user_ids = user_ids
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
let mut totals = BTreeMap::<String, StoredUsageUserTotals>::new();
|
||||
let mut aggregate_cutoffs = BTreeMap::<String, u64>::new();
|
||||
|
||||
let mut aggregate_builder = QueryBuilder::<MySql>::new(
|
||||
r#"
|
||||
SELECT
|
||||
user_id,
|
||||
COALESCE(SUM(total_requests), 0) AS request_count,
|
||||
COALESCE(
|
||||
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
|
||||
0
|
||||
) AS total_tokens,
|
||||
MAX(`date`) AS latest_date
|
||||
FROM stats_user_daily
|
||||
WHERE user_id IN (
|
||||
"#,
|
||||
);
|
||||
{
|
||||
let mut separated = aggregate_builder.separated(", ");
|
||||
for user_id in &unique_user_ids {
|
||||
separated.push_bind(user_id.clone());
|
||||
}
|
||||
}
|
||||
aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC");
|
||||
|
||||
let aggregate_rows = aggregate_builder
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
for row in aggregate_rows {
|
||||
let user_id: String = row.try_get("user_id").map_sql_err()?;
|
||||
let latest_date = row.try_get::<i64, _>("latest_date").map_sql_err()?.max(0) as u64;
|
||||
aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400));
|
||||
totals.insert(
|
||||
user_id.clone(),
|
||||
StoredUsageUserTotals {
|
||||
user_id,
|
||||
request_count: row_u64(&row, "request_count")?,
|
||||
total_tokens: row_u64(&row, "total_tokens")?,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let mut raw_builder = QueryBuilder::<MySql>::new(
|
||||
r#"
|
||||
SELECT
|
||||
`usage`.user_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(GREATEST(COALESCE(`usage`.total_tokens, 0), 0)), 0) AS total_tokens
|
||||
FROM `usage`
|
||||
JOIN (
|
||||
"#,
|
||||
);
|
||||
for (index, user_id) in unique_user_ids.iter().enumerate() {
|
||||
if index > 0 {
|
||||
raw_builder.push(" UNION ALL ");
|
||||
}
|
||||
let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default();
|
||||
raw_builder
|
||||
.push("SELECT ")
|
||||
.push_bind(user_id.clone())
|
||||
.push(" AS user_id, ")
|
||||
.push_bind(to_i64(cutoff, "usage aggregate cutoff")?)
|
||||
.push(" AS cutoff_unix_secs");
|
||||
}
|
||||
raw_builder.push(
|
||||
r#"
|
||||
) AS requested ON requested.user_id = `usage`.user_id
|
||||
WHERE `usage`.created_at_unix_ms >= requested.cutoff_unix_secs
|
||||
AND `usage`.status NOT IN ('pending', 'streaming')
|
||||
AND `usage`.provider_name NOT IN ('unknown', 'pending')
|
||||
GROUP BY `usage`.user_id
|
||||
ORDER BY `usage`.user_id ASC
|
||||
"#,
|
||||
);
|
||||
|
||||
let raw_rows = raw_builder
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
for row in raw_rows {
|
||||
let user_id: String = row.try_get("user_id").map_sql_err()?;
|
||||
let entry = totals
|
||||
.entry(user_id.clone())
|
||||
.or_insert_with(|| StoredUsageUserTotals {
|
||||
user_id,
|
||||
request_count: 0,
|
||||
total_tokens: 0,
|
||||
});
|
||||
entry.request_count = entry
|
||||
.request_count
|
||||
.saturating_add(row_u64(&row, "request_count")?);
|
||||
entry.total_tokens = entry
|
||||
.total_tokens
|
||||
.saturating_add(row_u64(&row, "total_tokens")?);
|
||||
}
|
||||
|
||||
Ok(totals.into_values().collect())
|
||||
}
|
||||
}
|
||||
|
||||
impl_materialized_usage_read_repository!(MysqlUsageReadRepository);
|
||||
@@ -990,6 +1455,25 @@ fn row_u64(row: &MySqlRow, field: &str) -> Result<u64, DataLayerError> {
|
||||
u64::try_from(value).map_err(|_| DataLayerError::UnexpectedValue(format!("{field} negative")))
|
||||
}
|
||||
|
||||
fn map_mysql_usage_daily_summary(
|
||||
row: &MySqlRow,
|
||||
) -> Result<StoredUsageDailySummary, DataLayerError> {
|
||||
Ok(StoredUsageDailySummary {
|
||||
date: row.try_get("date").map_sql_err()?,
|
||||
requests: row_u64(row, "requests")?,
|
||||
total_tokens: row_u64(row, "total_tokens")?,
|
||||
total_cost_usd: row.try_get("total_cost_usd").map_sql_err()?,
|
||||
actual_total_cost_usd: row.try_get("actual_total_cost_usd").map_sql_err()?,
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::{MysqlUsageReadRepository, MysqlUsageWriteRepository};
|
||||
@@ -1010,6 +1494,34 @@ mod tests {
|
||||
let _repository = MysqlUsageWriteRepository::new(pool);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mysql_usage_daily_heatmap_reads_imported_daily_aggregates() {
|
||||
let source = include_str!("mysql.rs");
|
||||
assert!(source.contains("summarize_usage_daily_heatmap_from_daily_aggregates"));
|
||||
assert!(source.contains("FROM stats_daily"));
|
||||
assert!(source.contains("FROM stats_user_daily"));
|
||||
assert!(source.contains("summaries.entry(item.date.clone()).or_insert(item)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mysql_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() {
|
||||
let source = include_str!("mysql.rs");
|
||||
assert!(source.contains("async fn summarize_usage_totals_by_user_ids"));
|
||||
assert!(source.contains("FROM stats_user_daily"));
|
||||
assert!(source.contains("MAX(`date`) AS latest_date"));
|
||||
assert!(source.contains("requested.cutoff_unix_secs"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn mysql_dashboard_reads_imported_daily_aggregates() {
|
||||
let source = include_str!("mysql.rs");
|
||||
assert!(source.contains("summarize_dashboard_usage_from_daily_aggregates"));
|
||||
assert!(source.contains("list_dashboard_daily_breakdown_from_daily_aggregates"));
|
||||
assert!(source.contains("FROM stats_daily"));
|
||||
assert!(source.contains("FROM stats_user_daily"));
|
||||
assert!(source.contains("'aggregate' AS model"));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn mysql_usage_write_repository_upserts_when_url_is_set() {
|
||||
let Some(database_url) = std::env::var("AETHER_TEST_MYSQL_URL")
|
||||
|
||||
@@ -1827,11 +1827,39 @@ LIMIT 1
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
|
||||
row.map(|row| {
|
||||
row.try_get::<DateTime<Utc>, _>("cutoff_date")
|
||||
.map_postgres_err()
|
||||
})
|
||||
.transpose()
|
||||
if let Some(row) = row {
|
||||
return row
|
||||
.try_get::<DateTime<Utc>, _>("cutoff_date")
|
||||
.map(Some)
|
||||
.map_postgres_err();
|
||||
}
|
||||
|
||||
let row = sqlx::query(
|
||||
r#"
|
||||
SELECT MAX(date) AS latest_date
|
||||
FROM (
|
||||
SELECT MAX(date) AS date
|
||||
FROM stats_daily
|
||||
WHERE total_requests > 0
|
||||
OR is_complete IS TRUE
|
||||
UNION ALL
|
||||
SELECT MAX(date) AS date
|
||||
FROM stats_user_daily
|
||||
WHERE total_requests > 0
|
||||
UNION ALL
|
||||
SELECT MAX(date) AS date
|
||||
FROM stats_daily_api_key
|
||||
WHERE total_requests > 0
|
||||
) AS imported_daily_aggregates
|
||||
"#,
|
||||
)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
let latest_date = row
|
||||
.try_get::<Option<DateTime<Utc>>, _>("latest_date")
|
||||
.map_postgres_err()?;
|
||||
Ok(latest_date.map(|value| value + chrono::Duration::days(1)))
|
||||
}
|
||||
|
||||
async fn read_stats_hourly_cutoff(&self) -> Result<Option<DateTime<Utc>>, DataLayerError> {
|
||||
@@ -1867,9 +1895,22 @@ WHERE is_complete IS TRUE
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
|
||||
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
|
||||
COALESCE(SUM(
|
||||
CASE
|
||||
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
|
||||
THEN input_tokens
|
||||
ELSE effective_input_tokens
|
||||
END
|
||||
), 0)::BIGINT AS effective_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
|
||||
COALESCE(SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
|
||||
COALESCE(SUM(
|
||||
CASE
|
||||
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
|
||||
THEN input_tokens
|
||||
ELSE effective_input_tokens
|
||||
END
|
||||
+ output_tokens + cache_creation_tokens + cache_read_tokens
|
||||
), 0)::BIGINT AS total_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
|
||||
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
|
||||
@@ -1898,9 +1939,22 @@ WHERE user_id = $1
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0)::BIGINT AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0)::BIGINT AS input_tokens,
|
||||
COALESCE(SUM(effective_input_tokens), 0)::BIGINT AS effective_input_tokens,
|
||||
COALESCE(SUM(
|
||||
CASE
|
||||
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
|
||||
THEN input_tokens
|
||||
ELSE effective_input_tokens
|
||||
END
|
||||
), 0)::BIGINT AS effective_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0)::BIGINT AS output_tokens,
|
||||
COALESCE(SUM(effective_input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
|
||||
COALESCE(SUM(
|
||||
CASE
|
||||
WHEN effective_input_tokens = 0 AND total_input_context = 0 AND input_tokens > 0
|
||||
THEN input_tokens
|
||||
ELSE effective_input_tokens
|
||||
END
|
||||
+ output_tokens + cache_creation_tokens + cache_read_tokens
|
||||
), 0)::BIGINT AS total_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens), 0)::BIGINT AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0)::BIGINT AS cache_read_tokens,
|
||||
COALESCE(SUM(total_input_context), 0)::BIGINT AS total_input_context,
|
||||
@@ -1926,6 +1980,75 @@ WHERE date >= $1
|
||||
decode_dashboard_summary_row(&row)
|
||||
}
|
||||
|
||||
async fn list_dashboard_daily_breakdown_from_daily_totals(
|
||||
&self,
|
||||
start_day_utc: DateTime<Utc>,
|
||||
end_day_utc: DateTime<Utc>,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||||
if start_day_utc >= end_day_utc {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let sql = if user_id.is_some() {
|
||||
r#"
|
||||
SELECT
|
||||
TO_CHAR(date, 'YYYY-MM-DD') AS date,
|
||||
'aggregate'::TEXT AS model,
|
||||
'aggregate'::TEXT AS provider,
|
||||
COALESCE(SUM(total_requests), 0)::BIGINT AS requests,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
|
||||
COALESCE(SUM(total_cost), 0)::DOUBLE PRECISION AS total_cost_usd,
|
||||
0::DOUBLE PRECISION AS response_time_sum_ms,
|
||||
0::BIGINT AS response_time_samples
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = $1
|
||||
AND date >= $2
|
||||
AND date < $3
|
||||
AND total_requests > 0
|
||||
GROUP BY date
|
||||
ORDER BY date ASC
|
||||
"#
|
||||
} else {
|
||||
r#"
|
||||
SELECT
|
||||
TO_CHAR(date, 'YYYY-MM-DD') AS date,
|
||||
'aggregate'::TEXT AS model,
|
||||
'aggregate'::TEXT AS provider,
|
||||
COALESCE(SUM(total_requests), 0)::BIGINT AS requests,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0)::BIGINT AS total_tokens,
|
||||
COALESCE(SUM(total_cost), 0)::DOUBLE PRECISION AS total_cost_usd,
|
||||
0::DOUBLE PRECISION AS response_time_sum_ms,
|
||||
0::BIGINT AS response_time_samples
|
||||
FROM stats_daily
|
||||
WHERE date >= $1
|
||||
AND date < $2
|
||||
AND total_requests > 0
|
||||
GROUP BY date
|
||||
ORDER BY date ASC
|
||||
"#
|
||||
};
|
||||
|
||||
let mut rows = if let Some(user_id) = user_id {
|
||||
sqlx::query(sql)
|
||||
.bind(user_id)
|
||||
.bind(start_day_utc)
|
||||
.bind(end_day_utc)
|
||||
.fetch(&self.pool)
|
||||
} else {
|
||||
sqlx::query(sql)
|
||||
.bind(start_day_utc)
|
||||
.bind(end_day_utc)
|
||||
.fetch(&self.pool)
|
||||
};
|
||||
|
||||
let mut items = Vec::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
items.push(decode_dashboard_daily_breakdown_row(&row)?);
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn summarize_dashboard_usage_raw(
|
||||
&self,
|
||||
created_from_unix_secs: u64,
|
||||
@@ -4241,8 +4364,19 @@ ORDER BY date ASC, total_cost_usd DESC, model ASC, provider_name ASC
|
||||
};
|
||||
|
||||
let mut items = Vec::new();
|
||||
let mut detailed_dates = std::collections::BTreeSet::<String>::new();
|
||||
while let Some(row) = rows.try_next().await.map_postgres_err()? {
|
||||
items.push(decode_dashboard_daily_breakdown_row(&row)?);
|
||||
let item = decode_dashboard_daily_breakdown_row(&row)?;
|
||||
detailed_dates.insert(item.date.clone());
|
||||
items.push(item);
|
||||
}
|
||||
for item in self
|
||||
.list_dashboard_daily_breakdown_from_daily_totals(start_day_utc, end_day_utc, user_id)
|
||||
.await?
|
||||
{
|
||||
if !detailed_dates.contains(&item.date) {
|
||||
items.push(item);
|
||||
}
|
||||
}
|
||||
Ok(items)
|
||||
}
|
||||
@@ -4350,12 +4484,53 @@ ORDER BY date ASC, total_cost_usd DESC, "usage".model ASC, "usage".provider_name
|
||||
Ok(items)
|
||||
}
|
||||
|
||||
async fn list_dashboard_daily_breakdown_aggregate_segments(
|
||||
&self,
|
||||
query: &UsageDashboardDailyBreakdownQuery,
|
||||
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||||
let cutoff_utc = match self.read_stats_daily_cutoff_date().await {
|
||||
Ok(value) => value,
|
||||
Err(err) if dashboard_should_fallback_to_raw_on_aggregate_error(&err) => {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
Err(err) => return Err(err),
|
||||
};
|
||||
let Some(cutoff_utc) = cutoff_utc else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
let start_utc = dashboard_unix_secs_to_utc(query.created_from_unix_secs);
|
||||
let end_utc = dashboard_unix_secs_to_utc(query.created_until_unix_secs);
|
||||
let split = split_dashboard_daily_aggregate_range(start_utc, end_utc, cutoff_utc);
|
||||
let Some((aggregate_start, aggregate_end)) = split.aggregate else {
|
||||
return Ok(Vec::new());
|
||||
};
|
||||
|
||||
self.list_dashboard_daily_breakdown_from_daily_aggregates(
|
||||
aggregate_start,
|
||||
aggregate_end,
|
||||
query.user_id.as_deref(),
|
||||
)
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn list_dashboard_daily_breakdown(
|
||||
&self,
|
||||
query: &UsageDashboardDailyBreakdownQuery,
|
||||
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||||
if query.tz_offset_minutes != 0 {
|
||||
return self.list_dashboard_daily_breakdown_raw(query).await;
|
||||
let mut items = self
|
||||
.list_dashboard_daily_breakdown_aggregate_segments(query)
|
||||
.await?;
|
||||
let mut aggregate_dates = items
|
||||
.iter()
|
||||
.map(|item| item.date.clone())
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
for item in self.list_dashboard_daily_breakdown_raw(query).await? {
|
||||
if aggregate_dates.insert(item.date.clone()) {
|
||||
items.push(item);
|
||||
}
|
||||
}
|
||||
return Ok(finalize_dashboard_daily_breakdown_rows(items));
|
||||
}
|
||||
|
||||
let cutoff_utc = match self.read_stats_daily_cutoff_date().await {
|
||||
@@ -7370,6 +7545,7 @@ ORDER BY api_key_id ASC
|
||||
|
||||
let mut totals = std::collections::BTreeMap::<String, StoredUsageUserTotals>::new();
|
||||
if let Some(cutoff_utc) = self.read_stats_daily_cutoff_date().await? {
|
||||
let mut summary_user_ids = std::collections::BTreeSet::<String>::new();
|
||||
let mut aggregate_rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
@@ -7392,6 +7568,50 @@ ORDER BY user_id ASC
|
||||
|
||||
while let Some(row) = aggregate_rows.try_next().await.map_postgres_err()? {
|
||||
let user_id = row.try_get::<String, _>("user_id").map_postgres_err()?;
|
||||
let request_count = row
|
||||
.try_get::<i64, _>("request_count")
|
||||
.map_postgres_err()?
|
||||
.max(0) as u64;
|
||||
let total_tokens = row
|
||||
.try_get::<i64, _>("total_tokens")
|
||||
.map_postgres_err()?
|
||||
.max(0) as u64;
|
||||
summary_user_ids.insert(user_id.clone());
|
||||
totals.insert(
|
||||
user_id.clone(),
|
||||
StoredUsageUserTotals {
|
||||
user_id,
|
||||
request_count,
|
||||
total_tokens,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let mut daily_rows = sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
user_id,
|
||||
COALESCE(SUM(total_requests), 0)::BIGINT AS request_count,
|
||||
COALESCE(
|
||||
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
|
||||
0
|
||||
)::BIGINT AS total_tokens
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ANY($1::TEXT[])
|
||||
AND date < $2
|
||||
GROUP BY user_id
|
||||
ORDER BY user_id ASC
|
||||
"#,
|
||||
)
|
||||
.bind(user_ids)
|
||||
.bind(cutoff_utc)
|
||||
.fetch(&self.pool);
|
||||
|
||||
while let Some(row) = daily_rows.try_next().await.map_postgres_err()? {
|
||||
let user_id = row.try_get::<String, _>("user_id").map_postgres_err()?;
|
||||
if summary_user_ids.contains(&user_id) {
|
||||
continue;
|
||||
}
|
||||
let request_count = row
|
||||
.try_get::<i64, _>("request_count")
|
||||
.map_postgres_err()?
|
||||
|
||||
@@ -436,6 +436,30 @@ fn usage_sql_summarize_usage_daily_heatmap_supports_daily_aggregates() {
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_daily_cutoff_falls_back_to_imported_stats_daily() {
|
||||
let source = include_str!("mod.rs");
|
||||
assert!(source.contains("FROM stats_summary"));
|
||||
assert!(source.contains("SELECT MAX(date) AS latest_date"));
|
||||
assert!(source.contains("FROM stats_daily"));
|
||||
assert!(source.contains("FROM stats_user_daily"));
|
||||
assert!(source.contains("FROM stats_daily_api_key"));
|
||||
assert!(source.contains("value + chrono::Duration::days(1)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_dashboard_daily_breakdown_falls_back_to_daily_totals() {
|
||||
let source = include_str!("mod.rs");
|
||||
assert!(source.contains("list_dashboard_daily_breakdown_aggregate_segments"));
|
||||
assert!(source.contains("list_dashboard_daily_breakdown_from_daily_totals"));
|
||||
assert!(source.contains("'aggregate'::TEXT AS model"));
|
||||
assert!(source.contains("FROM stats_daily"));
|
||||
assert!(source.contains("FROM stats_user_daily"));
|
||||
assert!(source.contains("detailed_dates.contains(&item.date)"));
|
||||
assert!(source.contains("query.tz_offset_minutes != 0"));
|
||||
assert!(source.contains("aggregate_dates.insert(item.date.clone())"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn usage_sql_summarize_usage_leaderboard_supports_daily_aggregates() {
|
||||
let source = include_str!("mod.rs");
|
||||
@@ -533,6 +557,9 @@ fn usage_sql_summarize_usage_totals_by_user_ids_supports_user_summary_aggregates
|
||||
let source = include_str!("mod.rs");
|
||||
assert!(source.contains("FROM stats_user_summary"));
|
||||
assert!(source.contains("all_time_input_tokens"));
|
||||
assert!(source.contains("FROM stats_user_daily"));
|
||||
assert!(source.contains("date < $2"));
|
||||
assert!(source.contains("summary_user_ids.contains(&user_id)"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use std::collections::{BTreeMap, HashSet};
|
||||
use std::io::Read;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
use aether_ai_formats::UPSTREAM_IS_STREAM_KEY;
|
||||
use aether_data_contracts::repository::usage::{parse_usage_body_ref, UsageBodyField};
|
||||
@@ -960,6 +961,102 @@ impl SqliteUsageReadRepository {
|
||||
Self { pool }
|
||||
}
|
||||
|
||||
async fn summarize_usage_daily_heatmap_raw_from_range(
|
||||
&self,
|
||||
created_from_unix_secs: u64,
|
||||
created_until_unix_secs: u64,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(format!(
|
||||
r#"
|
||||
SELECT
|
||||
date(created_at_unix_ms, 'unixepoch') AS date,
|
||||
COUNT(*) AS requests,
|
||||
COALESCE(SUM(
|
||||
MAX(COALESCE(input_tokens, 0), 0)
|
||||
+ MAX(COALESCE(output_tokens, 0), 0)
|
||||
+ {cache_creation_expr}
|
||||
+ MAX(COALESCE(cache_read_input_tokens, 0), 0)
|
||||
), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0)
|
||||
AS actual_total_cost_usd
|
||||
FROM "usage"
|
||||
"#,
|
||||
cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR
|
||||
));
|
||||
let mut has_where = false;
|
||||
push_sqlite_usage_where(&mut builder, &mut has_where);
|
||||
builder
|
||||
.push("created_at_unix_ms >= ")
|
||||
.push_bind(created_from_unix_secs as i64);
|
||||
push_sqlite_usage_where(&mut builder, &mut has_where);
|
||||
builder
|
||||
.push("created_at_unix_ms < ")
|
||||
.push_bind(created_until_unix_secs as i64);
|
||||
push_sqlite_usage_finalized_filter(&mut builder, &mut has_where);
|
||||
push_sqlite_usage_optional_text_filter(&mut builder, &mut has_where, "user_id", user_id);
|
||||
builder.push(" GROUP BY date ORDER BY date ASC");
|
||||
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
rows.iter().map(map_sqlite_usage_daily_summary).collect()
|
||||
}
|
||||
|
||||
async fn summarize_usage_daily_heatmap_from_daily_aggregates(
|
||||
&self,
|
||||
created_from_unix_secs: u64,
|
||||
created_until_unix_secs: u64,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||||
let rows = if let Some(user_id) = user_id {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
date("date", 'unixepoch') AS date,
|
||||
total_requests AS requests,
|
||||
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
|
||||
total_cost AS total_cost_usd,
|
||||
total_cost AS actual_total_cost_usd
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ?
|
||||
AND "date" >= ?
|
||||
AND "date" < ?
|
||||
AND total_requests > 0
|
||||
ORDER BY "date" ASC
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(created_from_unix_secs as i64)
|
||||
.bind(created_until_unix_secs as i64)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
date("date", 'unixepoch') AS date,
|
||||
total_requests AS requests,
|
||||
input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens AS total_tokens,
|
||||
total_cost AS total_cost_usd,
|
||||
actual_total_cost AS actual_total_cost_usd
|
||||
FROM stats_daily
|
||||
WHERE "date" >= ?
|
||||
AND "date" < ?
|
||||
AND total_requests > 0
|
||||
ORDER BY "date" ASC
|
||||
"#,
|
||||
)
|
||||
.bind(created_from_unix_secs as i64)
|
||||
.bind(created_until_unix_secs as i64)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
};
|
||||
|
||||
rows.iter().map(map_sqlite_usage_daily_summary).collect()
|
||||
}
|
||||
|
||||
async fn summarize_provider_performance_percentiles(
|
||||
&self,
|
||||
query: &UsageProviderPerformanceQuery,
|
||||
@@ -1463,38 +1560,266 @@ FROM "usage"
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(
|
||||
let unique_user_ids = user_ids
|
||||
.iter()
|
||||
.cloned()
|
||||
.collect::<std::collections::BTreeSet<_>>();
|
||||
let mut totals = BTreeMap::<String, StoredUsageUserTotals>::new();
|
||||
let mut aggregate_cutoffs = BTreeMap::<String, u64>::new();
|
||||
|
||||
let mut aggregate_builder = QueryBuilder::<Sqlite>::new(
|
||||
r#"
|
||||
SELECT
|
||||
user_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(MAX(COALESCE(total_tokens, 0), 0)), 0) AS total_tokens
|
||||
FROM "usage"
|
||||
COALESCE(SUM(total_requests), 0) AS request_count,
|
||||
COALESCE(
|
||||
SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens),
|
||||
0
|
||||
) AS total_tokens,
|
||||
MAX("date") AS latest_date
|
||||
FROM stats_user_daily
|
||||
WHERE user_id IN (
|
||||
"#,
|
||||
);
|
||||
let mut separated = builder.separated(", ");
|
||||
for user_id in user_ids {
|
||||
separated.push_bind(user_id.clone());
|
||||
{
|
||||
let mut separated = aggregate_builder.separated(", ");
|
||||
for user_id in &unique_user_ids {
|
||||
separated.push_bind(user_id.clone());
|
||||
}
|
||||
}
|
||||
separated.push_unseparated(
|
||||
r#")
|
||||
AND status NOT IN ('pending', 'streaming')
|
||||
AND provider_name NOT IN ('unknown', 'pending')
|
||||
GROUP BY user_id
|
||||
ORDER BY user_id ASC
|
||||
aggregate_builder.push(") GROUP BY user_id ORDER BY user_id ASC");
|
||||
|
||||
let aggregate_rows = aggregate_builder
|
||||
.build()
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?;
|
||||
for row in aggregate_rows {
|
||||
let user_id: String = row.try_get("user_id").map_sql_err()?;
|
||||
let latest_date = row.try_get::<i64, _>("latest_date").map_sql_err()?.max(0) as u64;
|
||||
aggregate_cutoffs.insert(user_id.clone(), latest_date.saturating_add(86_400));
|
||||
totals.insert(
|
||||
user_id.clone(),
|
||||
StoredUsageUserTotals {
|
||||
user_id,
|
||||
request_count: row_u64(&row, "request_count")?,
|
||||
total_tokens: row_u64(&row, "total_tokens")?,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(
|
||||
r#"
|
||||
SELECT
|
||||
"usage".user_id,
|
||||
COUNT(*) AS request_count,
|
||||
COALESCE(SUM(MAX(COALESCE("usage".total_tokens, 0), 0)), 0) AS total_tokens
|
||||
FROM "usage"
|
||||
JOIN (
|
||||
"#,
|
||||
);
|
||||
for (index, user_id) in unique_user_ids.iter().enumerate() {
|
||||
if index > 0 {
|
||||
builder.push(" UNION ALL ");
|
||||
}
|
||||
let cutoff = aggregate_cutoffs.get(user_id).copied().unwrap_or_default();
|
||||
builder
|
||||
.push("SELECT ")
|
||||
.push_bind(user_id.clone())
|
||||
.push(" AS user_id, ")
|
||||
.push_bind(to_i64(cutoff, "usage aggregate cutoff")?)
|
||||
.push(" AS cutoff_unix_secs");
|
||||
}
|
||||
builder.push(
|
||||
r#"
|
||||
) AS requested ON requested.user_id = "usage".user_id
|
||||
WHERE "usage".created_at_unix_ms >= requested.cutoff_unix_secs
|
||||
AND "usage".status NOT IN ('pending', 'streaming')
|
||||
AND "usage".provider_name NOT IN ('unknown', 'pending')
|
||||
GROUP BY "usage".user_id
|
||||
ORDER BY "usage".user_id ASC
|
||||
"#,
|
||||
);
|
||||
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
for row in rows {
|
||||
let user_id: String = row.try_get("user_id").map_sql_err()?;
|
||||
let entry = totals
|
||||
.entry(user_id.clone())
|
||||
.or_insert_with(|| StoredUsageUserTotals {
|
||||
user_id,
|
||||
request_count: 0,
|
||||
total_tokens: 0,
|
||||
});
|
||||
entry.request_count = entry
|
||||
.request_count
|
||||
.saturating_add(row_u64(&row, "request_count")?);
|
||||
entry.total_tokens = entry
|
||||
.total_tokens
|
||||
.saturating_add(row_u64(&row, "total_tokens")?);
|
||||
}
|
||||
Ok(totals.into_values().collect())
|
||||
}
|
||||
|
||||
async fn summarize_dashboard_usage_from_daily_aggregates(
|
||||
&self,
|
||||
query: &UsageDashboardSummaryQuery,
|
||||
) -> Result<Option<StoredUsageDashboardSummary>, DataLayerError> {
|
||||
let row = if let Some(user_id) = query.user_id.as_deref() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
|
||||
0.0 AS cache_creation_cost_usd,
|
||||
0.0 AS cache_read_cost_usd,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS actual_total_cost_usd,
|
||||
COALESCE(SUM(error_requests), 0) AS error_requests,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ?
|
||||
AND "date" >= ?
|
||||
AND "date" < ?
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(query.created_from_unix_secs as i64)
|
||||
.bind(query.created_until_unix_secs as i64)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
COALESCE(SUM(total_requests), 0) AS total_requests,
|
||||
COALESCE(SUM(input_tokens), 0) AS input_tokens,
|
||||
COALESCE(SUM(input_tokens), 0) AS effective_input_tokens,
|
||||
COALESCE(SUM(output_tokens), 0) AS output_tokens,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(cache_creation_tokens), 0) AS cache_creation_tokens,
|
||||
COALESCE(SUM(cache_read_tokens), 0) AS cache_read_tokens,
|
||||
COALESCE(SUM(input_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_input_context,
|
||||
0.0 AS cache_creation_cost_usd,
|
||||
0.0 AS cache_read_cost_usd,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(CAST(actual_total_cost AS REAL), 0)), 0) AS actual_total_cost_usd,
|
||||
COALESCE(SUM(error_requests), 0) AS error_requests,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_daily
|
||||
WHERE "date" >= ?
|
||||
AND "date" < ?
|
||||
"#,
|
||||
)
|
||||
.bind(query.created_from_unix_secs as i64)
|
||||
.bind(query.created_until_unix_secs as i64)
|
||||
.fetch_one(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
};
|
||||
|
||||
let total_requests = sqlite_aggregate_u64(&row, "total_requests")?;
|
||||
if total_requests == 0 {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
Ok(Some(StoredUsageDashboardSummary {
|
||||
total_requests,
|
||||
input_tokens: sqlite_aggregate_u64(&row, "input_tokens")?,
|
||||
effective_input_tokens: sqlite_aggregate_u64(&row, "effective_input_tokens")?,
|
||||
output_tokens: sqlite_aggregate_u64(&row, "output_tokens")?,
|
||||
total_tokens: sqlite_aggregate_u64(&row, "total_tokens")?,
|
||||
cache_creation_tokens: sqlite_aggregate_u64(&row, "cache_creation_tokens")?,
|
||||
cache_read_tokens: sqlite_aggregate_u64(&row, "cache_read_tokens")?,
|
||||
total_input_context: sqlite_aggregate_u64(&row, "total_input_context")?,
|
||||
cache_creation_cost_usd: sqlite_real(&row, "cache_creation_cost_usd")?,
|
||||
cache_read_cost_usd: sqlite_real(&row, "cache_read_cost_usd")?,
|
||||
total_cost_usd: sqlite_real(&row, "total_cost_usd")?,
|
||||
actual_total_cost_usd: sqlite_real(&row, "actual_total_cost_usd")?,
|
||||
error_requests: sqlite_aggregate_u64(&row, "error_requests")?,
|
||||
response_time_sum_ms: sqlite_real(&row, "response_time_sum_ms")?,
|
||||
response_time_samples: sqlite_aggregate_u64(&row, "response_time_samples")?,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn list_dashboard_daily_breakdown_from_daily_aggregates(
|
||||
&self,
|
||||
query: &UsageDashboardDailyBreakdownQuery,
|
||||
) -> Result<Vec<StoredUsageDashboardDailyBreakdownRow>, DataLayerError> {
|
||||
let rows = if let Some(user_id) = query.user_id.as_deref() {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
date("date", 'unixepoch') AS date,
|
||||
'aggregate' AS model,
|
||||
'aggregate' AS provider,
|
||||
COALESCE(SUM(total_requests), 0) AS requests,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_user_daily
|
||||
WHERE user_id = ?
|
||||
AND "date" >= ?
|
||||
AND "date" < ?
|
||||
AND total_requests > 0
|
||||
GROUP BY "date"
|
||||
ORDER BY "date" ASC
|
||||
"#,
|
||||
)
|
||||
.bind(user_id)
|
||||
.bind(query.created_from_unix_secs as i64)
|
||||
.bind(query.created_until_unix_secs as i64)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
} else {
|
||||
sqlx::query(
|
||||
r#"
|
||||
SELECT
|
||||
date("date", 'unixepoch') AS date,
|
||||
'aggregate' AS model,
|
||||
'aggregate' AS provider,
|
||||
COALESCE(SUM(total_requests), 0) AS requests,
|
||||
COALESCE(SUM(input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost AS REAL), 0)), 0) AS total_cost_usd,
|
||||
0.0 AS response_time_sum_ms,
|
||||
0 AS response_time_samples
|
||||
FROM stats_daily
|
||||
WHERE "date" >= ?
|
||||
AND "date" < ?
|
||||
AND total_requests > 0
|
||||
GROUP BY "date"
|
||||
ORDER BY "date" ASC
|
||||
"#,
|
||||
)
|
||||
.bind(query.created_from_unix_secs as i64)
|
||||
.bind(query.created_until_unix_secs as i64)
|
||||
.fetch_all(&self.pool)
|
||||
.await
|
||||
.map_sql_err()?
|
||||
};
|
||||
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
Ok(StoredUsageUserTotals {
|
||||
user_id: row.try_get::<String, _>("user_id").map_sql_err()?,
|
||||
request_count: row.try_get::<i64, _>("request_count").map_sql_err()?.max(0)
|
||||
as u64,
|
||||
total_tokens: row.try_get::<i64, _>("total_tokens").map_sql_err()?.max(0)
|
||||
as u64,
|
||||
Ok(StoredUsageDashboardDailyBreakdownRow {
|
||||
date: row.try_get("date").map_sql_err()?,
|
||||
model: row.try_get("model").map_sql_err()?,
|
||||
provider: row.try_get("provider").map_sql_err()?,
|
||||
requests: sqlite_aggregate_u64(row, "requests")?,
|
||||
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
|
||||
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
|
||||
response_time_sum_ms: sqlite_real(row, "response_time_sum_ms")?,
|
||||
response_time_samples: sqlite_aggregate_u64(row, "response_time_samples")?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
@@ -1885,6 +2210,13 @@ ORDER BY created_at_unix_ms ASC, id ASC
|
||||
return Ok(StoredUsageDashboardSummary::default());
|
||||
}
|
||||
|
||||
if let Some(summary) = self
|
||||
.summarize_dashboard_usage_from_daily_aggregates(query)
|
||||
.await?
|
||||
{
|
||||
return Ok(summary);
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(format!(
|
||||
r#"
|
||||
SELECT
|
||||
@@ -1959,6 +2291,13 @@ FROM "usage"
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
|
||||
let aggregate_rows = self
|
||||
.list_dashboard_daily_breakdown_from_daily_aggregates(query)
|
||||
.await?;
|
||||
if !aggregate_rows.is_empty() {
|
||||
return Ok(aggregate_rows);
|
||||
}
|
||||
|
||||
let date_expr = sqlite_usage_local_date_expr(query.tz_offset_minutes);
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(format!(
|
||||
r#"
|
||||
@@ -3227,53 +3566,54 @@ WHERE provider_id = ?
|
||||
&self,
|
||||
query: &UsageDailyHeatmapQuery,
|
||||
) -> Result<Vec<StoredUsageDailySummary>, DataLayerError> {
|
||||
let mut builder = QueryBuilder::<Sqlite>::new(format!(
|
||||
r#"
|
||||
SELECT
|
||||
date(created_at_unix_ms, 'unixepoch') AS date,
|
||||
COUNT(*) AS requests,
|
||||
COALESCE(SUM(
|
||||
MAX(COALESCE(input_tokens, 0), 0)
|
||||
+ MAX(COALESCE(output_tokens, 0), 0)
|
||||
+ {cache_creation_expr}
|
||||
+ MAX(COALESCE(cache_read_input_tokens, 0), 0)
|
||||
), 0) AS total_tokens,
|
||||
COALESCE(SUM(COALESCE(CAST(total_cost_usd AS REAL), 0)), 0) AS total_cost_usd,
|
||||
COALESCE(SUM(COALESCE(CAST(actual_total_cost_usd AS REAL), 0)), 0)
|
||||
AS actual_total_cost_usd
|
||||
FROM "usage"
|
||||
"#,
|
||||
cache_creation_expr = SQLITE_USAGE_CACHE_CREATION_TOKENS_EXPR
|
||||
));
|
||||
let mut has_where = false;
|
||||
push_sqlite_usage_where(&mut builder, &mut has_where);
|
||||
builder
|
||||
.push("created_at_unix_ms >= ")
|
||||
.push_bind(query.created_from_unix_secs as i64);
|
||||
push_sqlite_usage_finalized_filter(&mut builder, &mut has_where);
|
||||
push_sqlite_usage_optional_text_filter(
|
||||
&mut builder,
|
||||
&mut has_where,
|
||||
"user_id",
|
||||
query.user_id.as_deref(),
|
||||
);
|
||||
builder.push(" GROUP BY date ORDER BY date ASC");
|
||||
let created_until_unix_secs = usage_current_unix_secs().saturating_add(1);
|
||||
let user_id = query.user_id.as_deref();
|
||||
let mut summaries = BTreeMap::<String, StoredUsageDailySummary>::new();
|
||||
|
||||
let rows = builder.build().fetch_all(&self.pool).await.map_sql_err()?;
|
||||
rows.iter()
|
||||
.map(|row| {
|
||||
Ok(StoredUsageDailySummary {
|
||||
date: row.try_get("date").map_sql_err()?,
|
||||
requests: sqlite_aggregate_u64(row, "requests")?,
|
||||
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
|
||||
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
|
||||
actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?,
|
||||
})
|
||||
})
|
||||
.collect()
|
||||
for item in self
|
||||
.summarize_usage_daily_heatmap_from_daily_aggregates(
|
||||
query.created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
user_id,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
summaries.insert(item.date.clone(), item);
|
||||
}
|
||||
for item in self
|
||||
.summarize_usage_daily_heatmap_raw_from_range(
|
||||
query.created_from_unix_secs,
|
||||
created_until_unix_secs,
|
||||
user_id,
|
||||
)
|
||||
.await?
|
||||
{
|
||||
summaries.entry(item.date.clone()).or_insert(item);
|
||||
}
|
||||
|
||||
Ok(summaries.into_values().collect())
|
||||
}
|
||||
}
|
||||
|
||||
fn map_sqlite_usage_daily_summary(
|
||||
row: &SqliteRow,
|
||||
) -> Result<StoredUsageDailySummary, DataLayerError> {
|
||||
Ok(StoredUsageDailySummary {
|
||||
date: row.try_get("date").map_sql_err()?,
|
||||
requests: sqlite_aggregate_u64(row, "requests")?,
|
||||
total_tokens: sqlite_aggregate_u64(row, "total_tokens")?,
|
||||
total_cost_usd: sqlite_real(row, "total_cost_usd")?,
|
||||
actual_total_cost_usd: sqlite_real(row, "actual_total_cost_usd")?,
|
||||
})
|
||||
}
|
||||
|
||||
fn usage_current_unix_secs() -> u64 {
|
||||
SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map(|duration| duration.as_secs())
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
impl SqliteUsageWriteRepository {
|
||||
pub fn new(pool: SqlitePool) -> Self {
|
||||
Self { pool }
|
||||
@@ -4012,7 +4352,8 @@ mod tests {
|
||||
use super::{SqliteUsageReadRepository, SqliteUsageWriteRepository};
|
||||
use crate::lifecycle::migrate::run_sqlite_migrations;
|
||||
use crate::repository::usage::{
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageDashboardSummaryQuery, UsageReadRepository,
|
||||
UpsertUsageRecord, UsageAuditListQuery, UsageDailyHeatmapQuery,
|
||||
UsageDashboardDailyBreakdownQuery, UsageDashboardSummaryQuery, UsageReadRepository,
|
||||
UsageWriteRepository,
|
||||
};
|
||||
|
||||
@@ -4350,6 +4691,169 @@ INSERT INTO request_candidates (
|
||||
assert_eq!(summary.total_tokens, 10);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_usage_daily_heatmap_reads_imported_daily_aggregates() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("sqlite pool should connect");
|
||||
run_sqlite_migrations(&pool)
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO stats_daily (
|
||||
id, "date", total_requests, success_requests, error_requests,
|
||||
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
||||
total_cost, actual_total_cost, is_complete, created_at, updated_at
|
||||
) VALUES (
|
||||
'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1
|
||||
);
|
||||
INSERT INTO stats_user_daily (
|
||||
id, user_id, username, "date", total_requests, success_requests, error_requests,
|
||||
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
||||
total_cost, created_at, updated_at
|
||||
) VALUES (
|
||||
'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1
|
||||
);
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("daily aggregates should seed");
|
||||
|
||||
let reader = SqliteUsageReadRepository::new(pool);
|
||||
let admin = reader
|
||||
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
|
||||
created_from_unix_secs: 0,
|
||||
user_id: None,
|
||||
admin_mode: true,
|
||||
})
|
||||
.await
|
||||
.expect("admin heatmap should load");
|
||||
assert_eq!(admin.len(), 1);
|
||||
assert_eq!(admin[0].date, "1970-01-02");
|
||||
assert_eq!(admin[0].requests, 9);
|
||||
assert_eq!(admin[0].total_tokens, 37);
|
||||
assert_eq!(admin[0].actual_total_cost_usd, 1.0);
|
||||
|
||||
let user = reader
|
||||
.summarize_usage_daily_heatmap(&UsageDailyHeatmapQuery {
|
||||
created_from_unix_secs: 0,
|
||||
user_id: Some("user-1".to_string()),
|
||||
admin_mode: false,
|
||||
})
|
||||
.await
|
||||
.expect("user heatmap should load");
|
||||
assert_eq!(user.len(), 1);
|
||||
assert_eq!(user[0].date, "1970-01-02");
|
||||
assert_eq!(user[0].requests, 5);
|
||||
assert_eq!(user[0].total_tokens, 18);
|
||||
assert_eq!(user[0].actual_total_cost_usd, 0.75);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_usage_totals_by_user_ids_reads_imported_user_daily_aggregates() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("sqlite pool should connect");
|
||||
run_sqlite_migrations(&pool)
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO stats_user_daily (
|
||||
id, user_id, username, "date", total_requests, success_requests, error_requests,
|
||||
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
||||
total_cost, created_at, updated_at
|
||||
) VALUES (
|
||||
'user-daily-1', 'user-1', 'user one', 86400, 5, 5, 0, 7, 8, 2, 1, 0.75, 1, 1
|
||||
);
|
||||
INSERT INTO "usage" (
|
||||
request_id, id, user_id, api_key_id, provider_name, model, total_tokens,
|
||||
status, billing_status, created_at_unix_ms, updated_at_unix_secs
|
||||
) VALUES
|
||||
('raw-before-cutoff', 'usage-1', 'user-1', 'api-key-1', 'Provider One', 'model-1', 99,
|
||||
'completed', 'settled', 90000, 90000),
|
||||
('raw-after-cutoff', 'usage-2', 'user-1', 'api-key-1', 'Provider One', 'model-1', 7,
|
||||
'completed', 'settled', 172800, 172800);
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("usage totals fixtures should seed");
|
||||
|
||||
let reader = SqliteUsageReadRepository::new(pool);
|
||||
let totals = reader
|
||||
.summarize_usage_totals_by_user_ids(&["user-1".to_string()])
|
||||
.await
|
||||
.expect("user totals should load");
|
||||
|
||||
assert_eq!(totals.len(), 1);
|
||||
assert_eq!(totals[0].user_id, "user-1");
|
||||
assert_eq!(totals[0].request_count, 6);
|
||||
assert_eq!(totals[0].total_tokens, 25);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn sqlite_dashboard_daily_stats_reads_imported_daily_aggregates() {
|
||||
let pool = sqlx::sqlite::SqlitePoolOptions::new()
|
||||
.max_connections(1)
|
||||
.connect("sqlite::memory:")
|
||||
.await
|
||||
.expect("sqlite pool should connect");
|
||||
run_sqlite_migrations(&pool)
|
||||
.await
|
||||
.expect("sqlite migrations should run");
|
||||
|
||||
sqlx::query(
|
||||
r#"
|
||||
INSERT INTO stats_daily (
|
||||
id, "date", total_requests, success_requests, error_requests,
|
||||
input_tokens, output_tokens, cache_creation_tokens, cache_read_tokens,
|
||||
total_cost, actual_total_cost, is_complete, created_at, updated_at
|
||||
) VALUES (
|
||||
'daily-1', 86400, 9, 8, 1, 10, 20, 3, 4, 1.25, 1.0, 1, 1, 1
|
||||
);
|
||||
"#,
|
||||
)
|
||||
.execute(&pool)
|
||||
.await
|
||||
.expect("daily aggregates should seed");
|
||||
|
||||
let reader = SqliteUsageReadRepository::new(pool);
|
||||
let summary = reader
|
||||
.summarize_dashboard_usage(&UsageDashboardSummaryQuery {
|
||||
created_from_unix_secs: 0,
|
||||
created_until_unix_secs: 172800,
|
||||
user_id: None,
|
||||
})
|
||||
.await
|
||||
.expect("dashboard summary should load");
|
||||
assert_eq!(summary.total_requests, 9);
|
||||
assert_eq!(summary.total_tokens, 37);
|
||||
|
||||
let rows = reader
|
||||
.list_dashboard_daily_breakdown(&UsageDashboardDailyBreakdownQuery {
|
||||
created_from_unix_secs: 0,
|
||||
created_until_unix_secs: 172800,
|
||||
tz_offset_minutes: 480,
|
||||
user_id: None,
|
||||
})
|
||||
.await
|
||||
.expect("dashboard daily breakdown should load");
|
||||
assert_eq!(rows.len(), 1);
|
||||
assert_eq!(rows[0].date, "1970-01-02");
|
||||
assert_eq!(rows[0].model, "aggregate");
|
||||
assert_eq!(rows[0].requests, 9);
|
||||
assert_eq!(rows[0].total_tokens, 37);
|
||||
}
|
||||
|
||||
async fn seed_stats_targets(pool: &sqlx::SqlitePool) {
|
||||
sqlx::query(
|
||||
r#"
|
||||
|
||||
Reference in New Issue
Block a user