cli-app.js

import { pipeline } from '@xenova/transformers';
import pkg from '@pinecone-database/pinecone';
const { Pinecone } = pkg;
import fetch from 'node-fetch';
import fs from 'fs';
import path from 'path';
import crypto from 'crypto';
import readline from 'readline';
import { fileURLToPath } from 'url';

// Configuration
const PINECONE_API_KEY = 'pcsk_52AbGy_CDjDs6zTiABmKYor3xnxxh6qkYQFQEwsNmg9XCQwprpfv4NCWmk6TDnJ3jTECE4';
const PINECONE_INDEX = 'code-embeddings';
const MODEL_NAME = 'Salesforce/codet5-small';
const OLLAMA_URL = 'http://localhost:11434/api/generate';
const OLLAMA_MODEL = 'deepseek-r1:8b';

// Initialize the CLI interface
const rl = readline.createInterface({
  input: process.stdin,
  output: process.stdout
});

class CodeAIAssistant {
  constructor() {
    this.embeddingModel = null;
    this.pineconeClient = null;
    this.pineconeIndex = null;
    this.isInitialized = false;
  }

  async initialize() {
    console.log('Initializing Code AI Assistant...');
    
    try {
      // Initialize Pinecone
      console.log('Connecting to Pinecone...');
      this.pineconeClient = new Pinecone({
        apiKey: PINECONE_API_KEY,
      });
      
      // Get index
      try {
        const indexes = await this.pineconeClient.listIndexes();
        const indexExists = indexes.indexes?.some(idx => idx.name === PINECONE_INDEX);
        
        if (!indexExists) {
          console.log(`Warning: Pinecone index '${PINECONE_INDEX}' not found. Please create it in the Pinecone dashboard with dimension 768 (CodeT5)`);
          console.log('Will attempt to use the index anyway');
        } else {
          console.log(`Connected to Pinecone index '${PINECONE_INDEX}'`);
        }
        
        this.pineconeIndex = this.pineconeClient.index(PINECONE_INDEX);
      } catch (error) {
        console.error(`Error with Pinecone index: ${error.message}`);
        throw error;
      }
      
      // Initialize the CodeT5 model
      console.log('Loading CodeT5 model...');
      this.embeddingModel = await pipeline('feature-extraction', MODEL_NAME);
      console.log('CodeT5 model loaded successfully');
      
      this.isInitialized = true;
      console.log('Initialization complete!');
    } catch (error) {
      console.error('Failed to initialize:', error);
      throw error;
    }
  }

  async generateEmbedding(text) {
    if (!this.isInitialized) {
      throw new Error('AI Assistant not initialized');
    }
    
    const output = await this.embeddingModel(text, { pooling: 'mean', normalize: true });
    return Array.from(output.data);
  }

  async indexFile(filePath) {
    if (!this.isInitialized) {
      throw new Error('AI Assistant not initialized');
    }
    
    console.log(`Indexing file: ${filePath}`);
    
    try {
      // Read the file
      const fileContent = fs.readFileSync(filePath, 'utf8');
      const fileName = path.basename(filePath);
      const language = path.extname(filePath).substring(1); // Remove the dot
      
      // Split content into chunks
      const chunks = this.splitIntoChunks(fileContent, filePath, language);
      console.log(`Split file into ${chunks.length} chunks`);
      
      // Process each chunk
      for (const chunk of chunks) {
        // Generate unique ID for this chunk
        const id = crypto.createHash('md5').update(`${chunk.filePath}:${chunk.startLine}-${chunk.endLine}`).digest('hex');
        
        // Get embedding for the chunk
        const embedding = await this.generateEmbedding(chunk.content);
        
        // Store in Pinecone
        await this.pineconeIndex.upsert([{
          id,
          values: embedding,
          metadata: chunk
        }]);
        
        console.log(`Indexed chunk from lines ${chunk.startLine}-${chunk.endLine}`);
      }
      
      console.log(`Successfully indexed file: ${filePath}`);
      return true;
    } catch (error) {
      console.error(`Error indexing file ${filePath}:`, error);
      return false;
    }
  }

  async indexDirectory(directoryPath, fileExtensions = ['.js', '.ts', '.py', '.java', '.c', '.cpp', '.h', '.cs', '.go']) {
    if (!this.isInitialized) {
      throw new Error('AI Assistant not initialized');
    }
    
    console.log(`Indexing directory: ${directoryPath}`);
    
    try {
      const files = await this.getFilesRecursively(directoryPath, fileExtensions);
      console.log(`Found ${files.length} files to index`);
      
      let successCount = 0;
      let failCount = 0;
      
      for (let i = 0; i < files.length; i++) {
        const file = files[i];
        console.log(`[${i+1}/${files.length}] Processing file: ${file}`);
        
        const success = await this.indexFile(file);
        if (success) {
          successCount++;
        } else {
          failCount++;
        }
      }
      
      console.log(`Indexing complete. Successfully indexed ${successCount} files. Failed to index ${failCount} files.`);
      return { successCount, failCount, totalFiles: files.length };
    } catch (error) {
      console.error(`Error indexing directory ${directoryPath}:`, error);
      throw error;
    }
  }
  
  getFilesRecursively(dir, extensions) {
    let files = [];
    const entries = fs.readdirSync(dir, { withFileTypes: true });
    
    for (const entry of entries) {
      const fullPath = path.join(dir, entry.name);
      
      if (entry.isDirectory()) {
        // Skip node_modules, .git and other common directories to exclude
        const skipDirs = ['node_modules', '.git', 'dist', 'build', 'out'];
        if (!skipDirs.includes(entry.name)) {
          files = files.concat(this.getFilesRecursively(fullPath, extensions));
        }
      } else if (entry.isFile()) {
        const ext = path.extname(entry.name).toLowerCase();
        if (extensions.includes(ext)) {
          files.push(fullPath);
        }
      }
    }
    
    return files;
  }

