@hashgraphonline/hedera-agent-kit/src/utils/token-usage-tracker.ts

Build LLM-powered applications that interact with the Hedera Network. Create conversational agents that can understand user requests in natural language and execute Hedera transactions, or build backend systems that leverage AI for on-chain operations.

import { BaseCallbackHandler } from '@langchain/core/callbacks/base';
import { LLMResult } from '@langchain/core/outputs';
import { Logger } from '@hashgraphonline/standards-sdk';

/**
 * Token usage data structure
 */
export interface TokenUsage {
  promptTokens: number;
  completionTokens: number;
  totalTokens: number;
  modelName?: string | undefined;
  timestamp?: Date | undefined;
}

/**
 * Cost calculation result
 */
export interface CostCalculation {
  promptCost: number;
  completionCost: number;
  totalCost: number;
  currency: string;
}

/**
 * OpenRouter model pricing structure
 */
export interface OpenRouterModel {
  id: string;
  name: string;
  pricing: {
    prompt: string;
    completion: string;
  };
}

/**
 * Model pricing manager that fetches and caches pricing from OpenRouter API
 */
export class ModelPricingManager {
  private static instance: ModelPricingManager;
  private pricingCache: Map<string, { prompt: number; completion: number }> =
    new Map();
  private lastFetchTime: number = 0;
  private readonly CACHE_DURATION = 24 * 60 * 60 * 1000;
  private readonly OPENROUTER_API_URL = 'https://openrouter.ai/api/v1/models';
  private logger: Logger;

  private constructor() {
    this.logger = new Logger({ module: 'ModelPricingManager', level: 'info' });
    this.initializeFallbackPricing();
  }

  public static getInstance(): ModelPricingManager {
    if (!ModelPricingManager.instance) {
      ModelPricingManager.instance = new ModelPricingManager();
    }
    return ModelPricingManager.instance;
  }

  private initializeFallbackPricing(): void {
    const fallbackPricing = {
      'gpt-4o': { prompt: 0.005, completion: 0.015 },
      'gpt-4o-mini': { prompt: 0.00015, completion: 0.0006 },
      'gpt-4-turbo': { prompt: 0.01, completion: 0.03 },
      'gpt-4': { prompt: 0.03, completion: 0.06 },
      'gpt-3.5-turbo': { prompt: 0.0005, completion: 0.0015 },
    };

    for (const [model, pricing] of Object.entries(fallbackPricing)) {
      this.pricingCache.set(model, pricing);
    }
  }

  private async fetchPricingFromAPI(): Promise<void> {
    try {
      const response = await fetch(this.OPENROUTER_API_URL, {
        headers: {
          Accept: 'application/json',
          'User-Agent': 'hedera-agent-kit/1.0',
        },
      });

      if (!response.ok) {
        throw new Error(`HTTP ${response.status}: ${response.statusText}`);
      }

      const data = await response.json();
      const models: OpenRouterModel[] = data.data || [];

      for (const model of models) {
        if (model.pricing?.prompt && model.pricing?.completion) {
          const promptPrice = parseFloat(model.pricing.prompt);
          const completionPrice = parseFloat(model.pricing.completion);

          if (!isNaN(promptPrice) && !isNaN(completionPrice)) {
            this.pricingCache.set(model.id, {
              prompt: promptPrice * 1000,
              completion: completionPrice * 1000,
            });

            const simplifiedId = this.getSimplifiedModelId(model.id);
            if (simplifiedId !== model.id) {
              this.pricingCache.set(simplifiedId, {
                prompt: promptPrice * 1000,
                completion: completionPrice * 1000,
              });
            }
          }
        }
      }

      this.lastFetchTime = Date.now();
      this.logger.info(
        `Loaded pricing for ${models.length} models from OpenRouter API`
      );
    } catch (error) {
      this.logger.warn(
        'Failed to fetch pricing from API, using fallback pricing',
        error
      );
    }
  }

  private getSimplifiedModelId(fullId: string): string {
    const prefixes = [
      'openai/',
      'anthropic/',
      'google/',
      'mistralai/',
      'meta-llama/',
    ];
    for (const prefix of prefixes) {
      if (fullId.startsWith(prefix)) {
        return fullId.replace(prefix, '');
      }
    }
    return fullId;
  }

  private async ensurePricingLoaded(): Promise<void> {
    const now = Date.now();
    const shouldRefresh = now - this.lastFetchTime > this.CACHE_DURATION;

    if (shouldRefresh || this.pricingCache.size < 10) {
      await this.fetchPricingFromAPI();
    }
  }

  public async getPricing(
    modelName: string
  ): Promise<{ prompt: number; completion: number }> {
    await this.ensurePricingLoaded();

    const exactMatch = this.pricingCache.get(modelName);
    if (exactMatch) {
      return exactMatch;
    }

    const normalizedName = modelName.toLowerCase();
    for (const [key, value] of this.pricingCache.entries()) {
      if (
        key.toLowerCase().includes(normalizedName) ||
        normalizedName.includes(key.toLowerCase())
      ) {
        return value;
      }
    }

    return (
      this.pricingCache.get('gpt-4o-mini') || {
        prompt: 0.00015,
        completion: 0.0006,
      }
    );
  }

