use std::collections::HashMap;
use chrono::{Datelike, Duration, NaiveDate, Weekday};
use serde::{Deserialize, Serialize};
use crate::metrics::types::{DailyMetrics, TokenUsage};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct PeriodMetrics {
pub label: String,
pub period_start: NaiveDate,
pub period_end: NaiveDate,
pub total_sessions: u32,
pub total_rounds: u32,
pub total_token_usage: TokenUsage,
pub total_tool_calls: u32,
pub prompt_cached_tool_outputs: u64,
pub model_breakdown: HashMap<String, TokenUsage>,
pub tool_breakdown: HashMap<String, u32>,
}
pub fn aggregate_weekly(daily_metrics: &[DailyMetrics]) -> Vec<PeriodMetrics> {
aggregate_by_period(daily_metrics, start_of_week)
}
pub fn aggregate_monthly(daily_metrics: &[DailyMetrics]) -> Vec<PeriodMetrics> {
aggregate_by_period(daily_metrics, |date| date.with_day(1).unwrap_or(date))
}
fn aggregate_by_period<F>(
daily_metrics: &[DailyMetrics],
period_start_resolver: F,
) -> Vec<PeriodMetrics>
where
F: Fn(NaiveDate) -> NaiveDate,
{
let mut buckets: HashMap<NaiveDate, PeriodMetrics> = HashMap::new();
for day in daily_metrics {
let period_start = period_start_resolver(day.date);
let entry = buckets
.entry(period_start)
.or_insert_with(|| PeriodMetrics {
label: period_start.to_string(),
period_start,
period_end: period_start,
total_sessions: 0,
total_rounds: 0,
total_token_usage: TokenUsage::default(),
total_tool_calls: 0,
prompt_cached_tool_outputs: 0,
model_breakdown: HashMap::new(),
tool_breakdown: HashMap::new(),
});
if day.date > entry.period_end {
entry.period_end = day.date;
}
entry.total_sessions += day.total_sessions;
entry.total_rounds += day.total_rounds;
entry.total_token_usage.add_assign(day.total_token_usage);
entry.total_tool_calls += day.total_tool_calls;
entry.prompt_cached_tool_outputs += day.prompt_cached_tool_outputs;
for (model, usage) in &day.model_breakdown {
let model_entry = entry.model_breakdown.entry(model.clone()).or_default();
model_entry.add_assign(*usage);
}
for (tool, count) in &day.tool_breakdown {
*entry.tool_breakdown.entry(tool.clone()).or_insert(0) += count;
}
}
let mut periods: Vec<PeriodMetrics> = buckets.into_values().collect();
periods.sort_by_key(|period| period.period_start);
for period in &mut periods {
period.label = format_period_label(period.period_start, period.period_end);
}
periods
}
fn start_of_week(date: NaiveDate) -> NaiveDate {
let weekday = match date.weekday() {
Weekday::Mon => 0,
Weekday::Tue => 1,
Weekday::Wed => 2,
Weekday::Thu => 3,
Weekday::Fri => 4,
Weekday::Sat => 5,
Weekday::Sun => 6,
};
date - Duration::days(weekday)
}
fn format_period_label(start: NaiveDate, end: NaiveDate) -> String {
if start == end {
start.to_string()
} else {
format!("{}..{}", start, end)
}
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use chrono::NaiveDate;
use super::{aggregate_monthly, aggregate_weekly};
use crate::metrics::types::{DailyMetrics, TokenUsage};
#[test]
fn aggregate_weekly_combines_days_into_a_single_week_bucket() {
let input = vec![
DailyMetrics {
date: NaiveDate::from_ymd_opt(2026, 2, 9).expect("valid date"),
total_sessions: 2,
total_rounds: 3,
total_token_usage: TokenUsage {
prompt_tokens: 10,
completion_tokens: 20,
total_tokens: 30,
},
total_tool_calls: 4,
prompt_cached_tool_outputs: 3,
model_breakdown: HashMap::new(),
tool_breakdown: HashMap::new(),
},
DailyMetrics {
date: NaiveDate::from_ymd_opt(2026, 2, 10).expect("valid date"),
total_sessions: 1,
total_rounds: 2,
total_token_usage: TokenUsage {
prompt_tokens: 5,
completion_tokens: 5,
total_tokens: 10,
},
total_tool_calls: 1,
prompt_cached_tool_outputs: 2,
model_breakdown: HashMap::new(),
tool_breakdown: HashMap::new(),
},
];
let result = aggregate_weekly(&input);
assert_eq!(result.len(), 1);
assert_eq!(result[0].total_sessions, 3);
assert_eq!(result[0].total_rounds, 5);
assert_eq!(result[0].total_token_usage.total_tokens, 40);
assert_eq!(result[0].total_tool_calls, 5);
assert_eq!(result[0].prompt_cached_tool_outputs, 5);
}
#[test]
fn aggregate_monthly_groups_metrics_by_month_start() {
let input = vec![
DailyMetrics {
date: NaiveDate::from_ymd_opt(2026, 1, 31).expect("valid date"),
total_sessions: 1,
total_rounds: 1,
total_token_usage: TokenUsage {
prompt_tokens: 1,
completion_tokens: 1,
total_tokens: 2,
},
total_tool_calls: 1,
prompt_cached_tool_outputs: 0,
model_breakdown: HashMap::new(),
tool_breakdown: HashMap::new(),
},
DailyMetrics {
date: NaiveDate::from_ymd_opt(2026, 2, 1).expect("valid date"),
total_sessions: 2,
total_rounds: 2,
total_token_usage: TokenUsage {
prompt_tokens: 2,
completion_tokens: 3,
total_tokens: 5,
},
total_tool_calls: 3,
prompt_cached_tool_outputs: 1,
model_breakdown: HashMap::new(),
tool_breakdown: HashMap::new(),
},
];
let result = aggregate_monthly(&input);
assert_eq!(result.len(), 2);
assert_eq!(
result[0].period_start,
NaiveDate::from_ymd_opt(2026, 1, 1).expect("valid date")
);
assert_eq!(
result[1].period_start,
NaiveDate::from_ymd_opt(2026, 2, 1).expect("valid date")
);
}
}