@allternit/gizzi-code 1.0.0

package/src/runtime/services/tokenEstimation.ts

@@ -0,0 +1,547 @@
+/**
+ * @fileoverview Token estimation utilities for LLM interactions
+ * @module runtime/services/tokenEstimation
+ *
+ * Provides token counting and estimation functions for various LLM models.
+ * Supports both accurate encoding (when tiktoken is available) and
+ * rough estimation based on character counts.
+ */
+
+import type { Message } from '../../types/message.js';
+
+/**
+ * Supported LLM models for token estimation
+ * @typedef {'gpt-4' | 'gpt-4o' | 'gpt-4-turbo' | 'gpt-3.5-turbo' | 'claude-3-opus' | 'claude-3-sonnet' | 'claude-3-haiku' | 'claude-2' | 'gemini-pro' | 'gemini-ultra' | 'default'} SupportedModel
+ */
+export type SupportedModel =
+  | 'gpt-4'
+  | 'gpt-4o'
+  | 'gpt-4-turbo'
+  | 'gpt-3.5-turbo'
+  | 'claude-3-opus'
+  | 'claude-3-sonnet'
+  | 'claude-3-haiku'
+  | 'claude-2'
+  | 'gemini-pro'
+  | 'gemini-ultra'
+  | 'default';
+
+/**
+ * Token estimation configuration options
+ * @interface TokenEstimationOptions
+ */
+export interface TokenEstimationOptions {
+  /** Target model for estimation */
+  readonly model?: SupportedModel;
+
+  /** Whether to include formatting overhead in estimate */
+  readonly includeFormatting?: boolean;
+
+  /** Custom tokens per character ratio for rough estimation */
+  readonly customRatio?: number;
+}
+
+/**
+ * Token estimation result with metadata
+ * @interface TokenEstimationResult
+ */
+export interface TokenEstimationResult {
+  /** Estimated token count */
+  readonly tokens: number;
+
+  /** Method used for estimation */
+  readonly method: 'precise' | 'rough';
+
+  /** Model used for estimation */
+  readonly model: SupportedModel;
+
+  /** Confidence level (0-1) based on estimation method */
+  readonly confidence: number;
+}
+
+/**
+ * Model-specific token estimation ratios (characters per token)
+ * These are approximate averages and vary by language/content
+ * @constant {Readonly<Record<SupportedModel, number>>}
+ */
+export const MODEL_TOKEN_RATIOS: Readonly<Record<SupportedModel, number>> = {
+  'gpt-4': 3.5,
+  'gpt-4o': 3.5,
+  'gpt-4-turbo': 3.5,
+  'gpt-3.5-turbo': 3.5,
+  'claude-3-opus': 3.25,
+  'claude-3-sonnet': 3.25,
+  'claude-3-haiku': 3.25,
+  'claude-2': 3.25,
+  'gemini-pro': 3.5,
+  'gemini-ultra': 3.5,
+  default: 3.5,
+};
+
+/**
+ * Message format overhead in tokens (role markers, formatting, etc.)
+ * @constant {number}
+ */
+export const MESSAGE_OVERHEAD_TOKENS = 4;
+
+/**
+ * Base overhead for a conversation (system prompts, etc.)
+ * @constant {number}
+ */
+export const CONVERSATION_BASE_OVERHEAD = 3;
+
+/**
+ * Cache for tiktoken encoder to avoid re-initialization
+ * @type {unknown | null}
+ */
+let tiktokenEncoder: unknown | null = null;
+
+/**
+ * Try to load tiktoken for precise token counting
+ * @returns {Promise<unknown | null>} Tiktoken encoder or null if not available
+ */
+async function getTiktokenEncoder(): Promise<unknown | null> {
+  if (tiktokenEncoder) {
+    return tiktokenEncoder;
+  }
+
+  try {
+    // Dynamic import to avoid hard dependency
+    const tiktoken = await import('tiktoken');
+    // Use cl100k_base encoding (used by GPT-4 and GPT-3.5-turbo)
+    tiktokenEncoder = tiktoken.get_encoding('cl100k_base');
+    return tiktokenEncoder;
+  } catch {
+    // tiktoken not available, fall back to rough estimation
+    return null;
+  }
+}
+
+/**
+ * Estimate tokens in a single text string
+ *
+ * @param {string} text - Text to estimate tokens for
+ * @param {TokenEstimationOptions} options - Estimation options
+ * @returns {TokenEstimationResult} Token estimation with metadata
+ *
+ * @example
+ * ```typescript
+ * const result = estimateTokens('Hello, world!', { model: 'gpt-4' });
+ * console.log(result.tokens); // ~4 tokens
+ * ```
+ */
+export function estimateTokens(
+  text: string,
+  options: TokenEstimationOptions = {}
+): TokenEstimationResult {
+  const { model = 'default', customRatio } = options;
+
+  // Handle empty text
+  if (!text || text.length === 0) {
+    return {
+      tokens: 0,
+      method: 'precise',
+      model,
+      confidence: 1,
+    };
+  }
+
+  // Try synchronous tiktoken if available (pre-loaded)
+  if (tiktokenEncoder && typeof (tiktokenEncoder as { encode: (text: string) => unknown }).encode === 'function') {
+    try {
+      const tokens = (tiktokenEncoder as { encode: (text: string) => number[] }).encode(text);
+      return {
+        tokens: tokens.length,
+        method: 'precise',
+        model,
+        confidence: 0.95,
+      };
+    } catch {
+      // Fall through to rough estimation
+    }
+  }
+
+  // Rough estimation based on character count
+  const ratio = customRatio ?? MODEL_TOKEN_RATIOS[model] ?? MODEL_TOKEN_RATIOS.default;
+  const estimatedTokens = Math.ceil(text.length / ratio);
+
+  return {
+    tokens: estimatedTokens,
+    method: 'rough',
+    model,
+    confidence: 0.8,
+  };
+}
+
+/**
+ * Async version of estimateTokens that attempts to use tiktoken for precision
+ *
+ * @param {string} text - Text to estimate tokens for
+ * @param {TokenEstimationOptions} options - Estimation options
+ * @returns {Promise<TokenEstimationResult>} Token estimation with metadata
+ *
+ * @example
+ * ```typescript
+ * const result = await estimateTokensAsync('Hello, world!', { model: 'gpt-4' });
+ * console.log(result.method); // 'precise' if tiktoken available
+ * ```
+ */
+export async function estimateTokensAsync(
+  text: string,
+  options: TokenEstimationOptions = {}
+): Promise<TokenEstimationResult> {
+  const { model = 'default', customRatio } = options;
+
+  // Handle empty text
+  if (!text || text.length === 0) {
+    return {
+      tokens: 0,
+      method: 'precise',
+      model,
+      confidence: 1,
+    };
+  }
+
+  // Try to use tiktoken for precise counting
+  const encoder = await getTiktokenEncoder();
+  if (encoder && typeof (encoder as { encode: (text: string) => unknown }).encode === 'function') {
+    try {
+      const tokens = (encoder as { encode: (text: string) => number[] }).encode(text);
+      return {
+        tokens: tokens.length,
+        method: 'precise',
+        model,
+        confidence: 0.95,
+      };
+    } catch {
+      // Fall through to rough estimation
+    }
+  }
+
+  // Fall back to rough estimation
+  return estimateTokens(text, options);
+}
+
+/**
+ * Estimate tokens for an array of messages
+ *
+ * @param {readonly Message[]} messages - Messages to estimate tokens for
+ * @param {TokenEstimationOptions} options - Estimation options
+ * @returns {TokenEstimationResult} Total token estimation with metadata
+ *
+ * @example
+ * ```typescript
+ * const messages = [
+ *   { role: 'system', content: 'You are a helpful assistant.' },
+ *   { role: 'user', content: 'Hello!' }
+ * ];
+ * const result = estimateTokensForMessages(messages);
+ * console.log(result.tokens); // Total estimated tokens
+ * ```
+ */
+export function estimateTokensForMessages(
+  messages: readonly Message[],
+  options: TokenEstimationOptions = {}
+): TokenEstimationResult {
+  const { model = 'default', includeFormatting = true } = options;
+
+  if (!messages || messages.length === 0) {
+    return {
+      tokens: 0,
+      method: 'precise',
+      model,
+      confidence: 1,
+    };
+  }
+
+  let totalTokens = 0;
+  let preciseCount = 0;
+  let roughCount = 0;
+
+  for (const message of messages) {
+    // Estimate content tokens
+    const contentEstimate = estimateTokens(message.content, { model });
+    totalTokens += contentEstimate.tokens;
+
+    if (contentEstimate.method === 'precise') {
+      preciseCount++;
+    } else {
+      roughCount++;
+    }
+
+    // Add per-message overhead (formatting, role markers, etc.)
+    if (includeFormatting) {
+      totalTokens += MESSAGE_OVERHEAD_TOKENS;
+
+      // Add tokens for role name
+      totalTokens += estimateTokens(message.role, { model }).tokens;
+    }
+
+    // Add tokens for attachments if present
+    if (message.attachments && message.attachments.length > 0) {
+      for (const attachment of message.attachments) {
+        // Estimate based on attachment metadata
+        const attachmentText = `${attachment.name} ${attachment.mimeType} ${attachment.caption ?? ''}`;
+        totalTokens += estimateTokens(attachmentText, { model }).tokens;
+      }
+    }
+  }
+
+  // Add conversation base overhead
+  if (includeFormatting) {
+    totalTokens += CONVERSATION_BASE_OVERHEAD;
+  }
+
+  return {
+    tokens: totalTokens,
+    method: roughCount > 0 ? 'rough' : 'precise',
+    model,
+    confidence: preciseCount > roughCount ? 0.9 : 0.75,
+  };
+}
+
+/**
+ * Async version of estimateTokensForMessages with tiktoken support
+ *
+ * @param {readonly Message[]} messages - Messages to estimate tokens for
+ * @param {TokenEstimationOptions} options - Estimation options
+ * @returns {Promise<TokenEstimationResult>} Total token estimation with metadata
+ */
+export async function estimateTokensForMessagesAsync(
+  messages: readonly Message[],
+  options: TokenEstimationOptions = {}
+): Promise<TokenEstimationResult> {
+  const { model = 'default', includeFormatting = true } = options;
+
+  if (!messages || messages.length === 0) {
+    return {
+      tokens: 0,
+      method: 'precise',
+      model,
+      confidence: 1,
+    };
+  }
+
+  // Try to initialize tiktoken
+  const encoder = await getTiktokenEncoder();
+  const usePrecise = encoder !== null;
+
+  let totalTokens = 0;
+
+  for (const message of messages) {
+    // Use precise or rough estimation based on availability
+    if (usePrecise) {
+      const contentResult = await estimateTokensAsync(message.content, { model });
+      totalTokens += contentResult.tokens;
+    } else {
+      const contentResult = estimateTokens(message.content, { model });
+      totalTokens += contentResult.tokens;
+    }
+
+    // Add per-message overhead
+    if (includeFormatting) {
+      totalTokens += MESSAGE_OVERHEAD_TOKENS;
+      const roleResult = usePrecise
+        ? await estimateTokensAsync(message.role, { model })
+        : estimateTokens(message.role, { model });
+      totalTokens += roleResult.tokens;
+    }
+
+    // Add attachment estimates
+    if (message.attachments && message.attachments.length > 0) {
+      for (const attachment of message.attachments) {
+        const attachmentText = `${attachment.name} ${attachment.mimeType} ${attachment.caption ?? ''}`;
+        const attachResult = usePrecise
+          ? await estimateTokensAsync(attachmentText, { model })
+          : estimateTokens(attachmentText, { model });
+        totalTokens += attachResult.tokens;
+      }
+    }
+  }
+
+  if (includeFormatting) {
+    totalTokens += CONVERSATION_BASE_OVERHEAD;
+  }
+
+  return {
+    tokens: totalTokens,
+    method: usePrecise ? 'precise' : 'rough',
+    model,
+    confidence: usePrecise ? 0.95 : 0.75,
+  };
+}
+
+/**
+ * Rough token count estimation without any dependencies
+ * Uses a simple characters-per-token heuristic
+ *
+ * @param {string | readonly Message[]} input - Text or messages to estimate
+ * @param {SupportedModel} model - Model for estimation ratios
+ * @returns {number} Estimated token count
+ *
+ * @example
+ * ```typescript
+ * // Estimate text
+ * const textTokens = roughTokenCountEstimation('Hello, world!');
+ *
+ * // Estimate messages
+ * const messageTokens = roughTokenCountEstimation([
+ *   { role: 'user', content: 'Hello!' }
+ * ]);
+ * ```
+ */
+export function roughTokenCountEstimation(
+  input: string | readonly Message[],
+  model: SupportedModel = 'default'
+): number {
+  const ratio = MODEL_TOKEN_RATIOS[model] ?? MODEL_TOKEN_RATIOS.default;
+
+  // Handle string input
+  if (typeof input === 'string') {
+    if (!input) return 0;
+    return Math.ceil(input.length / ratio);
+  }
+
+  // Handle message array input
+  if (!Array.isArray(input) || input.length === 0) {
+    return 0;
+  }
+
+  let totalChars = 0;
+
+  for (const message of input) {
+    // Add content length
+    if (message.content) {
+      totalChars += message.content.length;
+    }
+
+    // Add role length
+    if (message.role) {
+      totalChars += message.role.length;
+    }
+
+    // Add overhead
+    totalChars += MESSAGE_OVERHEAD_TOKENS * ratio;
+
+    // Add attachment metadata
+    if (message.attachments) {
+      for (const attachment of message.attachments) {
+        totalChars += (attachment.name?.length ?? 0);
+        totalChars += (attachment.mimeType?.length ?? 0);
+        totalChars += (attachment.caption?.length ?? 0);
+      }
+    }
+  }
+
+  // Add conversation overhead
+  totalChars += CONVERSATION_BASE_OVERHEAD * ratio;
+
+  return Math.ceil(totalChars / ratio);
+}
+
+/**
+ * Check if text would exceed token limit
+ *
+ * @param {string} text - Text to check
+ * @param {number} limit - Token limit
+ * @param {TokenEstimationOptions} options - Estimation options
+ * @returns {{ withinLimit: true } | { withinLimit: false; exceededBy: number }} Limit check result
+ *
+ * @example
+ * ```typescript
+ * const result = checkTokenLimit('Long text...', 100);
+ * if (!result.withinLimit) {
+ *   console.log(`Exceeded by ${result.exceededBy} tokens`);
+ * }
+ * ```
+ */
+export function checkTokenLimit(
+  text: string,
+  limit: number,
+  options: TokenEstimationOptions = {}
+): { withinLimit: true } | { withinLimit: false; exceededBy: number } {
+  const estimate = estimateTokens(text, options);
+
+  if (estimate.tokens <= limit) {
+    return { withinLimit: true };
+  }
+
+  return {
+    withinLimit: false,
+    exceededBy: estimate.tokens - limit,
+  };
+}
+
+/**
+ * Truncate text to fit within token limit
+ *
+ * @param {string} text - Text to truncate
+ * @param {number} maxTokens - Maximum token count
+ * @param {TokenEstimationOptions} options - Estimation options
+ * @returns {string} Truncated text
+ *
+ * @example
+ * ```typescript
+ * const truncated = truncateToTokenLimit(longText, 100);
+ * ```
+ */
+export function truncateToTokenLimit(
+  text: string,
+  maxTokens: number,
+  options: TokenEstimationOptions = {}
+): string {
+  const estimate = estimateTokens(text, options);
+
+  if (estimate.tokens <= maxTokens) {
+    return text;
+  }
+
+  const ratio = options.customRatio ??
+    MODEL_TOKEN_RATIOS[options.model ?? 'default'] ??
+    MODEL_TOKEN_RATIOS.default;
+
+  // Calculate approximate character limit
+  const targetChars = Math.floor(maxTokens * ratio);
+
+  // Simple truncation (could be improved to respect word boundaries)
+  return text.substring(0, targetChars) + '...';
+}
+
+/**
+ * Get model-specific context window size
+ *
+ * @param {SupportedModel} model - Model name
+ * @returns {number} Context window size in tokens
+ */
+export function getModelContextWindow(model: SupportedModel): number {
+  const contextWindows: Readonly<Record<SupportedModel, number>> = {
+    'gpt-4': 8192,
+    'gpt-4o': 128000,
+    'gpt-4-turbo': 128000,
+    'gpt-3.5-turbo': 16385,
+    'claude-3-opus': 200000,
+    'claude-3-sonnet': 200000,
+    'claude-3-haiku': 200000,
+    'claude-2': 100000,
+    'gemini-pro': 1000000,
+    'gemini-ultra': 1000000,
+    default: 4096,
+  };
+
+  return contextWindows[model] ?? contextWindows.default;
+}
+
+// Default export
+export default {
+  estimateTokens,
+  estimateTokensAsync,
+  estimateTokensForMessages,
+  estimateTokensForMessagesAsync,
+  roughTokenCountEstimation,
+  checkTokenLimit,
+  truncateToTokenLimit,
+  getModelContextWindow,
+  MODEL_TOKEN_RATIOS,
+  MESSAGE_OVERHEAD_TOKENS,
+  CONVERSATION_BASE_OVERHEAD,
+};