token_count/output/
simple.rs1use crate::output::OutputFormatter;
4use crate::tokenizers::TokenizationResult;
5
6pub 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}