const { Tiktoken } = require('tiktoken/lite');
const { load } = require('tiktoken/load');
const registry = require('tiktoken/registry.json');
const models = require('tiktoken/model_to_encoding.json');

const countTokens = async (text = '', modelName = 'gpt-3.5-turbo') => {
  let encoder = null;
  try {
    const model = await load(registry[models[modelName]]);
    encoder = new Tiktoken(model.bpe_ranks, model.special_tokens, model.pat_str);
    const tokens = encoder.encode(text);
    encoder.free();
    return tokens.length;
  } catch (e) {
    console.error(e);
    if (encoder) {
      encoder.free();
    }
    return 0;
  }
};

module.exports = countTokens;