45 lines
906 B
TypeScript
45 lines
906 B
TypeScript
/**
|
|
* @license
|
|
* Copyright 2026 Google LLC
|
|
* SPDX-License-Identifier: Apache-2.0
|
|
*/
|
|
|
|
import {readFileSync} from 'node:fs';
|
|
import {parseArgs} from 'node:util';
|
|
|
|
import {GoogleGenAI} from '@google/genai';
|
|
|
|
const ai = new GoogleGenAI({apiKey: process.env.GEMINI_API_KEY});
|
|
|
|
const {values, positionals} = parseArgs({
|
|
options: {
|
|
model: {
|
|
type: 'string',
|
|
default: 'gemini-2.5-flash',
|
|
},
|
|
file: {
|
|
type: 'string',
|
|
short: 'f',
|
|
},
|
|
},
|
|
allowPositionals: true,
|
|
});
|
|
|
|
let contents = positionals[0];
|
|
|
|
if (values.file) {
|
|
contents = readFileSync(values.file, 'utf8');
|
|
}
|
|
|
|
if (!contents) {
|
|
console.error('Usage: npm run count-tokens -- [-f <file>] [<text>]');
|
|
process.exit(1);
|
|
}
|
|
|
|
const response = await ai.models.countTokens({
|
|
model: values.model,
|
|
contents,
|
|
});
|
|
console.log(`Input: ${values.file || positionals[0]}`);
|
|
console.log(`Tokens: ${response.totalTokens}`);
|