lc_core/token_counter/
counter.rs

1//! Token counter trait and usage statistics
2
3use lc_schema::Message;
4
5/// Token counter trait
6pub trait TokenCounter: Send + Sync {
7    /// Counts tokens in a text
8    fn count_tokens(&self, text: &str) -> u32;
9    /// Counts tokens in a message list
10    fn count_messages(&self, messages: &[Message]) -> u32;
11}
12
13/// Char-ratio estimate counter (zero-dependency fast path)
14///
15/// `count_tokens = len / ratio`, no tiktoken BPE model needed, usable offline / in tests.
16/// Default `ratio = 4` aligns with the old `len/4` rough estimate; `count_messages` reuses
17/// `TiktokenCounter`'s overhead structure (4 per message + name + 1000/image, 2 at the end),
18/// keeping the budget semantics consistent with the BPE accounting.
19#[derive(Debug, Clone)]
20pub struct CharRatioCounter {
21    ratio: u32,
22}
23
24impl CharRatioCounter {
25    /// Creates a char-ratio counter.
26    ///
27    /// `ratio` is the number of chars per token, at least 1.
28    pub fn new(ratio: u32) -> Self {
29        Self {
30            ratio: ratio.max(1),
31        }
32    }
33}
34
35impl TokenCounter for CharRatioCounter {
36    fn count_tokens(&self, text: &str) -> u32 {
37        text.len() as u32 / self.ratio
38    }
39
40    fn count_messages(&self, messages: &[Message]) -> u32 {
41        let mut total = 0u32;
42        for msg in messages {
43            total += 4; // OpenAI message-format overhead
44            total += self.count_tokens(&msg.content);
45            if let Some(name) = &msg.name {
46                total += self.count_tokens(name);
47            }
48            // images count roughly as 1000 tokens each
49            for _ in &msg.images {
50                total += 1000;
51            }
52        }
53        total += 2; // conversation boundary marker
54        total
55    }
56}
57
58#[cfg(test)]
59mod char_ratio_tests {
60    use super::*;
61
62    #[test]
63    fn test_char_ratio_counts_bytes() {
64        let counter = CharRatioCounter::new(4);
65        assert_eq!(counter.count_tokens(""), 0);
66        assert_eq!(counter.count_tokens("Hello"), 1); // 5 / 4 = 1
67        assert_eq!(counter.count_tokens("Hello World"), 2); // 11 / 4 = 2
68    }
69
70    #[test]
71    fn test_char_ratio_ratio_at_least_one() {
72        let counter = CharRatioCounter::new(0);
73        assert!(counter.count_tokens("x") >= 1);
74    }
75
76    #[test]
77    fn test_char_ratio_count_messages_matches_structure() {
78        let counter = CharRatioCounter::new(4);
79        let msgs = vec![Message::system("You are helpful."), Message::human("Hi")];
80        let n = counter.count_messages(&msgs);
81        // at least 2*4 overhead + 2 boundary
82        assert!(n >= 10);
83    }
84}
85
86/// Token usage statistics (internal type for the counter module)
87///
88/// Note: `language_models::TokenUsage` is the usage returned by the LLM API (fields are `usize`);
89/// this `TrackerTokenUsage` is locally tracked cumulative usage, **also `usize` fields**, so the two
90/// convert without precision loss (Q6: unified base type, no more `usize` vs `u32` naming clash).
91/// No `TokenUsage` alias is used here, to avoid confusion with `language_models::TokenUsage`.
92#[derive(Debug, Clone, Default, PartialEq)]
93pub struct TrackerTokenUsage {
94    /// Cumulative prompt token count
95    pub prompt_tokens: usize,
96    /// Cumulative completion token count
97    pub completion_tokens: usize,
98    /// Cumulative total token count (prompt + completion)
99    pub total_tokens: usize,
100}
101
102impl TrackerTokenUsage {
103    /// Creates empty usage statistics (all zero).
104    pub fn new() -> Self {
105        Self::default()
106    }
107
108    /// Accumulates usage
109    pub fn add(&mut self, prompt: usize, completion: usize) {
110        self.prompt_tokens += prompt;
111        self.completion_tokens += completion;
112        self.total_tokens = self.prompt_tokens + self.completion_tokens;
113    }
114
115    /// Resets usage statistics to all zero.
116    pub fn reset(&mut self) {
117        *self = Self::default();
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    #[test]
126    fn test_usage_add() {
127        let mut u = TrackerTokenUsage::new();
128        u.add(10, 20);
129        assert_eq!(u.prompt_tokens, 10);
130        assert_eq!(u.completion_tokens, 20);
131        assert_eq!(u.total_tokens, 30);
132    }
133
134    #[test]
135    fn test_usage_accumulate() {
136        let mut u = TrackerTokenUsage::new();
137        u.add(10, 20);
138        u.add(5, 5);
139        assert_eq!(u.prompt_tokens, 15);
140        assert_eq!(u.total_tokens, 40);
141    }
142
143    #[test]
144    fn test_usage_reset() {
145        let mut u = TrackerTokenUsage::new();
146        u.add(10, 20);
147        u.reset();
148        assert_eq!(u, TrackerTokenUsage::new());
149    }
150}
lc_core/token_counter/counter.rs

lc_core/token_counter/
counter.rs
