Skip to main content

token_count/output/
simple.rs

1//! Simple formatter - outputs just the token count
2
3use crate::output::OutputFormatter;
4use crate::tokenizers::TokenizationResult;
5
6/// Simple formatter that outputs only the token count
7pub struct SimpleFormatter;
8
9impl OutputFormatter for SimpleFormatter {
10    fn format(&self, result: &TokenizationResult) -> String {
11        result.token_count.to_string()
12    }
13}
14
15#[cfg(test)]
16mod tests {
17    use super::*;
18    use crate::tokenizers::ModelInfo;
19
20    #[test]
21    fn test_simple_formatter() {
22        let formatter = SimpleFormatter;
23        let result = TokenizationResult {
24            token_count: 42,
25            model_info: ModelInfo {
26                name: "gpt-4".to_string(),
27                encoding: "cl100k_base".to_string(),
28                context_window: 128000,
29                description: "GPT-4".to_string(),
30            },
31        };
32
33        let output = formatter.format(&result);
34        assert_eq!(output, "42");
35    }
36
37    #[test]
38    fn test_zero_tokens() {
39        let formatter = SimpleFormatter;
40        let result = TokenizationResult {
41            token_count: 0,
42            model_info: ModelInfo {
43                name: "gpt-4".to_string(),
44                encoding: "cl100k_base".to_string(),
45                context_window: 128000,
46                description: "GPT-4".to_string(),
47            },
48        };
49
50        let output = formatter.format(&result);
51        assert_eq!(output, "0");
52    }
53}