//! Efficient database aggregates for the desktop overview. use std::collections::BTreeMap; use chrono::Utc; use sqlx::Row; use crate::{ model::{Overview, OverviewMetrics, TokenUsageBucket, TokenUsageGranularity}, Result, }; use super::Store; const OVERVIEW_DAYS: u64 = 365; const MAX_RANGE_BUCKETS: i64 = 60; const MINUTE_MS: i64 = 60_000; const HOUR_MS: i64 = 60 * MINUTE_MS; const DAY_MS: i64 = 24 * HOUR_MS; impl Store { pub async fn overview( &self, start_ms: Option, end_ms: Option, model_hashes: Option<&str>, provider_ids: Option<&str>, ) -> Result { let call_row = sqlx::query( "SELECT COUNT(*) AS llm_calls, COALESCE(SUM(status = 'completed'), 0) AS successful_calls, COALESCE(SUM(status != 'completed'), 0) AS failed_calls FROM llm_calls WHERE status != 'running' AND (? IS NULL OR created_at_ms >= ?) AND (? IS NULL OR created_at_ms < ?) AND (? IS NULL OR model_hash IN (SELECT value FROM json_each(?))) AND (? IS NULL OR model_hash IN ( SELECT model_hash FROM provider_models WHERE provider_id IN (SELECT value FROM json_each(?)) ))", ) .bind(start_ms) .bind(start_ms) .bind(end_ms) .bind(end_ms) .bind(model_hashes) .bind(model_hashes) .bind(provider_ids) .bind(provider_ids) .fetch_one(&self.pool) .await?; let token_row = sqlx::query(&format!( "SELECT COALESCE(SUM({fresh_input}), 0) AS input_tokens, COALESCE(SUM(COALESCE(cache_read_tokens, 0)), 0) AS cache_read_tokens, COALESCE(SUM(COALESCE(cache_write_tokens, 0)), 0) AS cache_write_tokens, COALESCE(SUM(COALESCE(output_tokens, 0)), 0) AS output_tokens FROM llm_calls WHERE (? IS NULL OR created_at_ms >= ?) AND (? IS NULL OR created_at_ms < ?) AND (? IS NULL OR model_hash IN (SELECT value FROM json_each(?))) AND (? IS NULL OR model_hash IN ( SELECT model_hash FROM provider_models WHERE provider_id IN (SELECT value FROM json_each(?)) ))", fresh_input = fresh_input_sql(), )) .bind(start_ms) .bind(start_ms) .bind(end_ms) .bind(end_ms) .bind(model_hashes) .bind(model_hashes) .bind(provider_ids) .bind(provider_ids) .fetch_one(&self.pool) .await?; let input_tokens = non_negative(token_row.try_get("input_tokens")?); let cache_read_tokens = non_negative(token_row.try_get("cache_read_tokens")?); let cache_write_tokens = non_negative(token_row.try_get("cache_write_tokens")?); let output_tokens = non_negative(token_row.try_get("output_tokens")?); let prompt_tokens = saturating_sum(&[input_tokens, cache_read_tokens, cache_write_tokens]); let metrics = OverviewMetrics { llm_calls: call_row.try_get("llm_calls")?, successful_calls: call_row.try_get("successful_calls")?, failed_calls: call_row.try_get("failed_calls")?, token_usage: prompt_tokens.saturating_add(output_tokens), prompt_tokens, input_tokens, cache_read_tokens, cache_write_tokens, output_tokens, }; let (token_usage_granularity, bucket_ms, series_start_ms, bucket_count) = token_usage_buckets(start_ms, end_ms); let rows = sqlx::query(&format!( "SELECT (created_at_ms / {bucket_ms}) * {bucket_ms} AS bucket_start_ms, COALESCE(SUM({fresh_input}), 0) AS input_tokens, COALESCE(SUM(COALESCE(cache_read_tokens, 0)), 0) AS cache_read_tokens, COALESCE(SUM(COALESCE(cache_write_tokens, 0)), 0) AS cache_write_tokens, COALESCE(SUM(COALESCE(output_tokens, 0)), 0) AS output_tokens FROM llm_calls WHERE created_at_ms >= ? AND (? IS NULL OR created_at_ms < ?) AND (? IS NULL OR model_hash IN (SELECT value FROM json_each(?))) AND (? IS NULL OR model_hash IN ( SELECT model_hash FROM provider_models WHERE provider_id IN (SELECT value FROM json_each(?)) )) GROUP BY bucket_start_ms ORDER BY bucket_start_ms", fresh_input = fresh_input_sql(), )) .bind(start_ms.unwrap_or(series_start_ms).max(series_start_ms)) .bind(end_ms) .bind(end_ms) .bind(model_hashes) .bind(model_hashes) .bind(provider_ids) .bind(provider_ids) .fetch_all(&self.pool) .await?; let mut recorded = rows .into_iter() .map(|row| { let bucket_start_ms: i64 = row.try_get("bucket_start_ms")?; Ok(( bucket_start_ms, TokenUsageBucket { bucket_start_ms, input_tokens: non_negative(row.try_get("input_tokens")?), cache_read_tokens: non_negative(row.try_get("cache_read_tokens")?), cache_write_tokens: non_negative(row.try_get("cache_write_tokens")?), output_tokens: non_negative(row.try_get("output_tokens")?), }, )) }) .collect::>>()?; let token_usage_series = (0..bucket_count) .map(|offset| series_start_ms.saturating_add(offset.saturating_mul(bucket_ms))) .map(|bucket_start_ms| { recorded .remove(&bucket_start_ms) .unwrap_or(TokenUsageBucket { bucket_start_ms, ..TokenUsageBucket::default() }) }) .collect(); Ok(Overview { metrics, token_usage_granularity, token_usage_series, }) } } fn token_usage_buckets( start_ms: Option, end_ms: Option, ) -> (TokenUsageGranularity, i64, i64, i64) { if let (Some(start_ms), Some(end_ms)) = (start_ms, end_ms) { let duration_ms = end_ms.saturating_sub(start_ms).max(1); let (granularity, bucket_ms) = if duration_ms <= HOUR_MS { (TokenUsageGranularity::Minute, MINUTE_MS) } else if duration_ms <= MAX_RANGE_BUCKETS * HOUR_MS { (TokenUsageGranularity::Hour, HOUR_MS) } else { (TokenUsageGranularity::Day, DAY_MS) }; let last_bucket_ms = end_ms.saturating_sub(1).div_euclid(bucket_ms) * bucket_ms; let first_bucket_ms = start_ms.div_euclid(bucket_ms) * bucket_ms; let bucket_count = ((last_bucket_ms - first_bucket_ms).div_euclid(bucket_ms) + 1) .clamp(1, MAX_RANGE_BUCKETS); let series_start_ms = last_bucket_ms.saturating_sub((bucket_count - 1).saturating_mul(bucket_ms)); return (granularity, bucket_ms, series_start_ms, bucket_count); } let today_start_ms = Utc::now() .date_naive() .and_hms_opt(0, 0, 0) .map(|value| value.and_utc().timestamp_millis()) .unwrap_or(0); let series_start_ms = today_start_ms.saturating_sub( i64::try_from(OVERVIEW_DAYS - 1) .unwrap_or(0) .saturating_mul(DAY_MS), ); ( TokenUsageGranularity::Day, DAY_MS, series_start_ms, i64::try_from(OVERVIEW_DAYS).unwrap_or(0), ) } fn fresh_input_sql() -> &'static str { "CASE WHEN request_type = 'anthropic' THEN MAX(0, COALESCE(input_tokens, 0)) ELSE MAX(0, COALESCE(input_tokens, 0) - COALESCE(cache_read_tokens, 0) - COALESCE(cache_write_tokens, 0)) END" } fn non_negative(value: i64) -> i64 { value.max(0) } fn saturating_sum(values: &[i64]) -> i64 { values .iter() .fold(0_i64, |total, value| total.saturating_add(*value)) } #[cfg(test)] mod tests { use chrono::{Duration, Utc}; use super::*; #[test] fn one_hour_range_uses_sixty_minute_buckets() { let start_ms = 1_800_000_000_000; let (granularity, bucket_ms, series_start_ms, bucket_count) = token_usage_buckets(Some(start_ms), Some(start_ms + HOUR_MS)); assert_eq!(granularity, TokenUsageGranularity::Minute); assert_eq!(bucket_ms, MINUTE_MS); assert_eq!(series_start_ms, start_ms); assert_eq!(bucket_count, 60); } #[tokio::test] async fn overview_aggregates_llm_calls_and_normalizes_provider_usage() { let directory = tempfile::tempdir().unwrap(); let store = Store::connect(&format!( "sqlite://{}", directory.path().join("overview.db").display() )) .await .unwrap(); let now = Utc::now().timestamp_millis(); insert_call(&store, "openai", "openai-responses", now, [100, 20, 80, 0]).await; insert_call(&store, "anthropic", "anthropic", now, [30, 10, 50, 5]).await; sqlx::query("UPDATE llm_calls SET status = 'error' WHERE call_id = 'anthropic'") .execute(&store.pool) .await .unwrap(); insert_call( &store, "old", "openai-chat", (Utc::now() - Duration::days(400)).timestamp_millis(), [10, 5, 0, 0], ) .await; let overview = store.overview(None, None, None, None).await.unwrap(); assert_eq!(overview.metrics.llm_calls, 3); assert_eq!(overview.metrics.successful_calls, 2); assert_eq!(overview.metrics.failed_calls, 1); assert_eq!(overview.metrics.input_tokens, 60); assert_eq!(overview.metrics.cache_read_tokens, 130); assert_eq!(overview.metrics.cache_write_tokens, 5); assert_eq!(overview.metrics.output_tokens, 35); assert_eq!(overview.metrics.prompt_tokens, 195); assert_eq!(overview.metrics.token_usage, 230); assert_eq!(overview.token_usage_granularity, TokenUsageGranularity::Day); assert_eq!(overview.token_usage_series.len(), OVERVIEW_DAYS as usize); let today = overview.token_usage_series.last().unwrap(); assert_eq!(today.input_tokens, 50); assert_eq!(today.cache_read_tokens, 130); assert_eq!(today.cache_write_tokens, 5); assert_eq!(today.output_tokens, 30); assert_eq!(today.total_tokens(), 215); } #[tokio::test] async fn overview_filters_metrics_and_usage_by_time_range() { let directory = tempfile::tempdir().unwrap(); let store = Store::connect(&format!( "sqlite://{}", directory.path().join("ranged-overview.db").display() )) .await .unwrap(); let now = Utc::now().timestamp_millis(); insert_call(&store, "inside", "anthropic", now, [20, 5, 10, 2]).await; insert_call( &store, "outside", "anthropic", now - Duration::hours(2).num_milliseconds(), [100, 50, 40, 20], ) .await; let overview = store .overview(Some(now - 1_000), Some(now + 1_000), None, None) .await .unwrap(); assert_eq!(overview.metrics.llm_calls, 1); assert_eq!(overview.metrics.input_tokens, 20); assert_eq!(overview.metrics.cache_read_tokens, 10); assert_eq!(overview.metrics.cache_write_tokens, 2); assert_eq!(overview.metrics.output_tokens, 5); assert_eq!( overview.token_usage_granularity, TokenUsageGranularity::Minute ); assert_eq!(overview.token_usage_series.len(), 1); assert_eq!(overview.token_usage_series[0].total_tokens(), 37); let filtered = store .overview( Some(now - 1_000), Some(now + 1_000), Some(r#"["missing-model"]"#), None, ) .await .unwrap(); assert_eq!(filtered.metrics.llm_calls, 0); assert_eq!(filtered.metrics.token_usage, 0); assert_eq!(filtered.token_usage_series[0].total_tokens(), 0); } async fn insert_call( store: &Store, call_id: &str, request_type: &str, created_at_ms: i64, usage: [i64; 4], ) { let [input_tokens, output_tokens, cache_read_tokens, cache_write_tokens] = usage; sqlx::query( "INSERT INTO llm_calls( call_id, run_id, conversation_id, provider_call_index, provider_type, provider_url, request_type, request_url, model_id, display_name, status, created_at_ms, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, message_count, tool_count, detailed) VALUES (?, 'completed', 'conversation', 0, ?, '', ?, '', 'model', 'Model', 'completed', ?, ?, ?, ?, ?, 0, 0, 0)", ) .bind(call_id) .bind(request_type) .bind(request_type) .bind(created_at_ms) .bind(input_tokens) .bind(output_tokens) .bind(cache_read_tokens) .bind(cache_write_tokens) .execute(&store.pool) .await .unwrap(); } }