  public getPricingSync(modelName: string): {
    prompt: number;
    completion: number;
  } {
    const exactMatch = this.pricingCache.get(modelName);
    if (exactMatch) {
      return exactMatch;
    }

    const normalizedName = modelName.toLowerCase();
    for (const [key, value] of this.pricingCache.entries()) {
      if (
        key.toLowerCase().includes(normalizedName) ||
        normalizedName.includes(key.toLowerCase())
      ) {
        return value;
      }
    }

    return (
      this.pricingCache.get('gpt-4o-mini') || {
        prompt: 0.00015,
        completion: 0.0006,
      }
    );
  }
}

/**
 * Callback handler to track token usage from OpenAI API responses
 */
export class TokenUsageCallbackHandler extends BaseCallbackHandler {
  name = 'TokenUsageCallbackHandler';
  private tokenUsageHistory: TokenUsage[] = [];
  private currentTokenUsage: TokenUsage | undefined = undefined;
  private logger: Logger;
  private modelName?: string | undefined;

  constructor(modelName?: string, logger?: Logger) {
    super();
    this.modelName = modelName;
    this.logger =
      logger || new Logger({ module: 'TokenUsageTracker', level: 'info' });
  }

  override async handleLLMEnd(output: LLMResult): Promise<void> {
    try {
      if (output.llmOutput?.tokenUsage) {
        this.currentTokenUsage = {
          promptTokens: output.llmOutput.tokenUsage.promptTokens || 0,
          completionTokens: output.llmOutput.tokenUsage.completionTokens || 0,
          totalTokens: output.llmOutput.tokenUsage.totalTokens || 0,
          modelName: this.modelName,
          timestamp: new Date(),
        };

        this.tokenUsageHistory.push(this.currentTokenUsage);

        this.logger.debug('Token usage tracked:', {
          promptTokens: this.currentTokenUsage.promptTokens,
          completionTokens: this.currentTokenUsage.completionTokens,
          totalTokens: this.currentTokenUsage.totalTokens,
          model: this.modelName,
        });
      }
    } catch (error) {
      this.logger.error('Failed to track token usage:', error);
    }
  }

  getLatestTokenUsage(): TokenUsage | undefined {
    return this.currentTokenUsage;
  }

  getTokenUsageHistory(): TokenUsage[] {
    return [...this.tokenUsageHistory];
  }

  getTotalTokenUsage(): TokenUsage {
    const total = this.tokenUsageHistory.reduce(
      (acc, usage) => ({
        promptTokens: acc.promptTokens + usage.promptTokens,
        completionTokens: acc.completionTokens + usage.completionTokens,
        totalTokens: acc.totalTokens + usage.totalTokens,
      }),
      { promptTokens: 0, completionTokens: 0, totalTokens: 0 }
    );

    return {
      ...total,
      modelName: this.modelName,
      timestamp: new Date(),
    };
  }

  reset(): void {
    this.currentTokenUsage = undefined;
    this.tokenUsageHistory = [];
  }
}

/**
 * Calculate cost based on token usage and model with dynamic pricing
 */
export async function calculateTokenCost(
  tokenUsage: TokenUsage,
  modelName?: string
): Promise<CostCalculation> {
  const model = modelName || tokenUsage.modelName || 'gpt-4o-mini';
  const pricingManager = ModelPricingManager.getInstance();
  const pricing = await pricingManager.getPricing(model);

  const promptCost = (tokenUsage.promptTokens / 1000) * pricing.prompt;
  const completionCost =
    (tokenUsage.completionTokens / 1000) * pricing.completion;

  return {
    promptCost,
    completionCost,
    totalCost: promptCost + completionCost,
    currency: 'USD',
  };
}

/**
 * Synchronous version of calculateTokenCost using cached pricing
 */
export function calculateTokenCostSync(
  tokenUsage: TokenUsage,
  modelName?: string
): CostCalculation {
  const model = modelName || tokenUsage.modelName || 'gpt-4o-mini';
  const pricingManager = ModelPricingManager.getInstance();
  const pricing = pricingManager.getPricingSync(model);

  const promptCost = (tokenUsage.promptTokens / 1000) * pricing.prompt;
  const completionCost =
    (tokenUsage.completionTokens / 1000) * pricing.completion;

  return {
    promptCost,
    completionCost,
    totalCost: promptCost + completionCost,
    currency: 'USD',
  };
}

/**
 * Format cost for display
 */
export function formatCost(
  cost: CostCalculation,
  precision: number = 6
): string {
  return `$${cost.totalCost.toFixed(precision)} ${cost.currency}`;
}

/**
 * Estimate tokens from text (rough approximation)
 */
export function estimateTokens(text: string): number {
  return Math.ceil(text.length / 4);
}