Calculate tokens, handles 1 file copy.

src/cli.rs

@@ -6,4 +6,7 @@ use clap::Parser;
 pub struct Opts {
     #[clap(default_value = "./")]
     pub path: String,
+
+    #[clap(short, long)]
+    pub token: bool,
 }

src/fs.rs

@@ -2,13 +2,26 @@ use std::fs;
 use std::io::Read;
 use std::path::Path;
 
-pub fn read_directory_contents(dir: &str) -> anyhow::Result<String> {
+pub fn read_directory_contents(dir: &Path) -> anyhow::Result<String> {
     let mut combined_content = String::new();
-    let base_path = Path::new(dir);
-    read_directory_contents_recursive(base_path, base_path, &mut combined_content)?;
+
+    if dir.is_file() {
+        read_file(dir, &mut combined_content)?;
+    } else {
+        read_directory_contents_recursive(dir, dir, &mut combined_content)?;
+    }
+
     Ok(combined_content)
 }
 
+pub fn read_file(path: &Path, content: &mut String) -> anyhow::Result<()> {
+    let file_content = std::fs::read_to_string(path)?;
+
+    *content = file_content;
+
+    Ok(())
+}
+
 fn read_directory_contents_recursive(
     base_path: &Path,
     current_path: &Path,

src/run.rs

@@ -1,13 +1,21 @@
 use clap::Parser;
+use std::path::Path;
 
-use crate::{cli, clip::copy_to_clipboard, fs::read_directory_contents};
+use crate::{cli, clip::copy_to_clipboard, fs::read_directory_contents, tiktoken::count_tokens};
 
 pub fn run() -> anyhow::Result<()> {
     let opts = cli::Opts::parse();
-    let path = opts.path;
+    let (path, token) = (opts.path, opts.token);
+
+    let path = Path::new(&path);
+    let contents = read_directory_contents(path)?;
 
-    let contents = read_directory_contents(&path)?;
     copy_to_clipboard(&contents)?;
 
+    if token {
+        let tokens = count_tokens(&contents)?;
+        println!("{} GPT-4 tokens.", tokens);
+    }
+
     Ok(())
 }

src/tiktoken.rs

@@ -1,7 +1,5 @@
 use tiktoken_rs::o200k_base;
 
-// TODO: remove this dead code
-#[allow(dead_code)]
 pub fn count_tokens(string: &str) -> anyhow::Result<usize> {
     let bpe = o200k_base()?;
     let tokens = bpe.encode_with_special_tokens(string);