agg/cmds/
tok.rs

1use std::ffi::OsString;
2use std::io::Read;
3use std::path::PathBuf;
4
5use rootcause::prelude::ResultExt;
6use rootcause::report;
7use ytil_sys::pico_args::Arguments;
8
9const DEFAULT_ENCODING: &str = "o200k_base";
10
11#[derive(Debug, Eq, PartialEq)]
12pub struct Opts {
13    pub encoding: String,
14    pub input: Input,
15}
16
17#[derive(Debug, Eq, PartialEq)]
18pub enum Input {
19    File(PathBuf),
20    Stdin,
21    Text(String),
22}
23
24impl TryFrom<Vec<OsString>> for Opts {
25    type Error = rootcause::Report;
26
27    fn try_from(raw: Vec<OsString>) -> Result<Self, Self::Error> {
28        let mut args = Arguments::from_vec(raw);
29        let encoding = args
30            .opt_value_from_str::<_, String>("--encoding")?
31            .unwrap_or_else(|| DEFAULT_ENCODING.to_owned());
32        let text = args.opt_value_from_str::<_, String>("--text")?;
33        let positionals = args.finish();
34
35        if let Some(argument) = positionals
36            .iter()
37            .find(|argument| argument.to_string_lossy().starts_with('-') && argument.as_os_str() != "-")
38        {
39            return Err(report!("unknown agg tok option").attach(format!("option={}", argument.to_string_lossy())));
40        }
41
42        let input = match (text, positionals.as_slice()) {
43            (Some(text), []) => Input::Text(text),
44            (Some(_), _) => return Err(report!("agg tok cannot combine --text with a file or stdin")),
45            (None, [path]) if path == "-" => Input::Stdin,
46            (None, [path]) => Input::File(PathBuf::from(path)),
47            (None, []) => return Err(report!("agg tok requires a file, --text value, or - for stdin")),
48            (None, _) => return Err(report!("agg tok accepts exactly one file or stdin input")),
49        };
50
51        Ok(Self { encoding, input })
52    }
53}
54
55impl Input {
56    fn read_to_string(&self) -> rootcause::Result<String> {
57        match self {
58            Self::File(path) => Ok(std::fs::read_to_string(path)
59                .attach_with(|| format!("cannot read token input file | path={}", path.display()))?),
60            Self::Stdin => {
61                let mut text = String::new();
62                std::io::stdin()
63                    .read_to_string(&mut text)
64                    .context("cannot read token input from stdin")?;
65                Ok(text)
66            }
67            Self::Text(text) => Ok(text.clone()),
68        }
69    }
70}
71
72pub fn run(options: &Opts) -> rootcause::Result<()> {
73    let encoding = tiktoken::get_encoding(&options.encoding).ok_or_else(|| {
74        report!("unknown token encoding")
75            .attach(format!("encoding={}", options.encoding))
76            .attach(format!("supported={}", tiktoken::list_encodings().join(", ")))
77    })?;
78    let text = options.input.read_to_string()?;
79    let tokens = encoding.count(&text);
80
81    println!("{tokens}");
82    Ok(())
83}
84
85#[cfg(test)]
86mod tests {
87    use std::path::Path;
88
89    use test_that::prelude::*;
90
91    use super::*;
92
93    #[test]
94    fn test_input_when_text_returns_the_text() {
95        let input = Input::Text("hello world".to_owned());
96
97        assert_that!(input.read_to_string(), ok(eq("hello world")));
98    }
99
100    #[test]
101    fn test_input_when_file_reads_utf8_contents() {
102        let directory = tempfile::tempdir().expect("test directory should be created");
103        let path = directory.path().join("prompt.txt");
104        std::fs::write(&path, "hello world").expect("test input should be written");
105
106        assert_that!(Input::File(path).read_to_string(), ok(eq("hello world")));
107    }
108
109    #[test]
110    fn test_input_when_file_is_missing_reports_the_path() {
111        let path = Path::new("missing-prompt.txt");
112
113        assert_that!(
114            (Input::File(path.to_owned()).read_to_string()).map(|_| ()),
115            err(displays_as(all!(
116                contains_substring("cannot read token input file"),
117                contains_substring(path.to_str().expect("test path should be valid UTF-8"))
118            )))
119        );
120    }
121
122    #[test]
123    fn test_count_with_default_encoding_counts_known_text() {
124        let encoding = tiktoken::get_encoding(DEFAULT_ENCODING).expect("default encoding should be available");
125
126        assert_eq!(encoding.count("hello world"), 2);
127    }
128
129    #[test]
130    fn test_run_when_encoding_is_unknown_reports_supported_encodings() {
131        let options = Opts {
132            encoding: "unknown".to_owned(),
133            input: Input::Text("hello".to_owned()),
134        };
135
136        assert_that!(
137            run(&options),
138            err(displays_as(all!(
139                contains_substring("unknown token encoding"),
140                contains_substring("encoding=unknown"),
141                contains_substring(DEFAULT_ENCODING)
142            )))
143        );
144    }
145
146    #[rstest::rstest]
147    #[case::missing_input(&[])]
148    #[case::multiple_files(&["first.txt", "second.txt"])]
149    #[case::text_and_file(&["--text", "hello", "prompt.txt"])]
150    #[case::unknown_option(&["--unknown", "prompt.txt"])]
151    fn test_options_when_input_arguments_are_invalid_rejects_command(#[case] args: &[&str]) {
152        assert_that!(parse(args), err(anything()));
153    }
154
155    fn parse(args: &[&str]) -> rootcause::Result<Opts> {
156        Opts::try_from(args.iter().map(OsString::from).collect::<Vec<_>>())
157    }
158}
agg/cmds/tok.rs

agg/cmds/
tok.rs
