ai-index/lib/smart-query.js

AI-powered local code indexing and search system for any codebase

#!/usr/bin/env node

import fs from 'fs/promises';
import path from 'path';
import { fileURLToPath } from 'url';
import { loadConfig } from './config.js';
import { createLocalEmbedder } from './local-embedder.js';
import { createLocalVectorStore } from './local-vector-store.js';

const __filename = fileURLToPath(import.meta.url);
const __dirname = path.dirname(__filename);

export class SmartQuery {
  constructor(options = {}) {
    this.rootPath = options.rootPath || process.cwd();
    this.indexName = options.indexName || path.basename(this.rootPath).replace(/[^a-zA-Z0-9_-]/g, '_');
    
    this.embedder = null;
    this.vectorStore = null;
    this.symbolIndex = new Map();
    this.fileIndex = new Map();
  }

  async initialize() {
    const config = await loadConfig();
    this.embedder = await createLocalEmbedder(config);
    this.vectorStore = await createLocalVectorStore(config, this.indexName);
    
    // Load indexes if available
    await this.loadIndexes();
  }

  async loadIndexes() {
    try {
      const manifestPath = path.join(this.rootPath, 'ai_index/manifest.json');
      const manifest = JSON.parse(await fs.readFile(manifestPath, 'utf-8'));
      
      // Load symbol index if available
      const symbolIndexPath = path.join(this.rootPath, 'ai_index/symbol_index.json');
      try {
        const symbolData = JSON.parse(await fs.readFile(symbolIndexPath, 'utf-8'));
        Object.entries(symbolData).forEach(([id, data]) => {
          this.symbolIndex.set(id, data);
        });
      } catch (e) {
        // Symbol index might not exist in older versions
      }
    } catch (error) {
      console.error('Warning: Could not load indexes:', error.message);
    }
  }

  async search(query, options = {}) {
    const {
      k = 10,
      type = 'hybrid',
      symbolsOnly = false,
      includeContext = true,
      includeRelationships = true
    } = options;
    
    // Generate query embedding
    const queryEmbedding = await this.embedder.embed(query);
    
    // Perform vector search with enhanced metadata
    // Filter out import chunks to focus on meaningful logic
    const filter = symbolsOnly ? 
      { chunk_type: 'symbol' } : 
      { chunk_type: { $ne: 'imports' } }; // Exclude imports from all searches
      
    const results = await this.vectorStore.hybridSearch(query, queryEmbedding, {
      k: k * 2,
      filter
    });
    
    // Process and enhance results
    const enhanced = await this.enhanceResults(results, {
      query,
      includeContext,
      includeRelationships
    });
    
    // Format for AI consumption
    return this.formatForAI(enhanced, options);
  }

  async enhanceResults(results, options) {
    const enhanced = [];
    
    for (const result of results) {
      const item = {
        file: result.metadata.repo_path,
        score: result.finalScore || result.score,
        type: result.metadata.chunk_type,
        lines: [result.metadata.start_line, result.metadata.end_line],
        content: result.content
      };
      
      // Add symbol information
      if (result.metadata.symbol_name) {
        item.symbol = {
          name: result.metadata.symbol_name,
          type: result.metadata.symbol_type,
          async: result.metadata.async,
          params: result.metadata.params,
          complexity: result.metadata.complexity
        };
        
        // Add methods for classes
        if (result.metadata.methods?.length > 0) {
          item.symbol.methods = result.metadata.methods;
        }
        
        // Add inheritance
        if (result.metadata.extends) {
          item.symbol.extends = result.metadata.extends;
        }
      }
      
      // Add context if requested
      if (options.includeContext) {
        item.context = {
          imports: result.metadata.file_imports || [],
          exports: result.metadata.file_exports || [],
          symbols: result.metadata.file_symbols || []
        };
      }
      
      // Always include usage information for better context
      if (result.metadata.usages?.length > 0 || result.metadata.references?.length > 0) {
        item.usage = {
          calls: result.metadata.usages?.filter(u => u.type === 'function_call' || u.type === 'method_call') || [],
          references: result.metadata.references?.slice(0, 10) || [] // Limit references to avoid noise
        };
      }
      
      // Add relationships if requested
      if (options.includeRelationships) {
        item.relationships = await this.findRelationships(
          result.metadata.repo_path,
          result.metadata.symbol_name
        );
      }
      
      enhanced.push(item);
    }
    
    return enhanced;
  }

