language_barrier_core/
token.rs

1/// A simple token counter for tracking token usage in conversations
2///
3/// This is a very basic implementation that provides a minimal token counting
4/// utility. In a real-world scenario, you would want to use a proper tokenizer
5/// like `tiktoken-rs` for accurate token counting specific to the model being used.
6///
7/// # Examples
8///
9/// ```
10/// use language_barrier_core::TokenCounter;
11///
12/// let mut counter = TokenCounter::default();
13/// counter.observe("Hello, world!");
14/// assert_eq!(counter.total(), 2); // "Hello," and "world!" as two tokens
15///
16/// counter.subtract("Hello,");
17/// assert_eq!(counter.total(), 1); // "world!" as one token
18///
19/// assert!(counter.under_budget(10)); // 1 < 10
20/// ```
21#[derive(Default, Clone, Debug)]
22pub struct TokenCounter {
23    total: usize,
24}
25
26impl TokenCounter {
27    /// Creates a new `TokenCounter` with zero tokens
28    ///
29    /// # Examples
30    ///
31    /// ```
32    /// use language_barrier_core::TokenCounter;
33    ///
34    /// let counter = TokenCounter::new();
35    /// assert_eq!(counter.total(), 0);
36    /// ```
37    #[must_use]
38    pub fn new() -> Self {
39        Self { total: 0 }
40    }
41
42    /// Counts the number of tokens in a string (naive implementation)
43    ///
44    /// This is a very simple whitespace-based tokenization. In a real-world
45    /// implementation, you would want to use a proper tokenizer like `tiktoken-rs`.
46    ///
47    /// # Examples
48    ///
49    /// ```
50    /// use language_barrier_core::TokenCounter;
51    ///
52    /// assert_eq!(TokenCounter::count_tokens("Hello, world!"), 2);
53    /// assert_eq!(TokenCounter::count_tokens("one two three four"), 4);
54    /// ```
55    #[must_use]
56    pub fn count_tokens(text: &str) -> usize {
57        text.split_whitespace().count()
58    }
59
60    /// Adds the token count of the given text to the total
61    ///
62    /// # Examples
63    ///
64    /// ```
65    /// use language_barrier_core::TokenCounter;
66    ///
67    /// let mut counter = TokenCounter::default();
68    /// counter.observe("Hello, world!");
69    /// assert_eq!(counter.total(), 2);
70    /// ```
71    pub fn observe(&mut self, text: &str) {
72        self.total += Self::count_tokens(text);
73    }
74
75    /// Subtracts the token count of the given text from the total
76    ///
77    /// Will not go below zero (saturates at zero).
78    ///
79    /// # Examples
80    ///
81    /// ```
82    /// use language_barrier_core::TokenCounter;
83    ///
84    /// let mut counter = TokenCounter::default();
85    /// counter.observe("Hello, world!");
86    /// counter.subtract("Hello,");
87    /// assert_eq!(counter.total(), 1);
88    ///
89    /// // Won't go below zero
90    /// counter.subtract("world! and more");
91    /// assert_eq!(counter.total(), 0);
92    /// ```
93    pub fn subtract(&mut self, text: &str) {
94        self.total = self.total.saturating_sub(Self::count_tokens(text));
95    }
96
97    /// Returns the current total token count
98    ///
99    /// # Examples
100    ///
101    /// ```
102    /// use language_barrier_core::TokenCounter;
103    ///
104    /// let mut counter = TokenCounter::default();
105    /// counter.observe("Hello, world!");
106    /// assert_eq!(counter.total(), 2);
107    /// ```
108    #[must_use]
109    pub fn total(&self) -> usize {
110        self.total
111    }
112
113    /// Checks if the total token count is under the specified budget
114    ///
115    /// # Examples
116    ///
117    /// ```
118    /// use language_barrier_core::TokenCounter;
119    ///
120    /// let mut counter = TokenCounter::default();
121    /// counter.observe("Hello, world!");
122    /// assert!(counter.under_budget(5));
123    /// assert!(!counter.under_budget(1));
124    /// ```
125    #[must_use]
126    pub fn under_budget(&self, max: usize) -> bool {
127        self.total <= max
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134
135    #[test]
136    fn test_token_counter_new() {
137        let counter = TokenCounter::new();
138        assert_eq!(counter.total(), 0);
139    }
140
141    #[test]
142    fn test_token_counting() {
143        assert_eq!(TokenCounter::count_tokens("Hello world"), 2);
144        assert_eq!(TokenCounter::count_tokens(""), 0);
145        assert_eq!(TokenCounter::count_tokens("   "), 0);
146        assert_eq!(TokenCounter::count_tokens("one two three"), 3);
147    }
148
149    #[test]
150    fn test_observe_and_total() {
151        let mut counter = TokenCounter::default();
152        counter.observe("Hello world");
153        assert_eq!(counter.total(), 2);
154
155        counter.observe("another message");
156        assert_eq!(counter.total(), 4); // 2 + 2 tokens
157    }
158
159    #[test]
160    fn test_subtract() {
161        let mut counter = TokenCounter::default();
162        counter.observe("Hello world another message");
163        assert_eq!(counter.total(), 4);
164
165        counter.subtract("Hello world");
166        assert_eq!(counter.total(), 2);
167
168        // Test saturation (won't go below zero)
169        counter.subtract("way more tokens than we have now");
170        assert_eq!(counter.total(), 0);
171    }
172
173    #[test]
174    fn test_under_budget() {
175        let mut counter = TokenCounter::default();
176        counter.observe("Hello world");
177
178        assert!(counter.under_budget(2));
179        assert!(counter.under_budget(3));
180        assert!(!counter.under_budget(1));
181    }
182}
language_barrier_core/token.rs

language_barrier_core/
token.rs
