@claude-vector/core/src/embedding-generator.js

Core vector search engine for code intelligence

/**
 * Embedding generator using OpenAI API
 * Enhanced with comprehensive error resilience and recovery features
 */

import { OpenAI } from 'openai';
import { APIResilience } from './api-resilience.js';

export class EmbeddingGenerator {
  constructor(config = {}) {
    this.config = {
      model: config.model || 'text-embedding-3-small',
      batchSize: config.batchSize || 100,
      maxRetries: config.maxRetries || 5, // 増加
      retryDelay: config.retryDelay || 1000,
      
      // APIレジリエンス設定
      enableResilience: config.enableResilience ?? true,
      resilienceOptions: {
        maxRetries: config.maxRetries || 5,
        baseDelay: config.retryDelay || 1000,
        maxDelay: 30000,
        circuitBreakerEnabled: true,
        rateLimitDetection: true,
        adaptiveRateLimiting: true,
        verboseLogging: config.verboseLogging ?? false,
        ...config.resilienceOptions
      },
      
      ...config
    };

    if (!config.openaiApiKey && !process.env.OPENAI_API_KEY) {
      throw new Error('OpenAI API key is required');
    }

    this.openai = new OpenAI({
      apiKey: config.openaiApiKey || process.env.OPENAI_API_KEY
    });
    
    // APIレジリエンスシステム初期化
    if (this.config.enableResilience) {
      this.apiResilience = new APIResilience(this.config.resilienceOptions);
      
      // イベントリスナー設定
      this.setupResilienceEventListeners();
    }
    
    // 統計情報
    this.stats = {
      totalEmbeddings: 0,
      totalBatches: 0,
      failedBatches: 0,
      totalTokensProcessed: 0,
      averageLatency: 0
    };
  }

  /**
   * Generate embeddings for multiple texts
   */
  async generateEmbeddings(texts) {
    const embeddings = [];
    
    // Process in batches
    for (let i = 0; i < texts.length; i += this.config.batchSize) {
      const batch = texts.slice(i, i + this.config.batchSize);
      const batchEmbeddings = await this.generateBatch(batch);
      embeddings.push(...batchEmbeddings);
    }
    
    return embeddings;
  }

  /**
   * APIレジリエンスイベントリスナー設定
   */
  setupResilienceEventListeners() {
    this.apiResilience.on('circuit-breaker-open', (data) => {
      console.warn('⚠️  Circuit breaker opened due to repeated failures');
      console.warn(`   Next retry in ${Math.round((data.nextRetryTime - Date.now()) / 1000)}s`);
    });
    
    this.apiResilience.on('circuit-breaker-closed', () => {
      console.log('✅ Circuit breaker closed - service recovered');
    });
    
    this.apiResilience.on('rate-limit-adjusted', (data) => {
      if (this.config.resilienceOptions.verboseLogging) {
        console.log(`⚙️  Rate limit adjusted: ${data.oldRate} → ${data.newRate} requests/min`);
      }
    });
    
    this.apiResilience.on('retry-attempt', (data) => {
      if (this.config.resilienceOptions.verboseLogging) {
        console.log(`🔄 Retrying embedding generation (attempt ${data.attempt}/${this.config.maxRetries})`);
      }
    });
  }

  /**
   * Estimate token count for a text
   * @param {string} text - Text to estimate tokens for
   * @returns {number} Estimated token count
   */
  estimateTokens(text) {
    // Simple estimation: ~4 characters per token on average
    // This is a reasonable approximation for code and English text
    return Math.ceil(text.length / 4);
  }

  /**
   * Split large text into smaller chunks
   * @param {string} text - Text to split
   * @param {number} maxTokens - Maximum tokens per chunk
   * @returns {string[]} Array of text chunks
   */
  splitLargeText(text, maxTokens = 6000) {
    const estimatedTokens = this.estimateTokens(text);
    
    // If within limits, return as-is
    if (estimatedTokens <= maxTokens) {
      return [text];
    }
    
    // Calculate chunk size with safety margin
    const targetTokensPerChunk = maxTokens * 0.9; // 90% of max to be safe
    const charsPerChunk = Math.floor(text.length * (targetTokensPerChunk / estimatedTokens));
    
    const chunks = [];
    const lines = text.split('\n');
    let currentChunk = '';
    let currentTokens = 0;
    
    for (const line of lines) {
      const lineTokens = this.estimateTokens(line);
      
      // If single line exceeds limit, split by characters
      if (lineTokens > targetTokensPerChunk) {
        // Save current chunk if not empty
        if (currentChunk) {
          chunks.push(currentChunk);
          currentChunk = '';
          currentTokens = 0;
        }
        
        // Split the large line
        const lineChunks = this.splitLineByTokens(line, targetTokensPerChunk);
        chunks.push(...lineChunks.slice(0, -1)); // Add all but last
        currentChunk = lineChunks[lineChunks.length - 1];
        currentTokens = this.estimateTokens(currentChunk);
      } else if (currentTokens + lineTokens > targetTokensPerChunk) {
        // Start new chunk
        chunks.push(currentChunk);
        currentChunk = line;
        currentTokens = lineTokens;
      } else {
        // Add to current chunk
        currentChunk += (currentChunk ? '\n' : '') + line;
        currentTokens += lineTokens;
      }
    }
    
    // Add final chunk
    if (currentChunk) {
      chunks.push(currentChunk);
    }
    
    return chunks;
  }

