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)>, pub daily_usage: Vec<(String, u64, u64)>, pub weekly_usage: Vec<(String, u64, u64)>, pub monthly_usage: Vec<(String, u64, u64)>, pub yearly_usage: Vec<(String, u64, u64)>, 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(); 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 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 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 let model_entry = model_stats.entry(entry.model.clone()).or_insert((0, 0));
93 model_entry.0 += 1; model_entry.1 += total_entry_tokens; 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 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 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 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 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)); 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)); 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 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, 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 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 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 let _result = UsageAnalyzer::new();
358 }
361
362 #[test]
363 fn test_usage_stats_struct_creation() {
364 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 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}