  splitIntoChunks(text, filePath, language) {
    const lines = text.split('\n');
    const chunks = [];
    
    // Simple chunking strategy
    const maxChunkSize = 100; // Maximum lines per chunk
    const minChunkSize = 10;  // Minimum lines per chunk
    
    let currentChunk = [];
    let startLine = 0;
    
    for (let i = 0; i < lines.length; i++) {
      currentChunk.push(lines[i]);
      
      // Check if we need to create a new chunk
      if (currentChunk.length >= maxChunkSize) {
        chunks.push({
          content: currentChunk.join('\n'),
          filePath: filePath,
          startLine: startLine,
          endLine: i,
          language: language
        });
        
        currentChunk = [];
        startLine = i + 1;
      }
    }
    
    // Add the remaining lines as a chunk if not empty
    if (currentChunk.length >= minChunkSize) {
      chunks.push({
        content: currentChunk.join('\n'),
        filePath: filePath,
        startLine: startLine,
        endLine: lines.length - 1,
        language: language
      });
    }
    
    return chunks;
  }

  async explainCode(code, prompt) {
    if (!this.isInitialized) {
      throw new Error('AI Assistant not initialized');
    }
    
    try {
      console.log('Generating embedding for the code...');
      const embedding = await this.generateEmbedding(code);
      
      console.log('Querying similar code from Pinecone...');
      const queryResult = await this.pineconeIndex.query({
        vector: embedding,
        topK: 5,
        includeMetadata: true,
      });
      
      // Prepare context from matched chunks
      let context = '';
      if (queryResult.matches && queryResult.matches.length > 0) {
        context = queryResult.matches.map(result => {
          const metadata = result.metadata;
          return `From ${metadata.filePath} (lines ${metadata.startLine}-${metadata.endLine}):\n${metadata.content}`;
        }).join('\n\n');
        console.log(`Found ${queryResult.matches.length} similar code segments`);
      } else {
        context = 'No similar code found in the indexed files.';
        console.log('No similar code found');
      }
      
      // Call Ollama with the context and prompt
      console.log('Asking Ollama for explanation...');
      const ollamaPrompt = `You are a programming assistant. Based on this context:\n\n${context}\n\nPlease explain: ${prompt || 'Explain this code'}`;
      
      const response = await fetch(OLLAMA_URL, {
        method: 'POST',
        headers: {
          'Content-Type': 'application/json',
        },
        body: JSON.stringify({
          model: OLLAMA_MODEL,
          prompt: ollamaPrompt,
          stream: false
        }),
      });
      
      const data = await response.json();
      return data.response;
    } catch (error) {
      console.error('Error explaining code:', error);
      throw error;
    }
  }
}

// Main function to run the CLI
async function main() {
  const assistant = new CodeAIAssistant();
  
  try {
    await assistant.initialize();
    
    showMenu();
    
    async function showMenu() {
      console.log('\n------ Code AI Assistant CLI ------');
      console.log('1. Index a file');
      console.log('2. Index a directory');
      console.log('3. Explain code');
      console.log('4. Exit');
      console.log('----------------------------------\n');
      
      rl.question('Choose an option (1-4): ', async (answer) => {
        switch (answer) {
          case '1':
            rl.question('Enter the path to the file you want to index: ', async (filePath) => {
              try {
                await assistant.indexFile(filePath);
                showMenu();
              } catch (error) {
                console.error(`Failed to index file: ${error.message}`);
                showMenu();
              }
            });
            break;
            
          case '2':
            rl.question('Enter the path to the directory you want to index: ', async (dirPath) => {
              try {
                rl.question('Enter file extensions to index (comma separated, e.g. .js,.ts,.py): ', async (extensionsInput) => {
                  let extensions = ['.js', '.ts', '.py', '.java', '.c', '.cpp', '.h', '.cs', '.go'];
                  
                  if (extensionsInput.trim()) {
                    extensions = extensionsInput.split(',').map(ext => ext.trim().startsWith('.') ? ext.trim() : '.' + ext.trim());
                  }
                  
                  console.log(`Indexing files with extensions: ${extensions.join(', ')}`);
                  await assistant.indexDirectory(dirPath, extensions);
                  showMenu();
                });
              } catch (error) {
                console.error(`Failed to index directory: ${error.message}`);
                showMenu();
              }
            });
            break;
            
          case '3':
            rl.question('Enter the path to the file with code you want to explain: ', async (filePath) => {
              try {
                const code = fs.readFileSync(filePath, 'utf8');
                rl.question('What would you like to know about this code? ', async (prompt) => {
                  try {
                    console.log('\nGenerating explanation...');
                    const explanation = await assistant.explainCode(code, prompt);
                    console.log('\n------ Explanation ------');
                    console.log(explanation);
                    console.log('-------------------------\n');
                    showMenu();
                  } catch (error) {
                    console.error(`Failed to explain code: ${error.message}`);
                    showMenu();
                  }
                });
              } catch (error) {
                console.error(`Failed to read file: ${error.message}`);
                showMenu();
              }
            });
            break;
            
          case '4':
            console.log('Goodbye!');
            rl.close();
            process.exit(0);
            break;
            
          default:
            console.log('Invalid option. Please try again.');
            showMenu();
            break;
        }
      });
    }
    
  } catch (error) {
    console.error('Application failed:', error);
    rl.close();
    process.exit(1);
  }
}

// Run the application
main();