  /**
   * Split a single line by token count
   * @param {string} line - Line to split
   * @param {number} maxTokens - Maximum tokens per chunk
   * @returns {string[]} Array of line chunks
   */
  splitLineByTokens(line, maxTokens) {
    const chunks = [];
    const charsPerToken = 4; // Average
    const charsPerChunk = maxTokens * charsPerToken;
    
    for (let i = 0; i < line.length; i += charsPerChunk) {
      chunks.push(line.slice(i, i + charsPerChunk));
    }
    
    return chunks;
  }

  /**
   * Average multiple embeddings into a single embedding
   * @param {number[][]} embeddings - Array of embeddings to average
   * @returns {number[]} Averaged embedding
   */
  averageEmbeddings(embeddings) {
    if (embeddings.length === 0) {
      throw new Error('Cannot average empty embeddings array');
    }
    
    const dimensions = embeddings[0].length;
    const averaged = new Array(dimensions).fill(0);
    
    // Sum all embeddings
    for (const embedding of embeddings) {
      for (let i = 0; i < dimensions; i++) {
        averaged[i] += embedding[i];
      }
    }
    
    // Average
    for (let i = 0; i < dimensions; i++) {
      averaged[i] /= embeddings.length;
    }
    
    return averaged;
  }

  /**
   * Generate embeddings for a batch of texts
   */
  async generateBatch(texts) {
    const startTime = Date.now();
    this.stats.totalBatches++;
    
    // Check each text for token limits and split if necessary
    const processedTexts = [];
    const textMapping = []; // Track which original text each chunk belongs to
    
    for (let i = 0; i < texts.length; i++) {
      const text = texts[i];
      const estimatedTokens = this.estimateTokens(text);
      
      if (estimatedTokens > 8000) {
        // Split large text
        console.warn(`⚠️ Text ${i} has ~${estimatedTokens} tokens, splitting into smaller chunks...`);
        const chunks = this.splitLargeText(text);
        
        for (const chunk of chunks) {
          processedTexts.push(chunk);
          textMapping.push(i);
        }
      } else {
        processedTexts.push(text);
        textMapping.push(i);
      }
    }
    
    const apiCall = async () => {
      const response = await this.openai.embeddings.create({
        model: this.config.model,
        input: processedTexts
      });
      
      return response.data.map(item => item.embedding);
    };
    
    try {
      let embeddings;
      
      if (this.config.enableResilience && this.apiResilience) {
        // レジリエンス機能付きで実行
        embeddings = await this.apiResilience.executeWithResilience(apiCall, {
          batchSize: processedTexts.length,
          model: this.config.model,
          operation: 'generate_embeddings'
        });
      } else {
        // 従来のリトライロジック（レガシー互換性）
        embeddings = await this.executeLegacyRetry(apiCall);
      }
      
      // If texts were split, we need to reconstruct the result
      if (processedTexts.length > texts.length) {
        const reconstructedEmbeddings = [];
        
        for (let i = 0; i < texts.length; i++) {
          // Find all embeddings belonging to this original text
          const chunkEmbeddings = [];
          
          for (let j = 0; j < textMapping.length; j++) {
            if (textMapping[j] === i) {
              chunkEmbeddings.push(embeddings[j]);
            }
          }
          
          // For split texts, average the embeddings
          if (chunkEmbeddings.length > 1) {
            console.log(`📊 Averaging ${chunkEmbeddings.length} chunk embeddings for text ${i}`);
            const averaged = this.averageEmbeddings(chunkEmbeddings);
            reconstructedEmbeddings.push(averaged);
          } else {
            reconstructedEmbeddings.push(chunkEmbeddings[0]);
          }
        }
        
        embeddings = reconstructedEmbeddings;
      }
      
      // 統計更新
      this.updateStats(texts.length, Date.now() - startTime, true);
      
      return embeddings;
      
    } catch (error) {
      this.stats.failedBatches++;
      this.updateStats(texts.length, Date.now() - startTime, false);
      
      throw new Error(`Failed to generate embeddings for batch of ${texts.length} texts: ${error.message}`);
    }
  }

