@loxia-labs/loxia-autopilot-one 1.0.1

package/src/services/tokenCountingService.js

@@ -0,0 +1,536 @@
+/**
+ * TokenCountingService - Accurate token counting for conversation compactization
+ *
+ * Purpose:
+ * - Provide accurate token counting using tiktoken (OpenAI's tokenizer)
+ * - Support multiple models with different tokenization schemes
+ * - Determine when conversation compactization should be triggered
+ * - Cache token counts for performance optimization
+ * - Provide fallback estimation when accurate counting unavailable
+ *
+ * Key Features:
+ * - Accurate token counting via tiktoken library
+ * - Model-specific tokenization
+ * - Intelligent caching with TTL
+ * - Fast estimation fallback
+ * - Context window management
+ * - Compaction trigger detection
+ */
+
+import { encoding_for_model } from 'tiktoken';
+import {
+  MODELS,
+  COMPACTION_CONFIG,
+  TOKEN_COUNTING_MODES,
+} from '../utilities/constants.js';
+
+class TokenCountingService {
+  constructor(logger) {
+    this.logger = logger;
+
+    // Token count cache for performance
+    this.tokenCache = new Map();
+
+    // Tiktoken encoders cache (reuse encoders for efficiency)
+    this.encoders = new Map();
+
+    // Supported tiktoken models mapping
+    this.tiktokenModelMap = {
+      // Anthropic models use OpenAI's cl100k_base encoding
+      [MODELS.ANTHROPIC_SONNET]: 'gpt-4',
+      [MODELS.ANTHROPIC_OPUS]: 'gpt-4',
+      [MODELS.ANTHROPIC_HAIKU]: 'gpt-4',
+
+      // OpenAI models
+      [MODELS.GPT_4]: 'gpt-4',
+      [MODELS.GPT_4_MINI]: 'gpt-4',
+      'gpt-4o': 'gpt-4',
+      'gpt-4o-mini': 'gpt-4',
+      'gpt-4-turbo': 'gpt-4',
+      'gpt-3.5-turbo': 'gpt-4', // Uses cl100k_base encoding
+
+      // DeepSeek uses similar tokenization to GPT-4
+      [MODELS.DEEPSEEK_R1]: 'gpt-4',
+
+      // Phi models - fallback to GPT-4 encoding
+      [MODELS.PHI_4]: 'gpt-4',
+
+      // Azure AI Foundry models
+      'azure-ai-grok3': 'gpt-4',
+      'azure-ai-deepseek-r1': 'gpt-4',
+      'azure-openai-gpt-5': 'gpt-4',
+      'azure-openai-gpt-4': 'gpt-4',
+      'azure-openai-gpt-4o': 'gpt-4',
+
+      // Compaction models
+      [COMPACTION_CONFIG.COMPACTION_MODEL]: 'gpt-4',
+      [COMPACTION_CONFIG.COMPACTION_MODEL_FALLBACK]: 'gpt-4',
+    };
+
+    this.logger?.info('TokenCountingService initialized', {
+      supportedModels: Object.keys(this.tiktokenModelMap).length,
+      cacheEnabled: true
+    });
+  }
+
+  /**
+   * Count tokens in text using accurate tiktoken encoder
+   * @param {string} text - Text to count tokens in
+   * @param {string} model - Model name for appropriate tokenization
+   * @param {string} mode - Counting mode (accurate, estimated, cached)
+   * @returns {Promise<number>} Token count
+   */
+  async countTokens(text, model, mode = TOKEN_COUNTING_MODES.ACCURATE) {
+    // Validate input
+    if (!text || typeof text !== 'string') {
+      return 0;
+    }
+
+    // Check cache first if enabled
+    if (mode === TOKEN_COUNTING_MODES.CACHED) {
+      const cached = this._getCachedTokenCount(text, model);
+      if (cached !== null) {
+        return cached;
+      }
+      // Fall through to accurate counting if not cached
+    }
+
+    // Use fast estimation if requested
+    if (mode === TOKEN_COUNTING_MODES.ESTIMATED) {
+      return this._estimateTokens(text);
+    }
+
+    // Accurate counting with tiktoken
+    try {
+      const encoder = await this._getEncoder(model);
+      const tokens = encoder.encode(text);
+      const count = tokens.length;
+
+      // Cache the result
+      this._cacheTokenCount(text, model, count);
+
+      return count;
+
+    } catch (error) {
+      this.logger?.warn('Tiktoken encoding failed, falling back to estimation', {
+        model,
+        error: error.message
+      });
+
+      // Fallback to estimation
+      return this._estimateTokens(text);
+    }
+  }
+
+  /**
+   * Estimate conversation token count including all messages
+   * @param {Array} messages - Array of message objects with { role, content }
+   * @param {string} model - Model name for tokenization
+   * @param {string} mode - Counting mode
+   * @returns {Promise<number>} Total token count for conversation
+   */
+  async estimateConversationTokens(messages, model, mode = TOKEN_COUNTING_MODES.ACCURATE) {
+    if (!Array.isArray(messages) || messages.length === 0) {
+      return 0;
+    }
+
+    let totalTokens = 0;
+
+    // Count tokens for each message
+    for (const message of messages) {
+      if (!message.content) {
+        continue;
+      }
+
+      // Message formatting overhead (role + formatting)
+      const formattingOverhead = this._getMessageFormattingOverhead(model);
+
+      // Content tokens
+      const contentTokens = await this.countTokens(message.content, model, mode);
+
+      totalTokens += contentTokens + formattingOverhead;
+    }
+
+    // Add conversation-level overhead (system instructions, etc.)
+    const conversationOverhead = this._getConversationOverhead(messages.length);
+    totalTokens += conversationOverhead;
+
+    this.logger?.debug('Conversation token count', {
+      model,
+      messageCount: messages.length,
+      totalTokens,
+      mode
+    });
+
+    return totalTokens;
+  }
+
+  /**
+   * Get context window size for a model
+   * @param {string} model - Model name
+   * @returns {number} Context window size in tokens
+   */
+  getModelContextWindow(model) {
+    // Model context windows (from aiService model specs and vendor documentation)
+    const contextWindows = {
+      // Anthropic Claude models
+      [MODELS.ANTHROPIC_SONNET]: 200000,
+      [MODELS.ANTHROPIC_OPUS]: 200000,
+      [MODELS.ANTHROPIC_HAIKU]: 200000,
+
+      // OpenAI models
+      [MODELS.GPT_4]: 128000,
+      [MODELS.GPT_4_MINI]: 128000,
+      'gpt-4o': 128000,
+      'gpt-4o-mini': 128000,
+      'gpt-4-turbo': 128000,
+      'gpt-3.5-turbo': 16384,
+
+      // DeepSeek models
+      [MODELS.DEEPSEEK_R1]: 128000,
+
+      // Phi models
+      [MODELS.PHI_4]: 16384,
+
+      // Azure AI Foundry models
+      'azure-ai-grok3': 128000,
+      'azure-ai-deepseek-r1': 128000,
+      'azure-openai-gpt-5': 128000,
+      'azure-openai-gpt-4': 128000,
+      'azure-openai-gpt-4o': 128000,
+
+      // Compaction models
+      [COMPACTION_CONFIG.COMPACTION_MODEL]: 128000,
+      [COMPACTION_CONFIG.COMPACTION_MODEL_FALLBACK]: 128000,
+
+      // Router model
+      'autopilot-model-router': 16384,
+    };
+
+    const contextWindow = contextWindows[model];
+
+    if (!contextWindow) {
+      this.logger?.warn('Unknown model context window, using default', {
+        model,
+        defaultWindow: 128000
+      });
+      return 128000; // Default to 128k
+    }
+
+    return contextWindow;
+  }
+
+  /**
+   * Get maximum output tokens for a model
+   * @param {string} model - Model name
+   * @returns {number} Maximum output tokens
+   */
+  getModelMaxOutputTokens(model) {
+    // Max output tokens - increased to 8K-20K where supported for better responses
+    // Note: These are conservative estimates to ensure compaction triggers appropriately
+    const maxOutputTokens = {
+      // Anthropic Claude models - support up to 8K output
+      [MODELS.ANTHROPIC_SONNET]: 8192,
+      [MODELS.ANTHROPIC_OPUS]: 8192,
+      [MODELS.ANTHROPIC_HAIKU]: 8192,
+
+      // OpenAI models
+      [MODELS.GPT_4]: 8192,           // Supports up to 16K, using 8K for safety
+      [MODELS.GPT_4_MINI]: 16384,     // Already at max
+      'gpt-4o': 8192,                 // Supports up to 16K, using 8K
+      'gpt-4o-mini': 16384,           // Already at max
+      'gpt-4-turbo': 8192,            // Increased from default
+      'gpt-3.5-turbo': 4096,          // Smaller model, keep at 4K
+
+      // DeepSeek models
+      [MODELS.DEEPSEEK_R1]: 8192,     // Already at max
+
+      // Phi models - smaller architecture
+      [MODELS.PHI_4]: 4096,           // Increased from 2048
+
+      // Azure AI Foundry models
+      'azure-ai-grok3': 8192,         // Increased from 4K
+      'azure-ai-deepseek-r1': 8192,   // Already at max
+      'azure-openai-gpt-5': 8192,     // Increased from 4K
+      'azure-openai-gpt-4': 8192,     // Increased from default
+      'azure-openai-gpt-4o': 8192,    // Increased from default
+
+      // Compaction models - keep moderate for efficiency
+      [COMPACTION_CONFIG.COMPACTION_MODEL]: 8192,
+      [COMPACTION_CONFIG.COMPACTION_MODEL_FALLBACK]: 8192,
+
+      // Router model - keep small for fast routing
+      'autopilot-model-router': 2048,
+    };
+
+    return maxOutputTokens[model] || 8192; // Default increased to 8K
+  }
+
+  /**
+   * Determine if compaction should be triggered
+   * @param {number} currentTokens - Current conversation token count (K)
+   * @param {number} maxOutputTokens - Max tokens model can output (X)
+   * @param {number} contextWindow - Model's context window size (C)
+   * @param {number} threshold - Trigger threshold (default 0.8 = 80%)
+   * @returns {boolean} True if compaction should be triggered
+   */
+  shouldTriggerCompaction(currentTokens, maxOutputTokens, contextWindow, threshold = COMPACTION_CONFIG.DEFAULT_THRESHOLD) {
+    // Validate threshold
+    if (threshold < COMPACTION_CONFIG.MIN_THRESHOLD || threshold > COMPACTION_CONFIG.MAX_THRESHOLD) {
+      this.logger?.warn('Invalid compaction threshold, using default', {
+        provided: threshold,
+        default: COMPACTION_CONFIG.DEFAULT_THRESHOLD
+      });
+      threshold = COMPACTION_CONFIG.DEFAULT_THRESHOLD;
+    }
+
+    // Calculate: K + X >= threshold * C
+    const requiredTokens = currentTokens + maxOutputTokens;
+    const thresholdTokens = threshold * contextWindow;
+    const shouldTrigger = requiredTokens >= thresholdTokens;
+
+    this.logger?.debug('Compaction trigger check', {
+      currentTokens,
+      maxOutputTokens,
+      contextWindow,
+      threshold,
+      requiredTokens,
+      thresholdTokens,
+      shouldTrigger,
+      utilizationPercent: ((requiredTokens / contextWindow) * 100).toFixed(2)
+    });
+
+    return shouldTrigger;
+  }
+
+  /**
+   * Calculate how many tokens to target after compaction
+   * @param {number} contextWindow - Model's context window size
+   * @param {number} targetThreshold - Target threshold after compaction (default 85%)
+   * @returns {number} Target token count after compaction
+   */
+  calculateTargetTokenCount(contextWindow, targetThreshold = COMPACTION_CONFIG.MAX_ACCEPTABLE_TOKEN_COUNT_AFTER) {
+    return Math.floor(contextWindow * targetThreshold);
+  }
+
+  /**
+   * Validate that compaction achieved sufficient reduction
+   * @param {number} originalTokens - Token count before compaction
+   * @param {number} compactedTokens - Token count after compaction
+   * @param {number} contextWindow - Model's context window
+   * @returns {Object} Validation result { valid, reductionPercent, exceedsTarget }
+   */
+  validateCompaction(originalTokens, compactedTokens, contextWindow) {
+    const reductionPercent = ((originalTokens - compactedTokens) / originalTokens) * 100;
+    const targetTokens = this.calculateTargetTokenCount(contextWindow);
+    const exceedsTarget = compactedTokens > targetTokens;
+    const sufficientReduction = reductionPercent >= COMPACTION_CONFIG.MIN_REDUCTION_PERCENTAGE;
+
+    const valid = !exceedsTarget && sufficientReduction;
+
+    this.logger?.info('Compaction validation', {
+      originalTokens,
+      compactedTokens,
+      reductionPercent: reductionPercent.toFixed(2),
+      targetTokens,
+      exceedsTarget,
+      sufficientReduction,
+      valid
+    });
+
+    return {
+      valid,
+      reductionPercent,
+      exceedsTarget,
+      sufficientReduction,
+      targetTokens,
+      compactedTokens,
+      originalTokens
+    };
+  }
+
+  /**
+   * Clear token count cache
+   * @param {string} model - Optional: clear cache for specific model only
+   */
+  clearCache(model = null) {
+    if (model) {
+      // Clear cache entries for specific model
+      for (const [key, value] of this.tokenCache.entries()) {
+        if (value.model === model) {
+          this.tokenCache.delete(key);
+        }
+      }
+      this.logger?.debug('Token cache cleared for model', { model });
+    } else {
+      // Clear entire cache
+      this.tokenCache.clear();
+      this.logger?.debug('Token cache cleared completely');
+    }
+  }
+
+  /**
+   * Get encoder for model (with caching)
+   * @private
+   */
+  async _getEncoder(model) {
+    // Get tiktoken model name
+    const tiktokenModel = this.tiktokenModelMap[model] || 'gpt-4';
+
+    // Check if encoder is already cached
+    if (this.encoders.has(tiktokenModel)) {
+      return this.encoders.get(tiktokenModel);
+    }
+
+    // Create new encoder
+    const encoder = encoding_for_model(tiktokenModel);
+    this.encoders.set(tiktokenModel, encoder);
+
+    this.logger?.debug('Created tiktoken encoder', {
+      model,
+      tiktokenModel
+    });
+
+    return encoder;
+  }
+
+  /**
+   * Fast token estimation (character-based)
+   * @private
+   */
+  _estimateTokens(text) {
+    if (!text || typeof text !== 'string') {
+      return 0;
+    }
+
+    // Use configured estimation ratio
+    return Math.ceil(text.length / COMPACTION_CONFIG.CHARS_PER_TOKEN_ESTIMATE);
+  }
+
+  /**
+   * Get cached token count if available and not expired
+   * @private
+   */
+  _getCachedTokenCount(text, model) {
+    const cacheKey = this._getCacheKey(text, model);
+    const cached = this.tokenCache.get(cacheKey);
+
+    if (!cached) {
+      return null;
+    }
+
+    // Check if cache entry is expired
+    const now = Date.now();
+    if (now - cached.timestamp > COMPACTION_CONFIG.TOKEN_COUNT_CACHE_TTL_MS) {
+      this.tokenCache.delete(cacheKey);
+      return null;
+    }
+
+    return cached.count;
+  }
+
+  /**
+   * Cache token count with timestamp
+   * @private
+   */
+  _cacheTokenCount(text, model, count) {
+    const cacheKey = this._getCacheKey(text, model);
+
+    this.tokenCache.set(cacheKey, {
+      count,
+      model,
+      timestamp: Date.now()
+    });
+
+    // Prevent cache from growing indefinitely
+    if (this.tokenCache.size > 1000) {
+      // Remove oldest entries
+      const entries = Array.from(this.tokenCache.entries());
+      entries.sort((a, b) => a[1].timestamp - b[1].timestamp);
+
+      // Remove oldest 20%
+      const toRemove = Math.floor(entries.length * 0.2);
+      for (let i = 0; i < toRemove; i++) {
+        this.tokenCache.delete(entries[i][0]);
+      }
+    }
+  }
+
+  /**
+   * Generate cache key for text + model
+   * @private
+   */
+  _getCacheKey(text, model) {
+    // Use a simple hash for the text to avoid storing full text as key
+    const hash = this._simpleHash(text);
+    return `${model}:${hash}`;
+  }
+
+  /**
+   * Simple hash function for cache keys
+   * @private
+   */
+  _simpleHash(str) {
+    let hash = 0;
+    for (let i = 0; i < str.length; i++) {
+      const char = str.charCodeAt(i);
+      hash = ((hash << 5) - hash) + char;
+      hash = hash & hash; // Convert to 32-bit integer
+    }
+    return hash.toString(36);
+  }
+
+  /**
+   * Get message formatting overhead for model
+   * Accounts for role labels, XML tags, etc.
+   * @private
+   */
+  _getMessageFormattingOverhead(model) {
+    // Different models have different formatting
+    // Anthropic: ~10 tokens per message (role tags)
+    // OpenAI: ~5 tokens per message (JSON formatting)
+
+    if (model.includes('anthropic') || model.includes('claude')) {
+      return 10;
+    }
+
+    return 5; // Default for OpenAI-style models
+  }
+
+  /**
+   * Get conversation-level overhead
+   * Accounts for system prompts, special tokens, etc.
+   * @private
+   */
+  _getConversationOverhead(messageCount) {
+    // Base overhead for conversation structure
+    const baseOverhead = 50;
+
+    // Additional overhead scales with message count
+    const scalingOverhead = messageCount * 2;
+
+    return baseOverhead + scalingOverhead;
+  }
+
+  /**
+   * Clean up resources (close encoders)
+   */
+  async cleanup() {
+    // Free tiktoken encoders
+    for (const [model, encoder] of this.encoders.entries()) {
+      try {
+        encoder.free();
+        this.logger?.debug('Freed tiktoken encoder', { model });
+      } catch (error) {
+        this.logger?.warn('Failed to free encoder', { model, error: error.message });
+      }
+    }
+
+    this.encoders.clear();
+    this.tokenCache.clear();
+
+    this.logger?.info('TokenCountingService cleaned up');
+  }
+}
+
+export default TokenCountingService;