cc_token_usage/analysis/
trend.rs

1use std::collections::HashMap;
2
3use chrono::{Datelike, NaiveDate, Utc};
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        Utc::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.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        let all_turns = session.turns.iter().chain(session.agent_turns.iter());
38        for turn in all_turns {
39            let date = turn.timestamp.date_naive();
40            if date < cutoff {
41                continue;
42            }
43
44            let label = make_label(date, group_by_month);
45            let acc = accumulators.entry(label).or_insert_with(|| Accumulator {
46                first_date: date,
47                turn_count: 0,
48                tokens: AggregatedTokens::default(),
49                cost: 0.0,
50                models: HashMap::new(),
51                cost_by_category: CostByCategory::default(),
52            });
53
54            // Keep earliest date for sorting
55            if date < acc.first_date {
56                acc.first_date = date;
57            }
58
59            acc.turn_count += 1;
60            acc.tokens.add_usage(&turn.usage);
61
62            let pricing_cost = calc.calculate_turn_cost(&turn.model, &turn.usage);
63            acc.cost += pricing_cost.total;
64
65            // Accumulate cost by category
66            acc.cost_by_category.input_cost += pricing_cost.input_cost;
67            acc.cost_by_category.output_cost += pricing_cost.output_cost;
68            acc.cost_by_category.cache_write_5m_cost += pricing_cost.cache_write_5m_cost;
69            acc.cost_by_category.cache_write_1h_cost += pricing_cost.cache_write_1h_cost;
70            acc.cost_by_category.cache_read_cost += pricing_cost.cache_read_cost;
71
72            *acc.models.entry(turn.model.clone()).or_insert(0) +=
73                turn.usage.output_tokens.unwrap_or(0);
74        }
75    }
76
77    let mut entries: Vec<TrendEntry> = accumulators
78        .into_iter()
79        .map(|(label, acc)| TrendEntry {
80            label: label.clone(),
81            date: acc.first_date,
82            session_count: session_labels.get(&label).copied().unwrap_or(0),
83            turn_count: acc.turn_count,
84            tokens: acc.tokens,
85            cost: acc.cost,
86            models: acc.models,
87            cost_by_category: acc.cost_by_category,
88        })
89        .collect();
90
91    entries.sort_by_key(|e| e.date);
92
93    TrendResult {
94        entries,
95        group_label: if group_by_month { "Month" } else { "Day" }.to_string(),
96    }
97}
98
99fn make_label(date: NaiveDate, group_by_month: bool) -> String {
100    if group_by_month {
101        format!("{}-{:02}", date.year(), date.month())
102    } else {
103        date.format("%Y-%m-%d").to_string()
104    }
105}
106
107struct Accumulator {
108    first_date: NaiveDate,
109    turn_count: usize,
110    tokens: AggregatedTokens,
111    cost: f64,
112    models: HashMap<String, u64>,
113    cost_by_category: CostByCategory,
114}
cc_token_usage/analysis/trend.rs

cc_token_usage/analysis/
trend.rs
