cognis_llm/
usage.rs

1//! Running-total usage tracker for budget enforcement.
2
3use std::sync::atomic::{AtomicU64, Ordering};
4
5use crate::chat::Usage;
6
7/// Aggregates `Usage` across many `ChatResponse`s. Cheap to share via `Arc`.
8#[derive(Debug, Default)]
9pub struct UsageTracker {
10    prompt_tokens: AtomicU64,
11    completion_tokens: AtomicU64,
12    total_tokens: AtomicU64,
13    requests: AtomicU64,
14}
15
16impl UsageTracker {
17    /// Empty tracker.
18    pub fn new() -> Self {
19        Self::default()
20    }
21
22    /// Add one `Usage` to the running totals.
23    pub fn record(&self, u: &Usage) {
24        self.prompt_tokens
25            .fetch_add(u.prompt_tokens as u64, Ordering::Relaxed);
26        self.completion_tokens
27            .fetch_add(u.completion_tokens as u64, Ordering::Relaxed);
28        self.total_tokens
29            .fetch_add(u.total_tokens as u64, Ordering::Relaxed);
30        self.requests.fetch_add(1, Ordering::Relaxed);
31    }
32
33    /// Snapshot the totals as a `Usage` (truncates `u64 → u32`).
34    pub fn snapshot(&self) -> Usage {
35        Usage {
36            prompt_tokens: self.prompt_tokens.load(Ordering::Relaxed) as u32,
37            completion_tokens: self.completion_tokens.load(Ordering::Relaxed) as u32,
38            total_tokens: self.total_tokens.load(Ordering::Relaxed) as u32,
39        }
40    }
41
42    /// Number of recorded calls.
43    pub fn requests(&self) -> u64 {
44        self.requests.load(Ordering::Relaxed)
45    }
46
47    /// Reset all counters to zero.
48    pub fn reset(&self) {
49        self.prompt_tokens.store(0, Ordering::Relaxed);
50        self.completion_tokens.store(0, Ordering::Relaxed);
51        self.total_tokens.store(0, Ordering::Relaxed);
52        self.requests.store(0, Ordering::Relaxed);
53    }
54}
55
56#[cfg(test)]
57mod tests {
58    use super::*;
59
60    #[test]
61    fn record_and_snapshot() {
62        let t = UsageTracker::new();
63        t.record(&Usage {
64            prompt_tokens: 10,
65            completion_tokens: 5,
66            total_tokens: 15,
67        });
68        t.record(&Usage {
69            prompt_tokens: 20,
70            completion_tokens: 8,
71            total_tokens: 28,
72        });
73        let snap = t.snapshot();
74        assert_eq!(snap.prompt_tokens, 30);
75        assert_eq!(snap.completion_tokens, 13);
76        assert_eq!(snap.total_tokens, 43);
77        assert_eq!(t.requests(), 2);
78    }
79
80    #[test]
81    fn reset_zeros_everything() {
82        let t = UsageTracker::new();
83        t.record(&Usage {
84            prompt_tokens: 5,
85            completion_tokens: 5,
86            total_tokens: 10,
87        });
88        t.reset();
89        assert_eq!(t.requests(), 0);
90        assert_eq!(t.snapshot().total_tokens, 0);
91    }
92}
cognis_llm/usage.rs

cognis_llm/
usage.rs
