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
101
102
103
104
105
106
107
108
109
110
111
112
113

//! Command-line argument parsing

use clap::Parser;

/// Count tokens for LLM models using exact tokenization
#[derive(Parser, Debug)]
#[command(name = "token-count")]
#[command(version, about)]
#[command(after_help = "\
EXAMPLES:
    echo \"Hello world\" | token-count --model gpt-4
    token-count --model gpt-4 < file.txt
    token-count --list-models
")]
pub struct Cli {
    /// Model to use for tokenization (use --list-models to see all)
    #[arg(short, long, default_value = "gpt-3.5-turbo")]
    pub model: String,

    /// Increase output verbosity (-v, -vv, -vvv for debug)
    #[arg(short, long, action = clap::ArgAction::Count)]
    pub verbose: u8,

    /// List all supported models and exit
    #[arg(long)]
    pub list_models: bool,

    /// Use API for exact token counts (requires ANTHROPIC_API_KEY for Claude models)
    #[arg(long)]
    pub accurate: bool,

    /// Skip API consent prompt (for scripting/automation, requires --accurate)
    #[arg(short = 'y', long)]
    pub yes: bool,
}

impl Cli {
    /// Parse command-line arguments
    pub fn parse_args() -> Self {
        Self::parse()
    }

    /// Normalize model name (lowercase, trim whitespace)
    pub fn normalized_model(&self) -> String {
        self.model.trim().to_lowercase()
    }

    /// Get verbosity level (0-3)
    pub fn verbosity_level(&self) -> u8 {
        self.verbose.min(3)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_default_model() {
        let cli = Cli::parse_from(["token-count"]);
        assert_eq!(cli.model, "gpt-3.5-turbo");
        assert_eq!(cli.verbose, 0);
        assert!(!cli.list_models);
    }

    #[test]
    fn test_model_normalization() {
        let cli = Cli::parse_from(["token-count", "--model", "  GPT-4  "]);
        assert_eq!(cli.normalized_model(), "gpt-4");
    }

    #[test]
    fn test_verbosity_levels() {
        let cli1 = Cli::parse_from(["token-count", "-v"]);
        assert_eq!(cli1.verbosity_level(), 1);

        let cli2 = Cli::parse_from(["token-count", "-vv"]);
        assert_eq!(cli2.verbosity_level(), 2);

        let cli3 = Cli::parse_from(["token-count", "-vvv"]);
        assert_eq!(cli3.verbosity_level(), 3);

        let cli4 = Cli::parse_from(["token-count", "-vvvv"]);
        assert_eq!(cli4.verbosity_level(), 3); // Capped at 3
    }

    #[test]
    fn test_list_models_flag() {
        let cli = Cli::parse_from(["token-count", "--list-models"]);
        assert!(cli.list_models);
    }

    #[test]
    fn test_accurate_flag() {
        let cli = Cli::parse_from(["token-count", "--model", "claude", "--accurate"]);
        assert!(cli.accurate);
        assert!(!cli.yes);
    }

    #[test]
    fn test_yes_flag() {
        let cli = Cli::parse_from(["token-count", "--model", "claude", "--accurate", "-y"]);
        assert!(cli.accurate);
        assert!(cli.yes);
    }

    #[test]
    fn test_yes_long_form() {
        let cli = Cli::parse_from(["token-count", "--model", "claude", "--accurate", "--yes"]);
        assert!(cli.accurate);
        assert!(cli.yes);
    }
}