lc/
usage_stats.rs

1use anyhow::Result;
2use chrono::{DateTime, Utc, Duration, Datelike};
3use colored::Colorize;
4use std::collections::HashMap;
5use crate::database::{Database, ChatEntry};
6
7#[derive(Debug, Clone)]
8pub struct UsageStats {
9    pub total_tokens: u64,
10    pub total_requests: u64,
11    pub input_tokens: u64,
12    pub output_tokens: u64,
13    pub model_usage: Vec<(String, u64, u64)>, // (model, requests, tokens)
14    pub daily_usage: Vec<(String, u64, u64)>, // (date, requests, tokens)
15    pub weekly_usage: Vec<(String, u64, u64)>, // (week, requests, tokens)
16    pub monthly_usage: Vec<(String, u64, u64)>, // (month, requests, tokens)
17    pub yearly_usage: Vec<(String, u64, u64)>, // (year, requests, tokens)
18    pub date_range: Option<(DateTime<Utc>, DateTime<Utc>)>,
19}
20
21#[derive(Debug, Clone)]
22pub enum TimeFrame {
23    Daily,
24    Weekly,
25    Monthly,
26    Yearly,
27}
28
29pub struct UsageAnalyzer {
30    db: Database,
31}
32
33impl UsageAnalyzer {
34    pub fn new() -> Result<Self> {
35        Ok(Self {
36            db: Database::new()?,
37        })
38    }
39
40    pub fn get_usage_stats(&self, days_back: Option<u32>) -> Result<UsageStats> {
41        let entries = if let Some(days) = days_back {
42            let cutoff_date = Utc::now() - Duration::days(days as i64);
43            self.get_entries_since(cutoff_date)?
44        } else {
45            self.db.get_all_logs()?
46        };
47
48        if entries.is_empty() {
49            return Ok(UsageStats {
50                total_tokens: 0,
51                total_requests: 0,
52                input_tokens: 0,
53                output_tokens: 0,
54                model_usage: Vec::new(),
55                daily_usage: Vec::new(),
56                weekly_usage: Vec::new(),
57                monthly_usage: Vec::new(),
58                yearly_usage: Vec::new(),
59                date_range: None,
60            });
61        }
62
63        let mut total_input_tokens = 0u64;
64        let mut total_output_tokens = 0u64;
65        let mut model_stats: HashMap<String, (u64, u64)> = HashMap::new(); // (requests, tokens)
66        let mut daily_stats: HashMap<String, (u64, u64)> = HashMap::new();
67        let mut weekly_stats: HashMap<String, (u64, u64)> = HashMap::new();
68        let mut monthly_stats: HashMap<String, (u64, u64)> = HashMap::new();
69        let mut yearly_stats: HashMap<String, (u64, u64)> = HashMap::new();
70
71        let mut earliest_date = entries[0].timestamp;
72        let mut latest_date = entries[0].timestamp;
73
74        for entry in &entries {
75            // Update date range
76            if entry.timestamp < earliest_date {
77                earliest_date = entry.timestamp;
78            }
79            if entry.timestamp > latest_date {
80                latest_date = entry.timestamp;
81            }
82
83            // Calculate tokens
84            let input_tokens = entry.input_tokens.unwrap_or(0) as u64;
85            let output_tokens = entry.output_tokens.unwrap_or(0) as u64;
86            let total_entry_tokens = input_tokens + output_tokens;
87
88            total_input_tokens += input_tokens;
89            total_output_tokens += output_tokens;
90
91            // Model usage
92            let model_entry = model_stats.entry(entry.model.clone()).or_insert((0, 0));
93            model_entry.0 += 1; // requests
94            model_entry.1 += total_entry_tokens; // tokens
95
96            // Time-based usage
97            let date = entry.timestamp.date_naive();
98            let daily_key = date.format("%Y-%m-%d").to_string();
99            let daily_entry = daily_stats.entry(daily_key).or_insert((0, 0));
100            daily_entry.0 += 1;
101            daily_entry.1 += total_entry_tokens;
102
103            // Weekly usage (ISO week)
104            let year = entry.timestamp.year();
105            let week = entry.timestamp.iso_week().week();
106            let weekly_key = format!("{}-W{:02}", year, week);
107            let weekly_entry = weekly_stats.entry(weekly_key).or_insert((0, 0));
108            weekly_entry.0 += 1;
109            weekly_entry.1 += total_entry_tokens;
110
111            // Monthly usage
112            let monthly_key = date.format("%Y-%m").to_string();
113            let monthly_entry = monthly_stats.entry(monthly_key).or_insert((0, 0));
114            monthly_entry.0 += 1;
115            monthly_entry.1 += total_entry_tokens;
116
117            // Yearly usage
118            let yearly_key = year.to_string();
119            let yearly_entry = yearly_stats.entry(yearly_key).or_insert((0, 0));
120            yearly_entry.0 += 1;
121            yearly_entry.1 += total_entry_tokens;
122        }
123
124        // Convert to sorted vectors
125        let mut model_usage: Vec<(String, u64, u64)> = model_stats
126            .into_iter()
127            .map(|(model, (requests, tokens))| (model, requests, tokens))
128            .collect();
129        model_usage.sort_by(|a, b| b.2.cmp(&a.2)); // Sort by tokens descending
130
131        let mut daily_usage: Vec<(String, u64, u64)> = daily_stats
132            .into_iter()
133            .map(|(date, (requests, tokens))| (date, requests, tokens))
134            .collect();
135        daily_usage.sort_by(|a, b| a.0.cmp(&b.0)); // Sort by date ascending
136
137        let mut weekly_usage: Vec<(String, u64, u64)> = weekly_stats
138            .into_iter()
139            .map(|(week, (requests, tokens))| (week, requests, tokens))
140            .collect();
141        weekly_usage.sort_by(|a, b| a.0.cmp(&b.0));
142
143        let mut monthly_usage: Vec<(String, u64, u64)> = monthly_stats
144            .into_iter()
145            .map(|(month, (requests, tokens))| (month, requests, tokens))
146            .collect();
147        monthly_usage.sort_by(|a, b| a.0.cmp(&b.0));
148
149        let mut yearly_usage: Vec<(String, u64, u64)> = yearly_stats
150            .into_iter()
151            .map(|(year, (requests, tokens))| (year, requests, tokens))
152            .collect();
153        yearly_usage.sort_by(|a, b| a.0.cmp(&b.0));
154
155        Ok(UsageStats {
156            total_tokens: total_input_tokens + total_output_tokens,
157            total_requests: entries.len() as u64,
158            input_tokens: total_input_tokens,
159            output_tokens: total_output_tokens,
160            model_usage,
161            daily_usage,
162            weekly_usage,
163            monthly_usage,
164            yearly_usage,
165            date_range: Some((earliest_date, latest_date)),
166        })
167    }
168
169    fn get_entries_since(&self, cutoff_date: DateTime<Utc>) -> Result<Vec<ChatEntry>> {
170        // This would need a custom query in the database
171        // For now, we'll filter after getting all entries
172        let all_entries = self.db.get_all_logs()?;
173        Ok(all_entries
174            .into_iter()
175            .filter(|entry| entry.timestamp >= cutoff_date)
176            .collect())
177    }
178}
179
180
181pub struct BarChart;
182
183impl BarChart {
184    pub fn render_horizontal(
185        title: &str,
186        data: &[(String, u64, u64)],
187        value_type: &str, // "tokens" or "requests"
188        max_width: usize,
189        max_items: usize,
190    ) {
191        if data.is_empty() {
192            println!("{} No data available", "ℹ️".blue());
193            return;
194        }
195
196        println!("\n{}", title.bold().blue());
197        
198        let display_data: Vec<_> = data.iter().take(max_items).collect();
199        let max_value = display_data
200            .iter()
201            .map(|(_, requests, tokens)| {
202                if value_type == "tokens" { *tokens } else { *requests }
203            })
204            .max()
205            .unwrap_or(1);
206
207        let max_label_width = display_data
208            .iter()
209            .map(|(label, _, _)| label.len())
210            .max()
211            .unwrap_or(10);
212
213        for (label, requests, tokens) in display_data {
214            let value = if value_type == "tokens" { *tokens } else { *requests };
215            let bar_width = if max_value > 0 {
216                ((value as f64 / max_value as f64) * max_width as f64) as usize
217            } else {
218                0
219            };
220
221            let bar = "█".repeat(bar_width);
222            let formatted_value = if value_type == "tokens" {
223                Self::format_tokens(*tokens)
224            } else {
225                format!("{}", requests)
226            };
227
228            println!(
229                "  {:width$} │{:bar_width$} {} ({})",
230                label.bold(),
231                bar.green(),
232                formatted_value.yellow(),
233                if value_type == "tokens" {
234                    format!("{} req", requests)
235                } else {
236                    Self::format_tokens(*tokens)
237                },
238                width = max_label_width,
239                bar_width = max_width
240            );
241        }
242    }
243
244    pub fn render_time_series(
245        title: &str,
246        data: &[(String, u64, u64)],
247        value_type: &str,
248        max_width: usize,
249        max_items: usize,
250    ) {
251        if data.is_empty() {
252            println!("{} No data available", "ℹ️".blue());
253            return;
254        }
255
256        println!("\n{}", title.bold().blue());
257        
258        let display_data: Vec<_> = data.iter().rev().take(max_items).rev().collect();
259        let max_value = display_data
260            .iter()
261            .map(|(_, requests, tokens)| {
262                if value_type == "tokens" { *tokens } else { *requests }
263            })
264            .max()
265            .unwrap_or(1);
266
267        let max_label_width = display_data
268            .iter()
269            .map(|(label, _, _)| label.len())
270            .max()
271            .unwrap_or(10);
272
273        for (label, requests, tokens) in display_data {
274            let value = if value_type == "tokens" { *tokens } else { *requests };
275            let bar_width = if max_value > 0 {
276                ((value as f64 / max_value as f64) * max_width as f64) as usize
277            } else {
278                0
279            };
280
281            let bar = "▓".repeat(bar_width);
282            let formatted_value = if value_type == "tokens" {
283                Self::format_tokens(*tokens)
284            } else {
285                format!("{}", requests)
286            };
287
288            println!(
289                "  {:width$} │{:bar_width$} {} ({})",
290                label.bold(),
291                bar.cyan(),
292                formatted_value.yellow(),
293                if value_type == "tokens" {
294                    format!("{} req", requests)
295                } else {
296                    Self::format_tokens(*tokens)
297                },
298                width = max_label_width,
299                bar_width = max_width
300            );
301        }
302    }
303
304    fn format_tokens(tokens: u64) -> String {
305        if tokens >= 1_000_000 {
306            format!("{:.1}M", tokens as f64 / 1_000_000.0)
307        } else if tokens >= 1_000 {
308            format!("{:.1}k", tokens as f64 / 1_000.0)
309        } else {
310            format!("{}", tokens)
311        }
312    }
313}
314
315pub fn display_usage_overview(stats: &UsageStats) {
316    println!("\n{}", "📊 Usage Overview".bold().blue());
317    println!();
318
319    // Basic stats
320    println!("{} {}", "Total Requests:".bold(), stats.total_requests.to_string().green());
321    println!("{} {}", "Total Tokens:".bold(), BarChart::format_tokens(stats.total_tokens).green());
322    println!("{} {}", "Input Tokens:".bold(), BarChart::format_tokens(stats.input_tokens).cyan());
323    println!("{} {}", "Output Tokens:".bold(), BarChart::format_tokens(stats.output_tokens).yellow());
324
325    if let Some((earliest, latest)) = stats.date_range {
326        let duration = latest.signed_duration_since(earliest);
327        println!(
328            "{} {} to {} ({} days)",
329            "Date Range:".bold(),
330            earliest.format("%Y-%m-%d").to_string().dimmed(),
331            latest.format("%Y-%m-%d").to_string().dimmed(),
332            duration.num_days().max(1)
333        );
334    }
335
336    // Average tokens per request
337    if stats.total_requests > 0 {
338        let avg_tokens = stats.total_tokens / stats.total_requests;
339        let avg_input = stats.input_tokens / stats.total_requests;
340        let avg_output = stats.output_tokens / stats.total_requests;
341        println!();
342        println!("{}", "📈 Averages per Request".bold().blue());
343        println!("{} {}", "Total Tokens:".bold(), BarChart::format_tokens(avg_tokens).green());
344        println!("{} {}", "Input Tokens:".bold(), BarChart::format_tokens(avg_input).cyan());
345        println!("{} {}", "Output Tokens:".bold(), BarChart::format_tokens(avg_output).yellow());
346    }
347}
348
349#[cfg(test)]
350mod tests {
351    use super::*;
352
353    #[test]
354    fn test_usage_analyzer_creation() {
355        // Test that creating a UsageAnalyzer works on all platforms
356        // Note: This may fail if database setup fails, which is acceptable for this test
357        let _result = UsageAnalyzer::new();
358        // We don't assert success here because database setup might fail in test environment
359        // but we're primarily testing the compilation path
360    }
361
362    #[test]
363    fn test_usage_stats_struct_creation() {
364        // Test that we can create empty UsageStats (should work on all platforms)
365        let stats = UsageStats {
366            total_tokens: 0,
367            total_requests: 0,
368            input_tokens: 0,
369            output_tokens: 0,
370            model_usage: Vec::new(),
371            daily_usage: Vec::new(),
372            weekly_usage: Vec::new(),
373            monthly_usage: Vec::new(),
374            yearly_usage: Vec::new(),
375            date_range: None,
376        };
377        
378        assert_eq!(stats.total_tokens, 0);
379        assert_eq!(stats.total_requests, 0);
380        assert!(stats.model_usage.is_empty());
381    }
382
383    #[test]
384    fn test_bar_chart_format_tokens() {
385        // Test token formatting function (should work on all platforms)
386        assert_eq!(BarChart::format_tokens(500), "500");
387        assert_eq!(BarChart::format_tokens(1500), "1.5k");
388        assert_eq!(BarChart::format_tokens(1_500_000), "1.5M");
389    }
390}