openai-tokens-count
Version:
OpenAI tokens calculator with function calls, images, and messages in one call
218 lines • 9.3 kB
JavaScript
;
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.imageTokenCache = exports.toolTokenCache = exports.messageTokenCache = exports.estimateTokensInTools = exports.estimateTokensInMessages = exports.getCachedEncoding = exports.estimateTokens = void 0;
const js_tiktoken_1 = require("js-tiktoken");
const lru_cache_1 = require("lru-cache");
const murmurhash_1 = __importDefault(require("murmurhash"));
const argument_format_1 = require("./argument-format");
const function_format_1 = require("./function-format");
const tool_content_format_1 = require("./tool-content-format");
const sharp_1 = __importDefault(require("sharp"));
const vision_constants_1 = require("./vision-constants");
// Create stringifiers
const stringifyTools = JSON.stringify;
const stringifyMessage = JSON.stringify;
const stringifyImage = JSON.stringify;
const HASH = (str) => murmurhash_1.default.v3(str).toString(16);
const HASHKEY = (model, str) => `${model}-${HASH(str)}`;
// Global cache for encoding objects
const encodingCache = new Map();
// Global caches for token counts
const messageTokenCache = new lru_cache_1.LRUCache({ max: 1000000 });
exports.messageTokenCache = messageTokenCache;
const toolTokenCache = new lru_cache_1.LRUCache({ max: 1000000 });
exports.toolTokenCache = toolTokenCache;
const imageTokenCache = new lru_cache_1.LRUCache({ max: 1000000 });
exports.imageTokenCache = imageTokenCache;
function getCachedEncoding(model) {
const encodingName = model.startsWith('gpt-4o') ? "o200k_base" : (0, js_tiktoken_1.getEncodingNameForModel)(model);
if (!encodingCache.has(encodingName)) {
const encoding = (0, js_tiktoken_1.getEncoding)(encodingName);
encodingCache.set(encodingName, encoding);
}
return encodingCache.get(encodingName);
}
exports.getCachedEncoding = getCachedEncoding;
async function estimateTokens(request) {
var _a;
const messages = request.messages;
const tools = request.tools;
const toolChoice = request.tool_choice;
const chatModel = request.model;
let tokens = 0;
tokens += await estimateTokensInMessages(chatModel, messages, tools);
if (tools) {
tokens += estimateTokensInTools(chatModel, tools);
}
if (tools && messages.some((msg) => msg.role === "system")) {
tokens -= 4;
}
if (toolChoice && toolChoice !== "auto") {
if (toolChoice === "none") {
tokens += 1;
}
else if (typeof toolChoice === "object") {
const tc = toolChoice;
if ((_a = tc.function) === null || _a === void 0 ? void 0 : _a.name) {
tokens += countTokens(getCachedEncoding(chatModel), tc.function.name) + 4;
}
}
}
return tokens;
}
exports.estimateTokens = estimateTokens;
function estimateTokensInTools(chatModel, tools) {
const cacheKey = HASHKEY(chatModel, stringifyTools(tools));
if (toolTokenCache.has(cacheKey)) {
return toolTokenCache.get(cacheKey);
}
const definitions = (0, function_format_1.formatFunctionDefinitions)(tools);
let tokens = countTokens(getCachedEncoding(chatModel), definitions);
tokens += 2; // Additional tokens for function definition of tools
toolTokenCache.set(cacheKey, tokens);
return tokens;
}
exports.estimateTokensInTools = estimateTokensInTools;
async function estimateTokensInMessages(chatModel, messages, tools) {
let tokens = 0;
let paddedSystem = false;
for (const message of messages) {
const msg = { ...message };
if (msg.role === "system" && tools && !paddedSystem) {
if (typeof msg.content === "string") {
msg.content += "\n";
}
paddedSystem = true;
}
tokens += await estimateTokensInMessage(chatModel, msg, 1);
}
tokens += 3; // Each completion (vs message) seems to carry a 3-token overhead
return tokens;
}
exports.estimateTokensInMessages = estimateTokensInMessages;
async function estimateTokensInMessage(chatModel, message, toolMessageSize) {
var _a;
const cacheKey = HASHKEY(chatModel, stringifyMessage(message));
if (messageTokenCache.has(cacheKey)) {
return messageTokenCache.get(cacheKey);
}
let tokens = 0;
const encoding = getCachedEncoding(chatModel);
tokens += countTokens(encoding, message.role);
if (message.role === "tool") {
if (toolMessageSize === 1) {
tokens += countTokens(encoding, message.content);
}
else {
tokens += countTokens(encoding, (0, tool_content_format_1.formatToolContent)(message.content));
const contentJSON = (0, tool_content_format_1.tryFormatJSON)(message.content);
if (contentJSON) {
tokens -= Object.keys(contentJSON).length;
}
}
}
else if (typeof message.content === "string") {
tokens += countTokens(encoding, message.content);
}
else if (Array.isArray(message.content)) {
for (const item of message.content) {
if (item.type === "text") {
tokens += countTokens(encoding, item.text);
}
else if (item.type === "image_url" && item.image_url) {
tokens += await countImageTokens(item, chatModel);
}
}
}
// OpenAI bug
if (message.name &&
message.role !== "tool") {
tokens += countTokens(encoding, message.name) + 1; // +1 for the name
}
if (message.role === "assistant" && message.tool_calls) {
tokens += 2;
for (const toolCall of message.tool_calls) {
tokens += 3;
tokens += countTokens(encoding, toolCall.type);
if (toolCall.type === "function") {
if ((_a = toolCall.function) === null || _a === void 0 ? void 0 : _a.name) {
const nameToken = countTokens(encoding, toolCall.function.name);
tokens += nameToken * 2;
}
if (toolCall.function.arguments && toolCall.function.arguments !== "{}") {
// console.log(JSON.stringify(formatArguments(toolCall.function.arguments)));
tokens += countTokens(encoding, (0, argument_format_1.formatArguments)(toolCall.function.arguments));
}
}
}
if (message.tool_calls.length > 1) {
tokens += 15; // s1, add delta when multi tools is added
tokens -= message.tool_calls.length * 5 - 6; // s2
}
else {
tokens -= 2; // s1, s2
}
}
if (message.role === "tool") {
tokens += 2; // add 2 if role is "tool"
}
else {
tokens += 3; // Add three per message
}
messageTokenCache.set(cacheKey, tokens);
return tokens;
}
function countTokens(encoding, text) {
if (!text)
return 0;
return encoding.encode(text).length;
}
async function countImageTokens(contentPart, chatModel) {
var _a, _b;
if (((_a = contentPart.image_url) === null || _a === void 0 ? void 0 : _a.detail) === 'low') {
return (0, vision_constants_1.getFixedPrice)(chatModel);
}
const cacheKey = HASHKEY(chatModel, stringifyImage(contentPart.image_url));
if (imageTokenCache.has(cacheKey)) {
return imageTokenCache.get(cacheKey);
}
const { width, height } = await getImageSize((_b = contentPart.image_url) === null || _b === void 0 ? void 0 : _b.url);
const longSide = Math.max(width, height);
const scaleFactor1 = (longSide > vision_constants_1.longSideLimit) ? longSide / vision_constants_1.longSideLimit : 1;
const shortSide = Math.min(width, height);
const scaleFactor2 = (shortSide / scaleFactor1 > vision_constants_1.shortSideLimit)
? (shortSide / scaleFactor1) / vision_constants_1.shortSideLimit : 1;
const scaleFactor = scaleFactor1 * scaleFactor2;
const scaledWidth = Math.floor(width / scaleFactor);
const scaledHeight = Math.floor(height / scaleFactor);
const tilesCount = Math.ceil(scaledWidth / vision_constants_1.tileSize) * Math.ceil(scaledHeight / vision_constants_1.tileSize);
const tokens = (0, vision_constants_1.getFixedPrice)(chatModel) + tilesCount * (0, vision_constants_1.getTilePrice)(chatModel);
imageTokenCache.set(cacheKey, tokens);
return tokens;
}
async function getImageSize(url) {
let imageBuffer;
if (url.startsWith('data:')) {
const uri = url.split(';base64,').pop();
imageBuffer = Buffer.from(uri, 'base64');
}
if (url.startsWith('https:') || url.startsWith('http:')) {
const response = await fetch(url);
const arrayBuffer = await response.arrayBuffer();
imageBuffer = Buffer.from(arrayBuffer);
}
if (!imageBuffer) {
throw new Error('imageBuffer is not defined');
}
const image = (0, sharp_1.default)(imageBuffer);
const metadata = await image.metadata();
const { width, height } = metadata;
if (!width || !height) {
throw new Error('unprocessable image - no size available');
}
return { width, height };
}
//# sourceMappingURL=index.js.map