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
29impl Cli {
30    /// Parse command-line arguments
31    pub fn parse_args() -> Self {
32        Self::parse()
33    }
34
35    /// Normalize model name (lowercase, trim whitespace)
36    pub fn normalized_model(&self) -> String {
37        self.model.trim().to_lowercase()
38    }
39
40    /// Get verbosity level (0-3)
41    pub fn verbosity_level(&self) -> u8 {
42        self.verbose.min(3)
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::*;
49
50    #[test]
51    fn test_default_model() {
52        let cli = Cli::parse_from(["token-count"]);
53        assert_eq!(cli.model, "gpt-3.5-turbo");
54        assert_eq!(cli.verbose, 0);
55        assert!(!cli.list_models);
56    }
57
58    #[test]
59    fn test_model_normalization() {
60        let cli = Cli::parse_from(["token-count", "--model", "  GPT-4  "]);
61        assert_eq!(cli.normalized_model(), "gpt-4");
62    }
63
64    #[test]
65    fn test_verbosity_levels() {
66        let cli1 = Cli::parse_from(["token-count", "-v"]);
67        assert_eq!(cli1.verbosity_level(), 1);
68
69        let cli2 = Cli::parse_from(["token-count", "-vv"]);
70        assert_eq!(cli2.verbosity_level(), 2);
71
72        let cli3 = Cli::parse_from(["token-count", "-vvv"]);
73        assert_eq!(cli3.verbosity_level(), 3);
74
75        let cli4 = Cli::parse_from(["token-count", "-vvvv"]);
76        assert_eq!(cli4.verbosity_level(), 3); // Capped at 3
77    }
78
79    #[test]
80    fn test_list_models_flag() {
81        let cli = Cli::parse_from(["token-count", "--list-models"]);
82        assert!(cli.list_models);
83    }
84}
token_count/cli/args.rs

token_count/cli/
args.rs
