codesnap-analyzer/dist/utils/token-counter.js

Create comprehensive snapshots of your codebase with token counting for LLMs

"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.TokenCounter = void 0;
// src/_core/helper/codesnap/utils/token-counter.ts
const tiktoken_1 = require("@dqbd/tiktoken");
const pricing_1 = require("../constants/pricing");
class TokenCounter {
    // Initialize tiktoken encoders
    static async getTokenizer(model) {
        try {
            return (0, tiktoken_1.encoding_for_model)(model);
        }
        catch (error) {
            console.error(`Error initializing tokenizer: ${error}`);
            return (0, tiktoken_1.get_encoding)("cl100k_base"); // fallback to base encoding
        }
    }
    // GPT-3.5 tokens using tiktoken
    static async countGPT35Tokens(text) {
        const encoder = await this.getTokenizer("gpt-3.5-turbo");
        return encoder.encode(text).length;
    }
    // GPT-4 tokens using tiktoken
    static async countGPT4Tokens(text) {
        const encoder = await this.getTokenizer("gpt-4");
        return encoder.encode(text).length;
    }
    // Claude tokens using cl100k_base encoding
    static async countClaudeTokens(text) {
        const encoder = await (0, tiktoken_1.get_encoding)("cl100k_base");
        return encoder.encode(text).length;
    }
    // LLaMA 2 tokens (approximation using cl100k_base)
    static async countLlama2Tokens(text) {
        const encoder = await (0, tiktoken_1.get_encoding)("cl100k_base");
        return Math.ceil(encoder.encode(text).length * 1.1); // 10% margin for differences
    }
    // Count tokens for all models
    static async countTokens(text) {
        // Remove unwanted tokens
        const sanitizedText = text
            .replace(/<\|fim_prefix\|>/g, "")
            .replace(/<\|fim_middle\|>/g, "")
            .replace(/<\|fim_suffix\|>/g, "")
            .replace(/<\|endofprompt\|>/g, "")
            .replace(/<\|endoftext\|>/g, "");
        const [gpt35, gpt4, claude, llama2] = await Promise.all([
            this.countGPT35Tokens(sanitizedText),
            this.countGPT4Tokens(sanitizedText),
            this.countClaudeTokens(sanitizedText),
            this.countLlama2Tokens(sanitizedText),
        ]);
        return {
            gpt35,
            gpt4,
            claude,
            llama2,
        };
    }
    // Calculate token cost
    static calculateTokenCost(counts) {
        return {
            gpt35: `$${((counts.gpt35 / 1000) * pricing_1.PRICING.gpt35).toFixed(4)}`,
            gpt4: `$${((counts.gpt4 / 1000) * pricing_1.PRICING.gpt4).toFixed(4)}`,
            claude: `$${((counts.claude / 1000) * pricing_1.PRICING.claude).toFixed(4)}`,
            llama2: `$${((counts.llama2 / 1000) * pricing_1.PRICING.llama2).toFixed(4)}`,
        };
    }
    // Format token counts and costs in a readable way
    static formatTokenCounts(counts) {
        const costs = this.calculateTokenCost(counts);
        return [
            "Token counts and costs by model:",
            `   GPT-3.5: ${this.formatNumber(counts.gpt35)} tokens → ${costs.gpt35}`,
            `   GPT-4:   ${this.formatNumber(counts.gpt4)} tokens → ${costs.gpt4}`,
            `   Claude:  ${this.formatNumber(counts.claude)} tokens → ${costs.claude}`,
            `   LLaMA 2: ${this.formatNumber(counts.llama2)} tokens → ${costs.llama2}`,
        ].join("\n");
    }
    static formatNumber(num) {
        if (num >= 1000000) {
            return `${(num / 1000000).toFixed(2)}M`;
        }
        if (num >= 1000) {
            return `${(num / 1000).toFixed(1)}K`;
        }
        return num.toString();
    }
}
exports.TokenCounter = TokenCounter;