cascade_agent/memory/
tokenizer.rs

1use llm_cascade::Message;
2use tokenizers::Tokenizer;
3
4use crate::error::Result;
5
6/// Token counter backed by a HuggingFace `tokenizers` tokenizer.
7///
8/// Provides accurate token counting for LLM context window management.
9/// Falls back to a character-based heuristic if tokenization fails.
10pub struct TokenCounter {
11    tokenizer: Option<Tokenizer>,
12    model_identifier: String,
13}
14
15impl TokenCounter {
16    /// Load a HuggingFace tokenizer by model identifier or local file path.
17    ///
18    /// Accepts:
19    /// - HuggingFace model names like `"Xenova/gpt-4o"` (fetched via the Hub)
20    /// - Local file paths to a tokenizer.json file
21    ///
22    /// If loading fails (e.g., no network for Hub model, missing file), the counter
23    /// still works using a character-based estimate fallback.
24    pub fn new(model_identifier: &str) -> Result<Self> {
25        let tokenizer = match Tokenizer::from_pretrained(model_identifier, None) {
26            Ok(t) => {
27                tracing::info!("Loaded tokenizer: {}", model_identifier);
28                Some(t)
29            }
30            Err(e) => {
31                tracing::warn!(
32                    "Failed to load tokenizer '{}': {}. Will use char-based estimate fallback.",
33                    model_identifier,
34                    e
35                );
36                None
37            }
38        };
39
40        Ok(Self {
41            tokenizer,
42            model_identifier: model_identifier.to_string(),
43        })
44    }
45
46    /// Count the number of tokens in a text string.
47    ///
48    /// If the tokenizer is available, encodes the text and returns the token count.
49    /// Falls back to `text.len() / 4` as a rough estimate.
50    pub fn count_text(&self, text: &str) -> usize {
51        if let Some(ref tok) = self.tokenizer {
52            match tok.encode(text, false) {
53                Ok(encoding) => encoding.len(),
54                Err(e) => {
55                    tracing::debug!("Tokenization failed, using estimate: {}", e);
56                    estimate_tokens(text)
57                }
58            }
59        } else {
60            estimate_tokens(text)
61        }
62    }
63
64    /// Count tokens for a list of chat messages, including role prefix overhead.
65    ///
66    /// Adds estimated tokens for role prefixes and message separators:
67    /// - `system: ` ≈ 2 tokens
68    /// - `user: ` ≈ 2 tokens
69    /// - `assistant: ` ≈ 3 tokens
70    /// - `tool: ` ≈ 2 tokens
71    /// - Plus 1 separator token per message
72    pub fn count_messages(&self, messages: &[Message]) -> usize {
73        let mut total = 0;
74        for msg in messages {
75            total += self.count_text(&msg.content);
76            total += role_overhead(&msg.role);
77            total += 1; // separator token between messages
78        }
79        // Add 3 tokens for the overall conversation framing (bos, eos, etc.)
80        total += 3;
81        total
82    }
83
84    /// Returns the model identifier this counter was created with.
85    pub fn model_identifier(&self) -> &str {
86        &self.model_identifier
87    }
88}
89
90/// Estimated token overhead for a message role prefix.
91fn role_overhead(role: &llm_cascade::MessageRole) -> usize {
92    match role {
93        llm_cascade::MessageRole::System => 2,    // "system: "
94        llm_cascade::MessageRole::User => 2,      // "user: "
95        llm_cascade::MessageRole::Assistant => 3, // "assistant: "
96        llm_cascade::MessageRole::Tool => 2,      // "tool: "
97    }
98}
99
100/// Rough character-based token estimate (≈4 chars per token for English text).
101fn estimate_tokens(text: &str) -> usize {
102    (text.len() / 4).max(1)
103}
104
105#[cfg(test)]
106mod tests {
107    use super::*;
108
109    #[test]
110    fn test_estimate_tokens() {
111        assert_eq!(estimate_tokens("hello"), 1); // 5 chars / 4 = 1
112        assert_eq!(estimate_tokens("hello world this is a test"), 6); // 27 chars / 4 = 6
113        assert!(estimate_tokens("") >= 1); // empty → max(1) = 1
114    }
115
116    #[test]
117    fn test_count_messages_empty() {
118        // Create a counter that will use fallback (no real tokenizer needed)
119        let counter = TokenCounter {
120            tokenizer: None,
121            model_identifier: "test".into(),
122        };
123        assert_eq!(counter.count_messages(&[]), 3); // just the framing tokens
124    }
125
126    #[test]
127    fn test_count_messages_with_roles() {
128        let counter = TokenCounter {
129            tokenizer: None,
130            model_identifier: "test".into(),
131        };
132        let msgs = vec![
133            Message::system("You are helpful."),
134            Message::user("Hello!"),
135            Message::assistant("Hi there!"),
136        ];
137        let count = counter.count_messages(&msgs);
138        // "You are helpful." = 16/4=4 + 2 (system) + 1 (sep) = 7
139        // "Hello!" = 6/4=1 + 2 (user) + 1 (sep) = 4
140        // "Hi there!" = 9/4=2 + 3 (assistant) + 1 (sep) = 6
141        // + 3 framing = 20
142        assert_eq!(count, 20);
143    }
144}
cascade_agent/memory/tokenizer.rs

cascade_agent/memory/
tokenizer.rs
