diff --git a/CHANGELOG.md b/CHANGELOG.md index b3ed778..9759570 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Added + +- CLI `--tokenizer` option to use a local module's `count(text)` function for + analysis and compaction without adding runtime dependencies. + ### Fixed - Invalid or negative numeric CLI flags now fail clearly instead of producing diff --git a/README.md b/README.md index 64aaf74..368c390 100644 --- a/README.md +++ b/README.md @@ -91,6 +91,39 @@ be an integer; invalid values throw a `RangeError` before the payload changes. Token counts are a fast **estimate** (a chars/word blend), typically within ~10-15% of exact BPE counts, enough to find waste and compare before/after. For exact counts, pass your own counter: `analyzePayload(payload, { counter: myTokenizer })`. +### Use a tokenizer from the CLI + +`--tokenizer ` loads a local JavaScript module exporting a synchronous +`count(text)` function that returns a numeric token count. Relative paths resolve +from your current working directory, not from the payload file. The same counter +is used for analysis and compaction; without the flag, the built-in estimate is +unchanged. Only load modules you trust, since they execute in the CLI process. + +For example, install [gpt-tokenizer](https://github.com/niieani/gpt-tokenizer) in +your own project (it is not a tokencut dependency): + +```bash +npm install gpt-tokenizer +``` + +Save this as `my-counter.mjs` beside that project's `package.json`: + +```js +import { encode } from "gpt-tokenizer"; +export function count(text) { + return encode(text).length; +} +``` + +```bash +tokencut payload.json --tokenizer ./my-counter.mjs --json +tokencut payload.json --compact --max 8000 --tokenizer ./my-counter.mjs +``` + +Choose an encoding appropriate for your model. This replaces text-unit counts; +it does not add provider-specific message framing or replace image-token +estimates, and compaction still preserves protected messages. + ## CLI reference ``` @@ -101,6 +134,7 @@ tokencut --compact cut, print savings --no-dedupe keep duplicate blocks --out write the compacted payload --price $ per 1M input tokens for the estimate (default 3) + --tokenizer load a local module exporting count(text) -> number --json machine-readable output ``` diff --git a/bin/tokencut.mjs b/bin/tokencut.mjs index 37a14f1..34c2cfa 100644 --- a/bin/tokencut.mjs +++ b/bin/tokencut.mjs @@ -1,10 +1,12 @@ #!/usr/bin/env node import { readFileSync, writeFileSync } from "node:fs"; +import { resolve } from "node:path"; +import { pathToFileURL } from "node:url"; import { analyzePayload, compact } from "../src/index.mjs"; const args = process.argv.slice(2); // Value-taking flags: when one is seen, the next argv entry is its value, not a file path. -const VALUE_FLAGS = new Set(["--max", "--max-tool", "--price", "--out"]); +const VALUE_FLAGS = new Set(["--max", "--max-tool", "--price", "--out", "--tokenizer"]); const options = {}; const positionals = []; for (let i = 0; i < args.length; i++) { @@ -41,6 +43,7 @@ if (!file || has("--help")) { --no-dedupe keep duplicate context blocks --out write the compacted payload --price $ per 1M input tokens for the cost estimate (default 3) + --tokenizer load a local module exporting count(text) -> number --json machine-readable output --version, -V print package version @@ -70,6 +73,25 @@ let payload; try { payload = JSON.parse(readFileSync(file, "utf8")); } catch (e) { console.error(`could not read ${file}: ${e.message}`); process.exit(1); } +let counter; +if (has("--tokenizer")) { + const modulePath = flag("--tokenizer"); + if (typeof modulePath !== "string") { + console.error("--tokenizer requires a module path"); + process.exit(1); + } + try { + const tokenizer = await import(pathToFileURL(resolve(modulePath)).href); + if (typeof tokenizer.count !== "function") { + throw new TypeError("module must export a count(text) function"); + } + counter = tokenizer.count; + } catch (e) { + console.error(`could not load tokenizer ${modulePath}: ${e.message}`); + process.exit(1); + } +} + const k = (n) => (n >= 1000 ? (n / 1000).toFixed(1).replace(/\.0$/, "") + "k" : "" + n); const usd = (n) => "$" + n.toFixed(n < 0.01 ? 5 : 4); @@ -78,6 +100,7 @@ if (has("--compact")) { maxTokens, maxToolResultTokens, dropDuplicates: !has("--no-dedupe"), + counter, }); if (flag("--out", null)) { const outPath = String(flag("--out", null)); @@ -99,7 +122,7 @@ if (has("--compact")) { if (flag("--out", null)) console.log(` wrote ${flag("--out", null)}`); console.log(); } else { - const a = analyzePayload(payload, { pricePerMTok: price }); + const a = analyzePayload(payload, { pricePerMTok: price, counter }); if (has("--json")) { console.log(JSON.stringify(a, null, 2)); process.exit(0); } console.log(`\n tokencut ${k(a.totalTokens)} tokens (~${usd(a.costUSD)} at $${price}/M) across ${a.units} blocks\n`); const row = (label, obj) => { diff --git a/test/cli.test.mjs b/test/cli.test.mjs index ebf4f97..9cca653 100644 --- a/test/cli.test.mjs +++ b/test/cli.test.mjs @@ -3,6 +3,7 @@ import { mkdtempSync, readFileSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import { spawnSync } from "node:child_process"; +import { fileURLToPath } from "node:url"; import test, { after } from "node:test"; const fixtureDir = mkdtempSync(join(tmpdir(), "tokencut-cli-")); @@ -79,3 +80,72 @@ test("rejects compact budget flags without --compact", async (t) => { }); } }); + +const tokenizer = join(fixtureDir, "counter #1.mjs"); +writeFileSync(tokenizer, "export const count = (text) => text.length;\n"); + +test("analyze keeps the built-in estimate when no tokenizer is supplied", () => { + const result = run("--json"); + assert.equal(result.status, 0, result.stderr); + assert.equal(JSON.parse(result.stdout).totalTokens, 2); +}); + +test("analyze uses a tokenizer module, including spaces and URL-special characters", () => { + const result = run("--tokenizer", tokenizer, "--json"); + assert.equal(result.status, 0, result.stderr); + const report = JSON.parse(result.stdout); + assert.equal(report.totalTokens, 5); + assert.equal(report.biggest[0].tokens, 5); +}); + +test("resolves a relative tokenizer from the caller's directory before the payload", () => { + const result = spawnSync(process.execPath, [ + fileURLToPath(new URL("../bin/tokencut.mjs", import.meta.url)), + "--tokenizer", "./counter #1.mjs", fixture, "--json", + ], { cwd: fixtureDir, encoding: "utf8" }); + assert.equal(result.status, 0, result.stderr); + assert.equal(JSON.parse(result.stdout).totalTokens, 5); +}); + +test("compact uses the tokenizer for its budget and before/after report", () => { + const payload = join(fixtureDir, "history.json"); + const output = join(fixtureDir, "compacted.json"); + const messages = Array.from({ length: 6 }, (_, i) => ({ role: "user", content: `hello-${i}` })); + writeFileSync(payload, JSON.stringify(messages)); + const result = spawnSync(process.execPath, [ + "bin/tokencut.mjs", payload, "--tokenizer", tokenizer, + "--compact", "--max", "28", "--out", output, "--json", + ], { cwd: new URL("..", import.meta.url), encoding: "utf8" }); + assert.equal(result.status, 0, result.stderr); + const report = JSON.parse(result.stdout); + assert.equal(report.beforeTokens, 42); + assert.equal(report.afterTokens, 28); + assert.equal(report.savedTokens, 14); + assert.deepEqual(JSON.parse(readFileSync(output, "utf8")), messages.slice(-4)); +}); + +test("rejects a missing tokenizer argument without a stack trace", () => { + const result = run("--tokenizer", "--json"); + assert.equal(result.status, 1); + assert.equal(result.stdout, ""); + assert.equal(result.stderr, "--tokenizer requires a module path\n"); +}); + +test("reports tokenizer import and export errors without a stack trace", async (t) => { + for (const [name, source, expected] of [ + ["missing.mjs", null, /could not load tokenizer/], + ["no-count.mjs", "export default () => 1;", /must export a count\(text\) function/], + ["not-callable.mjs", "export const count = 1;", /must export a count\(text\) function/], + ["broken.mjs", 'throw new Error("broken module");', /broken module/], + ]) { + await t.test(name, () => { + const path = join(fixtureDir, name); + if (source !== null) writeFileSync(path, source); + const result = run("--tokenizer", path, "--json"); + assert.equal(result.status, 1); + assert.equal(result.stdout, ""); + assert.match(result.stderr, expected); + assert.doesNotMatch(result.stderr, /\n\s+at /); + }); + } +});