Add token counter

Author: yleflour

bun.lockb

@@ -35,6 +35,7 @@
     "ora": "^7.0.1",
     "pretty-bytes": "^6.1.1",
     "rxjs": "^7.8.1",
+    "tiktoken": "^1.0.15",
     "wretch": "^2.8.1",
     "xml2js": "^0.6.2",
     "zod": "^3.22.4",

package.json

@@ -78,7 +78,7 @@ export const promptAssistantConfig = async (
 
   const shouldEditInstructions = await confirm({
     message: "Edit instructions?",
-    default: false,
+    default: true,
   });
 
   let instructions = assistant?.instructions ?? "";

src/cli/assistant/assistant.utils.ts

@@ -9,14 +9,15 @@ import { listAssistantsAction } from "./assistant/list-assistants.cli.js";
 import { createAssistantAction } from "./assistant/create-assistant.cli.js";
 import { updateAssistantAction } from "./assistant/update-assistant.cli.js";
 import { deleteAssistantAction } from "./assistant/delete-assistant.cli.js";
+import { countTokensAction } from "./helpers/count-tokens.cli.js";
 
 const program = new Command();
 program
   .name("ai")
   .description("Use your OpenAI assistant from the command line")
   .version("2.0.0");
 
-const runCommand = program
+program
   .command("chat", { isDefault: true, hidden: true })
   .argument("[args...]")
   .description("Start a chat with an assistant")
@@ -29,6 +30,18 @@ const runCommand = program
     await appAction();
   });
 
+const utilsCommand = program
+  .command("utils")
+  .alias("u")
+  .description("Useful utilities");
+
+utilsCommand
+  .command("count-tokens")
+  .alias("ct")
+  .description("Count tokens in a project")
+  .argument("[globs]", "Globs to count tokens from")
+  .action(countTokensAction);
+
 const assistantsCommand = program
   .command("assistants")
   .alias("ass")

src/cli/entry.cli.ts

@@ -0,0 +1,15 @@
+import ora from "ora";
+import { countGlobToken } from "../../utils/tokens.utils.js";
+
+export const countTokensAction = async (globFilters?: string) => {
+  const spinner = ora({
+    text: "Counting tokens",
+    color: "blue",
+  }).start();
+
+  const res = await countGlobToken(globFilters);
+
+  spinner.stopAndPersist({
+    text: `Found ${res.tokensCount} tokens accross ${res.filesCount} files (OpenAI tokens)`,
+  });
+};

src/cli/helpers/count-tokens.cli.ts

@@ -0,0 +1,10 @@
+import { promises as fs } from "fs";
+import { join } from "path";
+
+export function readFile(path: string, cwd?: string) {
+  const fullPath = cwd ? join(cwd, path) : path;
+  return fs.readFile(fullPath, "utf-8").catch((e) => {
+    if (e.code === "ENOENT") throw new Error("File does not exist");
+    throw e;
+  });
+}

src/utils/fs.utils.ts

@@ -0,0 +1,33 @@
+import { get_encoding } from "tiktoken";
+import { readFile } from "./fs.utils.js";
+import { glob } from "zx";
+
+export function countTokens(text: string): number {
+  const encoding = get_encoding("cl100k_base");
+  const length = encoding.encode(text).length;
+  encoding.free();
+  return length;
+}
+
+export async function countFileToken(path: string, cwd?: string) {
+  const content = await readFile(path, cwd);
+  return countTokens(content);
+}
+
+export async function countGlobToken(globsString: string = "**/*") {
+  const globs = globsString.split(",");
+  const includes = globs.filter((g) => !g.startsWith("!"));
+  const excludes = globs.filter((g) => g.startsWith("!"));
+
+  const paths = await glob(includes, {
+    ignore: excludes,
+  });
+
+  if (!paths) throw new Error("No files found");
+
+  const sizes = await Promise.all(paths.map((path) => countFileToken(path)));
+  return {
+    filesCount: paths.length,
+    tokensCount: sizes.reduce((acc, size) => acc + size, 0),
+  };
+}