token-count 0.4.0

Count tokens for LLM models using exact tokenization
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100

use token_count::output::{
    BasicFormatter, DebugFormatter, OutputFormatter, SimpleFormatter, VerboseFormatter,
};
use token_count::tokenizers::{ModelInfo, TokenizationResult};

fn create_test_result() -> TokenizationResult {
    TokenizationResult {
        token_count: 2,
        model_info: ModelInfo {
            name: "gpt-4".to_string(),
            encoding: "cl100k_base".to_string(),
            context_window: 128000,
            description: "GPT-4 model".to_string(),
        },
        token_details: None,
    }
}

#[test]
fn test_simple_formatter_output() {
    let formatter = SimpleFormatter;
    let result = create_test_result();
    let output = formatter.format(&result);

    assert_eq!(output, "2");
    assert!(!output.contains('\n'));
}

#[test]
fn test_basic_formatter_output() {
    let formatter = BasicFormatter;
    let result = create_test_result();
    let output = formatter.format(&result);

    assert!(output.contains("Model: gpt-4"));
    assert!(output.contains("cl100k_base"));
    assert!(output.contains("Tokens: 2"));
    assert!(!output.contains("Context window"), "Should not contain context window");
    assert!(!output.contains('%'), "Should not contain percentage");

    // Should be multi-line
    assert!(output.contains('\n'));
}

#[test]
fn test_verbose_formatter_output() {
    let formatter = VerboseFormatter;
    let result = create_test_result();
    let output = formatter.format(&result);

    assert!(output.contains("Model: gpt-4"));
    assert!(output.contains("cl100k_base"));
    assert!(output.contains("Tokens: 2"));
    assert!(output.contains("Context window: 128000"));
    assert!(output.contains('%'));

    // Should be multi-line
    assert!(output.contains('\n'));
}

#[test]
fn test_debug_formatter_output() {
    let formatter = DebugFormatter;
    let result = create_test_result();
    let output = formatter.format(&result);

    assert!(output.contains("Model: gpt-4"));
    assert!(output.contains("Tokens: 2"));
    assert!(output.contains("Token IDs not available"));

    // Should be multi-line
    assert!(output.contains('\n'));
}

#[test]
fn test_formatter_selection() {
    let result = create_test_result();

    // Verbosity 0 -> Simple
    let f0 = token_count::output::select_formatter(0);
    let out0 = f0.format(&result);
    assert_eq!(out0, "2");

    // Verbosity 1 -> Basic (no percentage)
    let f1 = token_count::output::select_formatter(1);
    let out1 = f1.format(&result);
    assert!(out1.contains("Model:"));
    assert!(!out1.contains('%'), "Level 1 should not have percentage");

    // Verbosity 2 -> Verbose (with percentage)
    let f2 = token_count::output::select_formatter(2);
    let out2 = f2.format(&result);
    assert!(out2.contains("Model:"));
    assert!(out2.contains('%'), "Level 2 should have percentage");

    // Verbosity 3+ -> Debug
    let f3 = token_count::output::select_formatter(3);
    let out3 = f3.format(&result);
    assert!(out3.contains("Token IDs not available"));
}