  async findRelationships(filePath, symbolName) {
    const relationships = {
      imports_from: [],
      exports_to: [],
      used_by: [],
      uses: []
    };
    
    // Find what this file imports
    const fileData = this.fileIndex.get(filePath);
    if (fileData?.imports) {
      relationships.imports_from = fileData.imports;
    }
    
    // Find who imports this symbol
    if (symbolName) {
      this.fileIndex.forEach((data, file) => {
        if (file !== filePath && data.imports) {
          data.imports.forEach(imp => {
            if (imp.includes(symbolName)) {
              relationships.used_by.push(file);
            }
          });
        }
      });
    }
    
    return relationships;
  }

  formatForAI(results, options) {
    const output = {
      query: options.query || '',
      total_results: results.length,
      results: []
    };
    
    // Group by file for better context
    const fileGroups = new Map();
    
    results.forEach(result => {
      if (!fileGroups.has(result.file)) {
        fileGroups.set(result.file, {
          path: result.file,
          relevance: result.score,
          matches: []
        });
      }
      
      const group = fileGroups.get(result.file);
      
      // Create concise match entry
      const match = {
        lines: result.lines,
        type: result.type
      };
      
      // Add symbol info if present
      if (result.symbol) {
        match.symbol = `${result.symbol.type}:${result.symbol.name}`;
        
        if (result.symbol.params?.length > 0) {
          match.params = result.symbol.params;
        }
        
        if (result.symbol.methods?.length > 0) {
          match.methods = result.symbol.methods.map(m => m.name);
        }
        
        if (result.symbol.async) {
          match.async = true;
        }
      }
      
      // Add key relationships
      if (result.relationships?.used_by?.length > 0) {
        match.used_by = result.relationships.used_by.slice(0, 3);
      }
      
      group.matches.push(match);
    });
    
    // Convert to array and sort by relevance
    output.results = Array.from(fileGroups.values())
      .sort((a, b) => b.relevance - a.relevance)
      .slice(0, options.k || 10);
    
    // Add query understanding
    output.intent = this.analyzeQueryIntent(options.query || '');
    
    // Add navigation hints for AI
    if (output.results.length > 0) {
      output.navigation = {
        entry_files: this.findEntryPoints(output.results),
        key_symbols: this.extractKeySymbols(output.results),
        suggested_order: this.suggestExplorationOrder(output.results)
      };
    }
    
    return output;
  }

  analyzeQueryIntent(query) {
    const q = query.toLowerCase();
    
    return {
      looking_for_definition: /\b(what is|define|definition|interface|type|class)\b/.test(q),
      looking_for_implementation: /\b(how|implement|code|function|method)\b/.test(q),
      looking_for_usage: /\b(use|usage|example|call|invoke)\b/.test(q),
      looking_for_relationship: /\b(import|export|depend|relate|connect)\b/.test(q),
      looking_for_flow: /\b(flow|process|sequence|pipeline|workflow)\b/.test(q)
    };
  }

  findEntryPoints(results) {
    const entries = [];
    
    results.forEach(file => {
      // Check if file looks like an entry point
      const path = file.path.toLowerCase();
      if (
        path.includes('index.') ||
        path.includes('main.') ||
        path.includes('app.') ||
        path.includes('server.') ||
        path.includes('entry.')
      ) {
        entries.push(file.path);
      }
    });
    
    return entries.slice(0, 3);
  }

  extractKeySymbols(results) {
    const symbols = new Map();
    
    results.forEach(file => {
      file.matches.forEach(match => {
        if (match.symbol) {
          const [type, name] = match.symbol.split(':');
          if (!symbols.has(name)) {
            symbols.set(name, {
              name,
              type,
              occurrences: 0,
              files: []
            });
          }
          
          const symbol = symbols.get(name);
          symbol.occurrences++;
          if (!symbol.files.includes(file.path)) {
            symbol.files.push(file.path);
          }
        }
      });
    });
    
    // Return top symbols by occurrence
    return Array.from(symbols.values())
      .sort((a, b) => b.occurrences - a.occurrences)
      .slice(0, 5)
      .map(s => ({
        name: s.name,
        type: s.type,
        primary_file: s.files[0]
      }));
  }

