adaptive-memory-multi-model-router 2.2.5 → 2.2.7

package/src/cost/preCallCostEstimator.ts

@@ -0,0 +1,345 @@
+/**
+ * A3M Router - Pre-Call Cost Estimator
+ *
+ * Estimates cost BEFORE making an API call based on input features.
+ * Uses historical data patterns and token count estimation via character ratio.
+ *
+ * No external API calls - all estimation is local.
+ *
+ * Usage:
+ *   const estimator = new PreCallCostEstimator();
+ *   const estimate = estimator.estimate({
+ *     query: "What is Python?",
+ *     modelTier: 'mid',
+ *     provider: 'groq'
+ *   });
+ *   console.log(estimate); // { estimatedTokens: 24, estimatedCost: 0.0014, estimatedLatency: 200 }
+ */
+
+import { ProviderTier } from '../providers/providerConfig';
+
+// ============================================================
+// Types
+// ============================================================
+
+export interface CostEstimateInput {
+  /** Query text to estimate tokens for */
+  query: string;
+  /** Model tier (free|cheap|mid|premium|enterprise) */
+  modelTier?: ProviderTier;
+  /** Provider ID for provider-specific cost lookup */
+  provider?: string;
+  /** Optional: explicit token count (if already known) */
+  explicitTokens?: number;
+  /** System prompt length (if using a system prompt) */
+  systemPromptLength?: number;
+}
+
+export interface CostEstimate {
+  /** Estimated total tokens (input + output) */
+  estimatedTokens: number;
+  /** Estimated input tokens */
+  estimatedInputTokens: number;
+  /** Estimated output tokens */
+  estimatedOutputTokens: number;
+  /** Estimated cost in USD */
+  estimatedCost: number;
+  /** Estimated latency in milliseconds */
+  estimatedLatency: number;
+  /** Confidence score 0-1 */
+  confidence: number;
+  /** Breakdown of estimation */
+  breakdown: {
+    inputCostPerM: number;
+    outputCostPerM: number;
+    charToTokenRatio: number;
+  };
+}
+
+export interface ProviderCostConfig {
+  input: number;  // per 1M tokens
+  output: number; // per 1M tokens
+}
+
+// Default latency estimates per tier (ms)
+const TIER_LATENCY: Record<ProviderTier, { min: number; max: number; avg: number }> = {
+  free:   { min: 500,  max: 5000,  avg: 2000 },
+  cheap:  { min: 100,  max: 800,   avg: 300 },
+  mid:    { min: 200,  max: 1500,  avg: 600 },
+  premium:{ min: 300,  max: 2000,  avg: 800 },
+  enterprise: { min: 200, max: 1500, avg: 500 },
+};
+
+// ============================================================
+// PreCallCostEstimator
+// ============================================================
+
+export class PreCallCostEstimator {
+  // Historical data for regression: [charCount, actualTokens][] tuples
+  private historicalData: Array<{ chars: number; tokens: number }> = [];
+  // Linear regression coefficients
+  private slope = 0.25; // chars per token ratio
+  private intercept = 10;
+  // EWMA for latency estimation
+  private latencyHistory: number[] = [];
+  private ewmaAlpha = 0.3;
+
+  constructor(historicalData?: Array<{ chars: number; tokens: number }>) {
+    if (historicalData && historicalData.length > 0) {
+      this.fitLinearRegression(historicalData);
+    }
+  }
+
+  /**
+   * Main estimation method - estimates tokens, cost, and latency.
+   */
+  estimate(input: CostEstimateInput): CostEstimate {
+    const {
+      query,
+      modelTier = 'mid',
+      provider,
+      explicitTokens,
+      systemPromptLength = 0,
+    } = input;
+
+    // Token estimation
+    let inputTokens: number;
+    if (explicitTokens !== undefined) {
+      inputTokens = explicitTokens;
+    } else {
+      inputTokens = this.estimateTokens(query, systemPromptLength);
+    }
+
+    // Output token estimation based on query complexity
+    const outputTokens = this.estimateOutputTokens(query, inputTokens);
+
+    const totalTokens = inputTokens + outputTokens;
+
+    // Cost estimation
+    const costConfig = this.getCostConfig(provider, modelTier);
+    const estimatedCost = this.calculateCost(totalTokens, costConfig);
+
+    // Latency estimation
+    const estimatedLatency = this.estimateLatency(modelTier, inputTokens, totalTokens);
+
+    // Confidence based on amount of historical data we have
+    const confidence = this.calculateConfidence();
+
+    return {
+      estimatedTokens: totalTokens,
+      estimatedInputTokens: inputTokens,
+      estimatedOutputTokens: outputTokens,
+      estimatedCost: Math.round(estimatedCost * 1000000) / 1000000, // 6 decimal places
+      estimatedLatency,
+      confidence,
+      breakdown: {
+        inputCostPerM: costConfig.input,
+        outputCostPerM: costConfig.output,
+        charToTokenRatio: this.slope,
+      },
+    };
+  }
+
+  /**
+   * Estimate input tokens using character-to-token ratio.
+   * Uses linear regression if historical data is available.
+   */
+  estimateTokens(text: string, systemPromptLength = 0): number {
+    const totalChars = text.length + systemPromptLength;
+
+    if (this.historicalData.length >= 5) {
+      // Use linear regression
+      return Math.max(1, Math.round(this.slope * totalChars + this.intercept));
+    }
+
+    // Fallback: general English average ~4 chars per token
+    return Math.max(1, Math.round(totalChars / 4));
+  }
+
+  /**
+   * Estimate output tokens based on query complexity.
+   * More complex queries (code, analysis) tend to need more output.
+   */
+  estimateOutputTokens(query: string, inputTokens: number): number {
+    const lower = query.toLowerCase();
+
+    // Base estimate: ~30% of input tokens as output
+    let multiplier = 0.3;
+
+    // Complexity adjustments
+    if (this.containsKeyword(lower, ['code', 'implement', 'function', 'class', 'algorithm'])) {
+      multiplier = 0.5; // Code needs more output
+    } else if (this.containsKeyword(lower, ['explain', 'describe', 'what is', 'how does'])) {
+      multiplier = 0.35; // Explanations need moderate output
+    } else if (this.containsKeyword(lower, ['list', 'count', 'find all'])) {
+      multiplier = 0.4; // List queries need more output
+    } else if (this.containsKeyword(lower, ['yes', 'no', 'is', 'are', 'does'])) {
+      multiplier = 0.1; // Simple questions need minimal output
+    }
+
+    // Cap to reasonable bounds
+    return Math.min(
+      Math.max(10, Math.round(inputTokens * multiplier)),
+      4000 // Max 4k output tokens
+    );
+  }
+
+  /**
+   * Calculate cost in USD.
+   */
+  calculateCost(totalTokens: number, config: ProviderCostConfig): number {
+    const inputM = this.estimateInputTokensFromTotal(totalTokens) / 1_000_000;
+    const outputM = this.estimateOutputTokensFromTotal(totalTokens) / 1_000_000;
+    return inputM * config.input + outputM * config.output;
+  }
+
+  /**
+   * Estimate latency in milliseconds.
+   */
+  estimateLatency(tier: ProviderTier, inputTokens: number, totalTokens: number): number {
+    const tierLatency = TIER_LATENCY[tier];
+
+    // Base latency from tier
+    let latency = tierLatency.avg;
+
+    // Scale by token count (rough linear approximation)
+    const tokenScale = totalTokens / 100;
+    latency *= Math.max(0.5, Math.min(3, tokenScale));
+
+    // Adjust for historical EWMA if available
+    if (this.latencyHistory.length > 0) {
+      const ewmaLatency = this.calculateEwma();
+      latency = latency * 0.7 + ewmaLatency * 0.3;
+    }
+
+    return Math.round(latency);
+  }
+
+  /**
+   * Record actual tokens for future regression improvements.
+   */
+  recordActualTokens(queryLength: number, actualTokens: number): void {
+    this.historicalData.push({ chars: queryLength, tokens: actualTokens });
+
+    // Keep only last 100 data points
+    if (this.historicalData.length > 100) {
+      this.historicalData.shift();
+    }
+
+    // Refit regression periodically
+    if (this.historicalData.length % 10 === 0) {
+      this.fitLinearRegression(this.historicalData);
+    }
+  }
+
+  /**
+   * Record actual latency for EWMA updates.
+   */
+  recordActualLatency(latencyMs: number): void {
+    this.latencyHistory.push(latencyMs);
+    if (this.latencyHistory.length > 50) {
+      this.latencyHistory.shift();
+    }
+  }
+
+  /**
+   * Fit linear regression to historical data.
+   * Uses ordinary least squares.
+   */
+  fitLinearRegression(data: Array<{ chars: number; tokens: number }>): void {
+    if (data.length < 2) return;
+
+    const n = data.length;
+    let sumX = 0, sumY = 0, sumXY = 0, sumX2 = 0;
+
+    for (const { chars, tokens } of data) {
+      sumX += chars;
+      sumY += tokens;
+      sumXY += chars * tokens;
+      sumX2 += chars * chars;
+    }
+
+    const denominator = n * sumX2 - sumX * sumX;
+    if (denominator === 0) return;
+
+    this.slope = (n * sumXY - sumX * sumY) / denominator;
+    this.intercept = (sumY - this.slope * sumX) / n;
+
+    // Sanity check
+    if (this.slope <= 0 || this.slope > 1) {
+      this.slope = 0.25; // Reset to default if outlier
+      this.intercept = 10;
+    }
+  }
+
+  // ---- Private helpers ----
+
+  private getCostConfig(provider?: string, tier?: ProviderTier): ProviderCostConfig {
+    // Provider-specific costs (from providerConfig.ts patterns)
+    const providerCosts: Partial<Record<string, ProviderCostConfig>> = {
+      groq:      { input: 0.59,  output: 0.79 },
+      cerebras:  { input: 0.60,  output: 0.60 },
+      deepseek:  { input: 0.14,  output: 0.28 },
+      deepinfra: { input: 0.05,  output: 0.05 },
+      together:  { input: 0.18,  output: 0.18 },
+      fireworks: { input: 0.20,  output: 0.20 },
+      mistral:   { input: 0.20,  output: 0.60 },
+      openai:    { input: 2.50,  output: 10.00 },
+      anthropic: { input: 3.00,  output: 15.00 },
+    };
+
+    if (provider && providerCosts[provider]) {
+      return providerCosts[provider]!;
+    }
+
+    // Tier fallback
+    const tierCosts: Record<ProviderTier, ProviderCostConfig> = {
+      free:       { input: 0,     output: 0 },
+      cheap:      { input: 0.20,  output: 0.40 },
+      mid:        { input: 1.00,  output: 3.00 },
+      premium:    { input: 3.00,  output: 12.00 },
+      enterprise: { input: 5.00,  output: 20.00 },
+    };
+
+    return tier ? tierCosts[tier] : tierCosts.mid;
+  }
+
+  private containsKeyword(text: string, keywords: string[]): boolean {
+    return keywords.some(kw => text.includes(kw));
+  }
+
+  private calculateEwma(): number {
+    if (this.latencyHistory.length === 0) return 0;
+    let ewma = this.latencyHistory[0];
+    for (let i = 1; i < this.latencyHistory.length; i++) {
+      ewma = this.ewmaAlpha * this.latencyHistory[i] + (1 - this.ewmaAlpha) * ewma;
+    }
+    return ewma;
+  }
+
+  private calculateConfidence(): number {
+    // More historical data = higher confidence
+    const dataFactor = Math.min(this.historicalData.length / 50, 1.0);
+    // More latency history = higher confidence
+    const latencyFactor = Math.min(this.latencyHistory.length / 20, 1.0);
+    return Math.round((dataFactor * 0.6 + latencyFactor * 0.4) * 100) / 100;
+  }
+
+  private estimateInputTokensFromTotal(total: number): number {
+    return Math.round(total * 0.7); // Assume 70% input
+  }
+
+  private estimateOutputTokensFromTotal(total: number): number {
+    return Math.round(total * 0.3); // Assume 30% output
+  }
+}
+
+// ============================================================
+// Factory
+// ============================================================
+
+export function createPreCallCostEstimator(
+  historicalData?: Array<{ chars: number; tokens: number }>
+): PreCallCostEstimator {
+  return new PreCallCostEstimator(historicalData);
+}

