@claude-vector/core/src/engine.js

Core vector search engine for code intelligence

/**
 * Core Vector Search Engine
 */

import { OpenAI } from 'openai';
import fs from 'fs/promises';
import path from 'path';
import { SimpleCache } from './cache.js';

export class VectorSearchEngine {
  constructor(config = {}) {
    this.config = {
      openaiApiKey: config.openaiApiKey || process.env.OPENAI_API_KEY,
      embeddingModel: config.embeddingModel || 'text-embedding-3-small',
      searchThreshold: config.searchThreshold || 0.4,
      maxResults: config.maxResults || 10,
      cacheEnabled: config.cacheEnabled ?? true,
      cacheTTL: config.cacheTTL || 3600,
      ...config
    };

    if (!this.config.openaiApiKey) {
      throw new Error('OpenAI API key is required. Set OPENAI_API_KEY environment variable or pass it in config.');
    }

    this.openai = new OpenAI({ apiKey: this.config.openaiApiKey });
    this.embeddings = null;
    this.chunks = null;
    this.index = null;
    
    if (this.config.cacheEnabled) {
      const cacheDir = this.config.cacheDir || path.join(process.cwd(), '.claude-vector-cache');
      this.cache = new SimpleCache(path.join(cacheDir, 'queries'), this.config.cacheTTL);
    }
  }

  /**
   * Load pre-computed embeddings and chunks
   */
  async loadIndex(embeddingsPath, chunksPath) {
    const startTime = Date.now();
    
    try {
      // Load embeddings
      const embeddingsData = await fs.readFile(embeddingsPath, 'utf-8');
      this.embeddings = JSON.parse(embeddingsData);
      
      // Load chunks
      const chunksData = await fs.readFile(chunksPath, 'utf-8');
      this.chunks = JSON.parse(chunksData);
      
      // Validate data
      if (!Array.isArray(this.embeddings) || !Array.isArray(this.chunks)) {
        throw new Error('Invalid index format');
      }
      
      if (this.embeddings.length !== this.chunks.length) {
        throw new Error('Embeddings and chunks count mismatch');
      }
      
      const loadTime = Date.now() - startTime;
      console.log(`✓ Loaded ${this.embeddings.length} embeddings in ${loadTime}ms`);
      
      return {
        embeddingsCount: this.embeddings.length,
        loadTime
      };
    } catch (error) {
      throw new Error(`Failed to load index: ${error.message}`);
    }
  }

  /**
   * Generate embedding for a query
   */
  async generateQueryEmbedding(query) {
    try {
      const response = await this.openai.embeddings.create({
        model: this.config.embeddingModel,
        input: query
      });
      
      return response.data[0].embedding;
    } catch (error) {
      throw new Error(`Failed to generate embedding: ${error.message}`);
    }
  }

  /**
   * Calculate cosine similarity between two vectors
   */
  cosineSimilarity(a, b) {
    if (a.length !== b.length) {
      throw new Error('Vectors must have the same length');
    }
    
    let dotProduct = 0;
    let normA = 0;
    let normB = 0;
    
    for (let i = 0; i < a.length; i++) {
      dotProduct += a[i] * b[i];
      normA += a[i] * a[i];
      normB += b[i] * b[i];
    }
    
    normA = Math.sqrt(normA);
    normB = Math.sqrt(normB);
    
    if (normA === 0 || normB === 0) {
      return 0;
    }
    
    return dotProduct / (normA * normB);
  }

  /**
   * Search for similar chunks
   */
  async search(query, options = {}) {
    if (!this.embeddings || !this.chunks) {
      throw new Error('Index not loaded. Call loadIndex() first.');
    }
    
    const config = { ...this.config, ...options };
    
    // Check cache if enabled
    if (this.cache && !options.noCache) {
      const cacheKey = `search:${query}:${JSON.stringify(options)}`;
      const cached = await this.cache.get(cacheKey);
      if (cached) {
        return cached;
      }
    }
    
    const startTime = Date.now();
    
    // Generate query embedding
    const queryEmbedding = await this.generateQueryEmbedding(query);
    
    // Calculate similarities
    const results = [];
    
    for (let i = 0; i < this.embeddings.length; i++) {
      const similarity = this.cosineSimilarity(queryEmbedding, this.embeddings[i]);
      
      if (similarity >= config.searchThreshold) {
        results.push({
          chunk: this.chunks[i],
          score: similarity,
          index: i
        });
      }
    }
    
    // Sort by score
    results.sort((a, b) => b.score - a.score);
    
    // Limit results
    const limitedResults = results.slice(0, config.maxResults);
    
    const searchTime = Date.now() - startTime;
    
    const response = {
      query,
      results: limitedResults,
      totalMatches: results.length,
      searchTime,
      config: {
        threshold: config.searchThreshold,
        maxResults: config.maxResults
      }
    };
    
    // Cache results if enabled
    if (this.cache && !options.noCache) {
      const cacheKey = `search:${query}:${JSON.stringify(options)}`;
      await this.cache.set(cacheKey, response);
    }
    
    return response;
  }

  /**
   * Find related chunks by similarity to a given chunk
   */
  async findRelated(chunkIndex, options = {}) {
    if (!this.embeddings || !this.chunks) {
      throw new Error('Index not loaded. Call loadIndex() first.');
    }
    
    if (chunkIndex < 0 || chunkIndex >= this.embeddings.length) {
      throw new Error('Invalid chunk index');
    }
    
    const config = { ...this.config, ...options };
    const embedding = this.embeddings[chunkIndex];
    
    const results = [];
    for (let i = 0; i < this.embeddings.length; i++) {
      if (i === chunkIndex) continue; // Skip self
      
      const similarity = this.cosineSimilarity(embedding, this.embeddings[i]);
      
      if (similarity >= config.searchThreshold) {
        results.push({
          chunk: this.chunks[i],
          score: similarity,
          index: i
        });
      }
    }
    
    results.sort((a, b) => b.score - a.score);
    
    return {
      sourceChunk: this.chunks[chunkIndex],
      relatedChunks: results.slice(0, config.maxResults),
      totalMatches: results.length
    };
  }

  /**
   * Get index statistics
   */
  getStats() {
    if (!this.embeddings || !this.chunks) {
      return { loaded: false };
    }
    
    const totalTokens = this.chunks.reduce((sum, chunk) => sum + (chunk.tokens || 0), 0);
    const avgTokensPerChunk = totalTokens / this.chunks.length;
    
    return {
      loaded: true,
      totalChunks: this.chunks.length,
      totalTokens,
      avgTokensPerChunk: Math.round(avgTokensPerChunk),
      embeddingDimensions: this.embeddings[0]?.length || 0,
      indexSizeEstimate: {
        embeddings: `${(JSON.stringify(this.embeddings).length / 1024 / 1024).toFixed(2)} MB`,
        chunks: `${(JSON.stringify(this.chunks).length / 1024 / 1024).toFixed(2)} MB`
      }
    };
  }

  /**
   * Clear loaded index from memory
   */
  clearIndex() {
    this.embeddings = null;
    this.chunks = null;
    this.index = null;
  }
}