token_count/cli/
args.rs

1//! Command-line argument parsing
2
3use clap::Parser;
4
5/// Count tokens for LLM models using exact tokenization
6#[derive(Parser, Debug)]
7#[command(name = "token-count")]
8#[command(version, about)]
9#[command(after_help = "\
10EXAMPLES:
11    echo \"Hello world\" | token-count --model gpt-4
12    token-count --model gpt-4 < file.txt
13    token-count --list-models
14")]
15pub struct Cli {
16    /// Model to use for tokenization (use --list-models to see all)
17    #[arg(short, long, default_value = "gpt-3.5-turbo")]
18    pub model: String,
19
20    /// Increase output verbosity (-v, -vv, -vvv for debug)
21    #[arg(short, long, action = clap::ArgAction::Count)]
22    pub verbose: u8,
23
24    /// List all supported models and exit
25    #[arg(long)]
26    pub list_models: bool,
27
28    /// Use API for exact token counts (requires ANTHROPIC_API_KEY for Claude models)
29    #[arg(long)]
30    pub accurate: bool,
31
32    /// Skip API consent prompt (for scripting/automation, requires --accurate)
33    #[arg(short = 'y', long)]
34    pub yes: bool,
35}
36
37impl Cli {
38    /// Parse command-line arguments
39    pub fn parse_args() -> Self {
40        Self::parse()
41    }
42
43    /// Normalize model name (lowercase, trim whitespace)
44    pub fn normalized_model(&self) -> String {
45        self.model.trim().to_lowercase()
46    }
47
48    /// Get verbosity level (0-3)
49    pub fn verbosity_level(&self) -> u8 {
50        self.verbose.min(3)
51    }
52}
53
54#[cfg(test)]
55mod tests {
56    use super::*;
57
58    #[test]
59    fn test_default_model() {
60        let cli = Cli::parse_from(["token-count"]);
61        assert_eq!(cli.model, "gpt-3.5-turbo");
62        assert_eq!(cli.verbose, 0);
63        assert!(!cli.list_models);
64    }
65
66    #[test]
67    fn test_model_normalization() {
68        let cli = Cli::parse_from(["token-count", "--model", "  GPT-4  "]);
69        assert_eq!(cli.normalized_model(), "gpt-4");
70    }
71
72    #[test]
73    fn test_verbosity_levels() {
74        let cli1 = Cli::parse_from(["token-count", "-v"]);
75        assert_eq!(cli1.verbosity_level(), 1);
76
77        let cli2 = Cli::parse_from(["token-count", "-vv"]);
78        assert_eq!(cli2.verbosity_level(), 2);
79
80        let cli3 = Cli::parse_from(["token-count", "-vvv"]);
81        assert_eq!(cli3.verbosity_level(), 3);
82
83        let cli4 = Cli::parse_from(["token-count", "-vvvv"]);
84        assert_eq!(cli4.verbosity_level(), 3); // Capped at 3
85    }
86
87    #[test]
88    fn test_list_models_flag() {
89        let cli = Cli::parse_from(["token-count", "--list-models"]);
90        assert!(cli.list_models);
91    }
92
93    #[test]
94    fn test_accurate_flag() {
95        let cli = Cli::parse_from(["token-count", "--model", "claude", "--accurate"]);
96        assert!(cli.accurate);
97        assert!(!cli.yes);
98    }
99
100    #[test]
101    fn test_yes_flag() {
102        let cli = Cli::parse_from(["token-count", "--model", "claude", "--accurate", "-y"]);
103        assert!(cli.accurate);
104        assert!(cli.yes);
105    }
106
107    #[test]
108    fn test_yes_long_form() {
109        let cli = Cli::parse_from(["token-count", "--model", "claude", "--accurate", "--yes"]);
110        assert!(cli.accurate);
111        assert!(cli.yes);
112    }
113}
token_count/cli/args.rs

token_count/cli/
args.rs