  suggestExplorationOrder(results) {
    const order = [];
    
    // Start with type definitions
    const typeFiles = results.filter(f => 
      f.path.includes('/types/') || f.path.endsWith('.d.ts')
    );
    
    // Then entry points
    const entryFiles = results.filter(f => {
      const p = f.path.toLowerCase();
      return p.includes('index.') || p.includes('main.') || p.includes('app.');
    });
    
    // Then high-relevance files
    const otherFiles = results.filter(f => 
      !typeFiles.includes(f) && !entryFiles.includes(f)
    );
    
    // Combine in suggested order
    [...typeFiles, ...entryFiles, ...otherFiles].forEach(f => {
      if (!order.includes(f.path)) {
        order.push(f.path);
      }
    });
    
    return order.slice(0, 5);
  }

  async findSymbol(symbolName, options = {}) {
    // Direct symbol search
    const results = await this.search(symbolName, {
      ...options,
      symbolsOnly: true
    });
    
    // Filter to exact matches if requested
    if (options.exact) {
      results.results = results.results.filter(r => 
        r.matches.some(m => 
          m.symbol && m.symbol.includes(`:${symbolName}`)
        )
      );
    }
    
    return results;
  }

  async explainCodeFlow(startSymbol, endSymbol = null) {
    // Find the flow between symbols
    const flow = {
      start: startSymbol,
      end: endSymbol,
      path: [],
      files: []
    };
    
    // Search for start symbol
    const startResults = await this.findSymbol(startSymbol, { exact: true });
    
    if (startResults.results.length === 0) {
      flow.error = `Could not find symbol: ${startSymbol}`;
      return flow;
    }
    
    const startFile = startResults.results[0].path;
    flow.files.push(startFile);
    flow.path.push({
      file: startFile,
      symbol: startSymbol,
      type: 'start'
    });
    
    // If end symbol specified, trace the path
    if (endSymbol) {
      const endResults = await this.findSymbol(endSymbol, { exact: true });
      
      if (endResults.results.length > 0) {
        const endFile = endResults.results[0].path;
        
        // Find connection through imports/exports
        // This is simplified - a real implementation would use graph traversal
        flow.path.push({
          file: endFile,
          symbol: endSymbol,
          type: 'end'
        });
        
        flow.files.push(endFile);
      }
    }
    
    return flow;
  }
}

// CLI support
if (import.meta.url === `file://${process.argv[1]}`) {
  const args = process.argv.slice(2);
  
  let query = null;
  let options = {
    k: 10,
    type: 'hybrid'
  };
  
  for (let i = 0; i < args.length; i++) {
    switch (args[i]) {
      case '--q':
      case '--query':
        query = args[++i];
        break;
      case '--k':
        options.k = parseInt(args[++i]);
        break;
      case '--symbol':
        options.symbolsOnly = true;
        break;
      case '--exact':
        options.exact = true;
        break;
      case '--help':
        console.log(`
Smart Query - AI-optimized code search

Usage:
  smart-query --q "query" [options]

Options:
  --q, --query <text>  Search query (required)
  --k <number>         Number of results (default: 10)
  --symbol             Search symbols only
  --exact              Exact symbol name match
  --help               Show this help

Examples:
  smart-query --q "authentication"
  smart-query --q "User" --symbol --exact
  smart-query --q "database connection" --k 5
        `);
        process.exit(0);
      default:
        if (!args[i].startsWith('-')) {
          query = args[i];
        }
    }
  }
  
  if (!query) {
    console.error('Error: Query required. Use --help for usage.');
    process.exit(1);
  }
  
  const searcher = new SmartQuery();
  
  searcher.initialize().then(() => {
    return searcher.search(query, options);
  }).then(results => {
    console.log(JSON.stringify(results, null, 2));
  }).catch(console.error);
}