lc_core/token_counter/
tiktoken.rs

1//! Tiktoken counter (OpenAI tokenizer)
2
3use tiktoken_rs::CoreBPE;
4
5use lc_schema::Message;
6
7use super::counter::TokenCounter;
8use super::TokenCounterError;
9
10/// Tiktoken counter (uses cl100k_base, for GPT-3.5 / 4 / 4o)
11pub struct TiktokenCounter {
12    encoder: CoreBPE,
13}
14
15impl TiktokenCounter {
16    /// Creates a Tiktoken counter (loads the cl100k_base BPE encoder).
17    pub fn new() -> Result<Self, TokenCounterError> {
18        let encoder = tiktoken_rs::cl100k_base()
19            .map_err(|e| TokenCounterError::EncoderLoad(e.to_string()))?;
20        Ok(Self { encoder })
21    }
22}
23
24// H35: Removed Default impl that could panic.
25// Use `TiktokenCounter::new()` instead of `TiktokenCounter::default()`.
26
27impl TokenCounter for TiktokenCounter {
28    fn count_tokens(&self, text: &str) -> u32 {
29        self.encoder.encode_with_special_tokens(text).len() as u32
30    }
31
32    fn count_messages(&self, messages: &[Message]) -> u32 {
33        let mut total = 0u32;
34        for msg in messages {
35            total += 4; // OpenAI message-format overhead
36            total += self.count_tokens(&msg.content);
37            if let Some(name) = &msg.name {
38                total += self.count_tokens(name);
39            }
40            // images count roughly as 1000 tokens each
41            for _ in &msg.images {
42                total += 1000;
43            }
44        }
45        total += 2; // conversation boundary marker
46        total
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53
54    #[test]
55    fn test_count_tokens_nonempty() {
56        let counter = TiktokenCounter::new().unwrap();
57        let n = counter.count_tokens("Hello, world!");
58        assert!(n > 0);
59    }
60
61    #[test]
62    fn test_count_tokens_empty() {
63        let counter = TiktokenCounter::new().unwrap();
64        assert_eq!(counter.count_tokens(""), 0);
65    }
66
67    #[test]
68    fn test_count_messages_includes_overhead() {
69        let counter = TiktokenCounter::new().unwrap();
70        let msgs = vec![Message::system("You are helpful."), Message::human("Hi")];
71        let n = counter.count_messages(&msgs);
72        // at least 2*4 overhead + 2 boundary + per-message tokens
73        assert!(n >= 10);
74    }
75
76    #[test]
77    fn test_count_messages_with_image() {
78        let counter = TiktokenCounter::new().unwrap();
79        let msg = Message::human_with_image("看图", "https://example.com/x.png");
80        let n = counter.count_messages(&[msg]);
81        assert!(n >= 1000); // image = 1000 tokens
82    }
83
84    #[test]
85    fn test_longer_text_more_tokens() {
86        let counter = TiktokenCounter::new().unwrap();
87        let short = counter.count_tokens("hi");
88        let long = counter.count_tokens("This is a much longer sentence with many words.");
89        assert!(long > short);
90    }
91}
lc_core/token_counter/tiktoken.rs

lc_core/token_counter/
tiktoken.rs
