Skip to main content

cc_token_usage/analysis/
trend.rs

1use std::collections::HashMap;
2
3use chrono::{Datelike, Local, NaiveDate};
4
5use crate::data::models::SessionData;
6use crate::pricing::calculator::PricingCalculator;
7
8use super::{AggregatedTokens, CostByCategory, TrendEntry, TrendResult};
9
10pub fn analyze_trend(
11    sessions: &[SessionData],
12    calc: &PricingCalculator,
13    days: u32,
14    group_by_month: bool,
15) -> TrendResult {
16    // days=0 means all history
17    let cutoff = if days == 0 {
18        chrono::NaiveDate::from_ymd_opt(2000, 1, 1).unwrap()
19    } else {
20        Local::now().date_naive() - chrono::Duration::days(days as i64)
21    };
22
23    let mut accumulators: HashMap<String, Accumulator> = HashMap::new();
24    let mut session_labels: HashMap<String, usize> = HashMap::new();
25
26    for session in sessions {
27        // Count session by its first_timestamp
28        if let Some(first_ts) = session.first_timestamp {
29            let date = first_ts.with_timezone(&Local).date_naive();
30            if date >= cutoff {
31                let label = make_label(date, group_by_month);
32                *session_labels.entry(label).or_insert(0) += 1;
33            }
34        }
35
36        // Process all turns
37        for turn in session.all_responses() {
38            let date = turn.timestamp.with_timezone(&Local).date_naive();
39            if date < cutoff {
40                continue;
41            }
42
43            let label = make_label(date, group_by_month);
44            let acc = accumulators.entry(label).or_insert_with(|| Accumulator {
45                first_date: date,
46                turn_count: 0,
47                tokens: AggregatedTokens::default(),
48                cost: 0.0,
49                models: HashMap::new(),
50                cost_by_category: CostByCategory::default(),
51            });
52
53            // Keep earliest date for sorting
54            if date < acc.first_date {
55                acc.first_date = date;
56            }
57
58            acc.turn_count += 1;
59            acc.tokens.add_usage(&turn.usage);
60
61            let pricing_cost = calc.calculate_turn_cost(&turn.model, &turn.usage);
62            acc.cost += pricing_cost.total;
63
64            // Accumulate cost by category
65            acc.cost_by_category.input_cost += pricing_cost.input_cost;
66            acc.cost_by_category.output_cost += pricing_cost.output_cost;
67            acc.cost_by_category.cache_write_5m_cost += pricing_cost.cache_write_5m_cost;
68            acc.cost_by_category.cache_write_1h_cost += pricing_cost.cache_write_1h_cost;
69            acc.cost_by_category.cache_read_cost += pricing_cost.cache_read_cost;
70
71            *acc.models.entry(turn.model.clone()).or_insert(0) +=
72                turn.usage.output_tokens.unwrap_or(0);
73        }
74    }
75
76    let mut entries: Vec<TrendEntry> = accumulators
77        .into_iter()
78        .map(|(label, acc)| TrendEntry {
79            label: label.clone(),
80            date: acc.first_date,
81            session_count: session_labels.get(&label).copied().unwrap_or(0),
82            turn_count: acc.turn_count,
83            tokens: acc.tokens,
84            cost: acc.cost,
85            models: acc.models,
86            cost_by_category: acc.cost_by_category,
87        })
88        .collect();
89
90    entries.sort_by_key(|e| e.date);
91
92    TrendResult {
93        entries,
94        group_label: if group_by_month { "Month" } else { "Day" }.to_string(),
95    }
96}
97
98fn make_label(date: NaiveDate, group_by_month: bool) -> String {
99    if group_by_month {
100        format!("{}-{:02}", date.year(), date.month())
101    } else {
102        date.format("%Y-%m-%d").to_string()
103    }
104}
105
106struct Accumulator {
107    first_date: NaiveDate,
108    turn_count: usize,
109    tokens: AggregatedTokens,
110    cost: f64,
111    models: HashMap<String, u64>,
112    cost_by_category: CostByCategory,
113}