  /**
   * レガシーリトライロジック（後方互換性用）
   */
  async executeLegacyRetry(apiCall) {
    let attempts = 0;
    
    while (attempts < this.config.maxRetries) {
      try {
        return await apiCall();
      } catch (error) {
        attempts++;
        
        if (attempts >= this.config.maxRetries) {
          throw error;
        }
        
        // Wait before retry
        await new Promise(resolve => setTimeout(resolve, this.config.retryDelay * attempts));
      }
    }
  }

  /**
   * 統計情報更新
   */
  updateStats(textCount, latency, success) {
    this.stats.totalEmbeddings += textCount;
    
    // トークン数推定（平均300トークン/テキスト）
    const estimatedTokens = textCount * 300;
    this.stats.totalTokensProcessed += estimatedTokens;
    
    // 移動平均でレイテンシ更新
    const alpha = 0.1;
    this.stats.averageLatency = this.stats.averageLatency * (1 - alpha) + latency * alpha;
  }

  /**
   * Generate embedding for a single text
   */
  async generateEmbedding(text) {
    const [embedding] = await this.generateEmbeddings([text]);
    return embedding;
  }

  /**
   * Estimate cost for generating embeddings
   */
  estimateCost(textCount) {
    const costs = {
      'text-embedding-3-small': 0.00002,  // per 1K tokens
      'text-embedding-3-large': 0.00013,  // per 1K tokens
      'text-embedding-ada-002': 0.00010   // per 1K tokens
    };
    
    const costPer1K = costs[this.config.model] || costs['text-embedding-3-small'];
    // Rough estimate: average 500 tokens per text
    const estimatedTokens = textCount * 500;
    const estimatedCost = (estimatedTokens / 1000) * costPer1K;
    
    return {
      textCount,
      estimatedTokens,
      estimatedCost,
      model: this.config.model
    };
  }

  /**
   * 包括的統計情報取得
   */
  getComprehensiveStats() {
    const basicStats = { ...this.stats };
    
    // APIレジリエンス統計
    const resilienceStats = this.apiResilience ? this.apiResilience.getStats() : null;
    
    // 成功率計算
    const successRate = this.stats.totalBatches > 0 ? 
      ((this.stats.totalBatches - this.stats.failedBatches) / this.stats.totalBatches) * 100 : 100;
    
    return {
      embeddings: basicStats,
      resilience: resilienceStats,
      performance: {
        successRate: Math.round(successRate * 100) / 100,
        averageLatency: Math.round(this.stats.averageLatency),
        embeddingsPerSecond: this.stats.averageLatency > 0 ? 
          Math.round(1000 / this.stats.averageLatency * this.config.batchSize) : 0
      },
      configuration: {
        model: this.config.model,
        batchSize: this.config.batchSize,
        resilienceEnabled: this.config.enableResilience,
        maxRetries: this.config.maxRetries
      }
    };
  }

  /**
   * ヘルスチェック
   */
  getHealthStatus() {
    const stats = this.getComprehensiveStats();
    
    let status = 'healthy';
    let issues = [];
    
    // APIレジリエンスのヘルスチェック
    if (this.apiResilience) {
      const resilienceHealth = this.apiResilience.getHealthStatus();
      if (resilienceHealth.status !== 'healthy') {
        status = resilienceHealth.status;
        issues.push(...resilienceHealth.issues);
      }
    }
    
    // 成功率チェック
    if (stats.performance.successRate < 90) {
      status = status === 'healthy' ? 'degraded' : status;
      issues.push(`Low success rate: ${stats.performance.successRate}%`);
    }
    
    // レイテンシチェック
    if (stats.performance.averageLatency > 10000) { // 10秒以上
      status = status === 'healthy' ? 'degraded' : status;
      issues.push(`High latency: ${stats.performance.averageLatency}ms`);
    }
    
    return {
      status,
      issues,
      stats: stats,
      timestamp: Date.now()
    };
  }

  /**
   * 統計リセット
   */
  resetStats() {
    this.stats = {
      totalEmbeddings: 0,
      totalBatches: 0,
      failedBatches: 0,
      totalTokensProcessed: 0,
      averageLatency: 0
    };
    
    if (this.apiResilience) {
      this.apiResilience.resetStats();
    }
  }

  /**
   * APIレジリエンス機能の有効/無効切り替え
   */
  setResilienceEnabled(enabled) {
    this.config.enableResilience = enabled;
    
    if (enabled && !this.apiResilience) {
      this.apiResilience = new APIResilience(this.config.resilienceOptions);
      this.setupResilienceEventListeners();
    }
  }

  /**
   * 設定更新
   */
  updateConfig(newConfig) {
    this.config = { ...this.config, ...newConfig };
    
    // APIレジリエンス設定の更新
    if (this.apiResilience && newConfig.resilienceOptions) {
      // 新しいインスタンスを作成（設定の動的更新は複雑なため）
      this.apiResilience = new APIResilience({
        ...this.config.resilienceOptions,
        ...newConfig.resilienceOptions
      });
      this.setupResilienceEventListeners();
    }
  }
}