@nuggetwise/cli/src/core/semantic-search.ts

Magic Nuggetwise CLI for Cursor IDE integration

import { GeminiClient } from './gemini-client.js';
import { ComponentInfo } from '../types/index.js';

export interface SearchResult {
  component: ComponentInfo;
  score: number;
  relevance: string;
  semanticScore: number;
  keywordScore: number;
}

export interface ComponentEmbedding {
  component: ComponentInfo;
  embedding: number[];
  text: string;
}

export class SemanticSearch {
  private gemini: GeminiClient;
  private componentEmbeddings: ComponentEmbedding[] = [];
  private isIndexed: boolean = false;

  constructor(apiKey: string) {
    this.gemini = new GeminiClient(apiKey);
  }

  async indexComponents(components: ComponentInfo[]): Promise<void> {
    console.log(`📚 Indexing ${components.length} components for semantic search`);
    
    this.componentEmbeddings = [];
    
    for (const component of components) {
      try {
        // Create searchable text from component metadata
        const searchableText = this.createSearchableText(component);
        
        // Generate embedding
        const embedding = await this.gemini.generateEmbeddings(searchableText);
        
        if (embedding.length > 0) {
          this.componentEmbeddings.push({
            component,
            embedding,
            text: searchableText
          });
        }
        
        // Add small delay to avoid rate limits
        await new Promise(resolve => setTimeout(resolve, 50));
        
      } catch (error) {
        console.warn(`Failed to index component ${component.name}:`, error);
      }
    }
    
    this.isIndexed = true;
    console.log(`✅ Successfully indexed ${this.componentEmbeddings.length} components for semantic search`);
  }

  async search(query: string, limit: number = 5): Promise<SearchResult[]> {
    if (!this.isIndexed || this.componentEmbeddings.length === 0) {
      console.warn('No components indexed for semantic search');
      return [];
    }

    try {
      // Generate query embedding
      const queryEmbedding = await this.gemini.generateEmbeddings(query);
      
      if (queryEmbedding.length === 0) {
        console.warn('Failed to generate query embedding, falling back to keyword search');
        return this.keywordSearch(query, limit);
      }

      // Perform hybrid search
      const results = await this.hybridSearch(query, queryEmbedding, limit);
      
      return results;
    } catch (error) {
      console.error('Semantic search error:', error);
      // Fallback to keyword search
      return this.keywordSearch(query, limit);
    }
  }

  private async hybridSearch(
    query: string, 
    queryEmbedding: number[], 
    limit: number
  ): Promise<SearchResult[]> {
    const queryLower = query.toLowerCase();
    const results: SearchResult[] = [];

    for (const item of this.componentEmbeddings) {
      // Calculate semantic similarity
      const semanticScore = this.cosineSimilarity(queryEmbedding, item.embedding);
      
      // Calculate keyword score
      const keywordScore = this.calculateKeywordScore(queryLower, item.text);
      
      // Combine scores (weighted average)
      const combinedScore = (semanticScore * 0.7) + (keywordScore * 0.3);
      
      if (combinedScore > 0.1) { // Minimum relevance threshold
        results.push({
          component: item.component,
          score: combinedScore,
          semanticScore,
          keywordScore,
          relevance: this.generateRelevanceExplanation(query, item.component, semanticScore, keywordScore)
        });
      }
    }

    // Sort by combined score and limit results
    return results
      .sort((a, b) => b.score - a.score)
      .slice(0, limit);
  }

  private keywordSearch(query: string, limit: number): SearchResult[] {
    const queryLower = query.toLowerCase();
    const results: SearchResult[] = [];

    for (const item of this.componentEmbeddings) {
      const keywordScore = this.calculateKeywordScore(queryLower, item.text);
      
      if (keywordScore > 0.1) {
        results.push({
          component: item.component,
          score: keywordScore,
          semanticScore: 0,
          keywordScore,
          relevance: this.generateRelevanceExplanation(query, item.component, 0, keywordScore)
        });
      }
    }

    return results
      .sort((a, b) => b.score - a.score)
      .slice(0, limit);
  }

  private createSearchableText(component: ComponentInfo): string {
    const parts = [
      component.name,
      component.description || '',
      component.category || '',
      ...(component.tags || []),
      component.complexity || '',
      component.source || ''
    ];
    
    return parts.filter(Boolean).join(' ').toLowerCase();
  }

  private calculateKeywordScore(query: string, componentText: string): number {
    const queryWords = query.split(/\s+/).filter(word => word.length >= 2);
    if (queryWords.length === 0) return 0;
    
    let score = 0;
    let exactMatches = 0;
    let partialMatches = 0;
    
    for (const word of queryWords) {
      // Exact word match (highest score)
      if (componentText.includes(word)) {
        score += 2;
        exactMatches++;
      }
      // Partial match (lower score)
      else if (componentText.includes(word.substring(0, Math.max(2, word.length - 1)))) {
        score += 0.5;
        partialMatches++;
      }
    }
    
    // Normalize score
    const normalizedScore = score / (queryWords.length * 2);
    
    // Bonus for multiple exact matches
    if (exactMatches > 1) {
      return Math.min(1, normalizedScore + (exactMatches * 0.1));
    }
    
    return normalizedScore;
  }

  private cosineSimilarity(vecA: number[], vecB: number[]): number {
    if (vecA.length !== vecB.length) {
      console.warn('Vector length mismatch, using fallback similarity');
      return 0;
    }

    let dotProduct = 0;
    let normA = 0;
    let normB = 0;

    for (let i = 0; i < vecA.length; i++) {
      dotProduct += vecA[i] * vecB[i];
      normA += vecA[i] * vecA[i];
      normB += vecB[i] * vecB[i];
    }

    if (normA === 0 || normB === 0) return 0;

    return dotProduct / (Math.sqrt(normA) * Math.sqrt(normB));
  }

  private generateRelevanceExplanation(
    query: string, 
    component: ComponentInfo, 
    semanticScore: number, 
    keywordScore: number
  ): string {
    const explanations = [];
    
    if (semanticScore > 0.8) {
      explanations.push('High semantic similarity');
    } else if (semanticScore > 0.6) {
      explanations.push('Good semantic match');
    }
    
    if (keywordScore > 0.8) {
      explanations.push('Exact keyword matches');
    } else if (keywordScore > 0.5) {
      explanations.push('Partial keyword matches');
    }
    
    if (explanations.length === 0) {
      return 'Semantic similarity';
    }
    
    return explanations.join(', ');
  }

  async getComponentSuggestions(prompt: string): Promise<ComponentInfo[]> {
    const results = await this.search(prompt, 3);
    return results.map(r => r.component);
  }

  getIndexStats(): { totalComponents: number; indexedComponents: number; isIndexed: boolean } {
    return {
      totalComponents: this.componentEmbeddings.length,
      indexedComponents: this.componentEmbeddings.length,
      isIndexed: this.isIndexed
    };
  }

  clearIndex(): void {
    this.componentEmbeddings = [];
    this.isIndexed = false;
    console.log('🗑️ Semantic search index cleared');
  }
}