package/src/inference/speculativeDecoding.ts

@@ -0,0 +1,373 @@
+/**
+ * A3M Router - Speculative Decoding Interface
+ *
+ * Interface for integrating speculative decoding.
+ * Currently a stub/interface for future Medusa/Lookahead integration.
+ *
+ * Speculative decoding uses a smaller "draft" model to predict
+ * multiple tokens ahead, which are then verified in parallel by
+ * the main model. This can provide 2-3x speedup in generation.
+ *
+ * Usage:
+ *   const specDec = new SpeculativeDecoding();
+ *   if (specDec.shouldUse(true)) {
+ *     const draftModel = specDec.getDraftModel('medusa');
+ *     // ... use draft model to generate and verify
+ *   }
+ */
+
+import { ProviderTier } from '../providers/providerConfig';
+
+// ============================================================
+// Types
+// ============================================================
+
+export type DraftModelType = 'medusa' | 'lookahead' | 'eagle' | 'spark';
+
+export interface SpeculativeConfig {
+  /** Enable speculative decoding */
+  enabled: boolean;
+  /** Draft model type */
+  draftModelType?: DraftModelType;
+  /** Number of speculative tokens to generate */
+  speculationWindow?: number;
+  /** Temperature for draft model */
+  temperature?: number;
+  /** Provider to use for draft model */
+  provider?: string;
+  /** Model to use for draft model */
+  model?: string;
+}
+
+export interface DraftModel {
+  /** Model identifier */
+  id: string;
+  /** Model type */
+  type: DraftModelType;
+  /** Provider for this model */
+  provider: string;
+  /** Model size description */
+  size: string;
+  /** Supported speculation windows */
+  supportedWindows: number[];
+}
+
+export interface SpeculativeResult {
+  /** Whether speculative decoding was used */
+  used: boolean;
+  /** Number of tokens in draft */
+  draftTokens: number;
+  /** Number of tokens accepted */
+  acceptedTokens: number;
+  /** Acceptance rate */
+  acceptanceRate: number;
+  /** Time saved (ms, estimated) */
+  estimatedTimeSaved: number;
+}
+
+export interface VerificationResult {
+  /** All tokens verified successfully */
+  allAccepted: boolean;
+  /** Indices of accepted tokens */
+  acceptedIndices: number[];
+  /** Indices of rejected tokens */
+  rejectedIndices: number[];
+  /** Actual tokens to use (may differ from draft) */
+  actualTokens: string[];
+  /** Number of tokens to rewind */
+  rewindCount: number;
+}
+
+// ============================================================
+// Speculative Decoding Interface
+// ============================================================
+
+export class SpeculativeDecoding {
+  private config: SpeculativeConfig;
+  private draftModels: Map<string, DraftModel> = new Map();
+
+  constructor(config?: Partial<SpeculativeConfig>) {
+    this.config = {
+      enabled: config?.enabled ?? false,
+      draftModelType: config?.draftModelType ?? 'medusa',
+      speculationWindow: config?.speculationWindow ?? 4,
+      temperature: config?.temperature ?? 0.0,
+      provider: config?.provider,
+      model: config?.model,
+    };
+
+    // Register available draft models
+    this.registerDraftModels();
+  }
+
+  /**
+   * Determine if speculative decoding should be used.
+   * Checks config, provider support, and model availability.
+   */
+  shouldUse(forceEnable?: boolean): boolean {
+    if (forceEnable) return true;
+    if (!this.config.enabled) return false;
+
+    // Speculative decoding beneficial for longer outputs
+    // Check if provider supports it
+    return this.config.enabled;
+  }
+
+  /**
+   * Get the draft model configuration.
+   */
+  getDraftModel(type?: DraftModelType): DraftModel | null {
+    const modelType = type || this.config.draftModelType || 'medusa';
+
+    const model = this.draftModels.get(modelType);
+    if (!model) {
+      // Return a generic configuration
+      return {
+        id: `draft-${modelType}`,
+        type: modelType,
+        provider: this.config.provider || 'auto',
+        size: 'small',
+        supportedWindows: [2, 4, 6, 8],
+      };
+    }
+
+    return model;
+  }
+
+  /**
+   * Generate draft tokens using the draft model.
+   * This is an async method to support API-based draft models.
+   */
+  async generateDraft(
+    prompt: string,
+    maxTokens: number
+  ): Promise<{ tokens: string[]; scores: number[] }> {
+    const draftModel = this.getDraftModel();
+    if (!draftModel) {
+      throw new Error('No draft model available');
+    }
+
+    // Stub: In production, this would call the draft model API
+    // For now, return empty draft
+    return {
+      tokens: [],
+      scores: [],
+    };
+  }
+
+  /**
+   * Verify draft tokens against the main model.
+   * Returns which tokens were accepted and corrections if needed.
+   */
+  async verifyDraft(
+    prompt: string,
+    draftTokens: string[]
+  ): Promise<VerificationResult> {
+    if (draftTokens.length === 0) {
+      return {
+        allAccepted: true,
+        acceptedIndices: [],
+        rejectedIndices: [],
+        actualTokens: [],
+        rewindCount: 0,
+      };
+    }
+
+    // Stub: In production, this would verify each draft token
+    // For now, accept all tokens (optimistic)
+    return {
+      allAccepted: true,
+      acceptedIndices: draftTokens.map((_, i) => i),
+      rejectedIndices: [],
+      actualTokens: draftTokens,
+      rewindCount: 0,
+    };
+  }
+
+  /**
+   * Calculate speedup from speculative decoding results.
+   */
+  calculateSpeedup(result: SpeculativeResult): number {
+    if (!result.used || result.draftTokens === 0) return 1.0;
+
+    // Approximate speedup based on acceptance rate
+    // Higher acceptance = more speedup
+    const acceptanceRate = result.acceptanceRate;
+
+    // Theoretical 2-3x speedup at 90%+ acceptance
+    if (acceptanceRate >= 0.9) return 2.5;
+    if (acceptanceRate >= 0.8) return 2.0;
+    if (acceptanceRate >= 0.7) return 1.7;
+    if (acceptanceRate >= 0.5) return 1.4;
+    if (acceptanceRate >= 0.3) return 1.2;
+    return 1.1;
+  }
+
+  /**
+   * Update configuration.
+   */
+  updateConfig(updates: Partial<SpeculativeConfig>): void {
+    this.config = { ...this.config, ...updates };
+  }
+
+  /**
+   * Get current configuration.
+   */
+  getConfig(): SpeculativeConfig {
+    return { ...this.config };
+  }
+
+  /**
+   * Check if a provider supports speculative decoding.
+   */
+  supportsSpeculativeDecoding(provider: string): boolean {
+    // Local providers (Ollama, vLLM) typically support it
+    const localProviders = ['ollama', 'lmstudio', 'vllm'];
+    if (localProviders.includes(provider)) return true;
+
+    // Check if provider has known speculative support
+    const supportedProviders = ['together', 'fireworks', 'anyscale'];
+    return supportedProviders.includes(provider);
+  }
+
+  /**
+   * Get recommended speculation window based on model size.
+   */
+  getRecommendedWindow(modelSize: 'small' | 'medium' | 'large'): number {
+    switch (modelSize) {
+      case 'small': return 6;
+      case 'medium': return 4;
+      case 'large': return 2;
+      default: return 4;
+    }
+  }
+
+  // ---- Private helpers ----
+
+  private registerDraftModels(): void {
+    // Medusa-style models (multiple draft heads)
+    this.draftModels.set('medusa', {
+      id: 'medusa-7b',
+      type: 'medusa',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6, 8, 10],
+    });
+
+    // Lookahead decoding
+    this.draftModels.set('lookahead', {
+      id: 'lookahead-7b',
+      type: 'lookahead',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6],
+    });
+
+    // EAGLE decoding
+    this.draftModels.set('eagle', {
+      id: 'eagle-7b',
+      type: 'eagle',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6, 8],
+    });
+
+    // Spark (speculative decoding for transformers)
+    this.draftModels.set('spark', {
+      id: 'spark-7b',
+      type: 'spark',
+      provider: 'auto',
+      size: '7B equivalent',
+      supportedWindows: [2, 4, 6, 8, 10, 12],
+    });
+  }
+}
+
+// ============================================================
+// Speculative Decoding Wrapper for LLM calls
+// ============================================================
+
+export class SpeculativeDecodingWrapper {
+  private specDec: SpeculativeDecoding;
+  private mainModelCall: (prompt: string, options?: any) => Promise<string>;
+
+  constructor(
+    mainModelCall: (prompt: string, options?: any) => Promise<string>,
+    config?: Partial<SpeculativeConfig>
+  ) {
+    this.mainModelCall = mainModelCall;
+    this.specDec = new SpeculativeDecoding(config);
+  }
+
+  /**
+   * Generate with speculative decoding.
+   */
+  async generate(
+    prompt: string,
+    maxTokens: number,
+    options?: { temperature?: number; useSpecDec?: boolean }
+  ): Promise<{ text: string; result: SpeculativeResult }> {
+    if (!this.specDec.shouldUse(options?.useSpecDec)) {
+      // Fallback to normal generation
+      const text = await this.mainModelCall(prompt, options);
+      return {
+        text,
+        result: {
+          used: false,
+          draftTokens: 0,
+          acceptedTokens: 0,
+          acceptanceRate: 0,
+          estimatedTimeSaved: 0,
+        },
+      };
+    }
+
+    // Speculative generation
+    const startTime = Date.now();
+    const draftTokens = await this.specDec.generateDraft(prompt, maxTokens);
+    const verification = await this.specDec.verifyDraft(prompt, draftTokens.tokens);
+
+    let text: string;
+
+    if (verification.allAccepted) {
+      // All draft tokens accepted, continue with main model
+      text = verification.actualTokens.join('') +
+        await this.mainModelCall(prompt, options);
+    } else {
+      // Some tokens rejected, need to rewind
+      text = verification.actualTokens.join('');
+      // Re-prompt with corrected context
+      text += await this.mainModelCall(prompt + text, options);
+    }
+
+    const endTime = Date.now();
+    const speedup = this.specDec.calculateSpeedup({
+      used: true,
+      draftTokens: draftTokens.tokens.length,
+      acceptedTokens: verification.acceptedIndices.length,
+      acceptanceRate: verification.acceptedIndices.length / draftTokens.tokens.length,
+      estimatedTimeSaved: endTime - startTime,
+    });
+
+    return {
+      text,
+      result: {
+        used: true,
+        draftTokens: draftTokens.tokens.length,
+        acceptedTokens: verification.acceptedIndices.length,
+        acceptanceRate: verification.acceptedIndices.length / draftTokens.tokens.length,
+        estimatedTimeSaved: Math.round((1 - 1 / speedup) * (endTime - startTime)),
+      },
+    };
+  }
+}
+
+// ============================================================
+// Factory
+// ============================================================
+
+export function createSpeculativeDecoding(
+  config?: Partial<SpeculativeConfig>
+): SpeculativeDecoding {
+  return new SpeculativeDecoding(config);
+}