code_indexer.py

import torch
import numpy as np
import faiss
from pathlib import Path
import os
import asttokens
import ast

class CodeIndexer:
    def __init__(self, tokenizer, model):
        self.code_embedding_tokenizer = tokenizer
        self.code_embedding_model = model

        self.files_content = {}
        self.file_components = {}
        self.code_file_idx = {}
        self.code_file_embeddings = []
        self.filepath_cembeddings = {}
        self.code_file_index = None
        self.file_paths = []

    def load_folder(self, folder_path):
        folder_path = Path(folder_path)
        for file_name in os.listdir(folder_path):
            file_path = folder_path / file_name
            if file_path.is_file():
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                        relative_path = file_path.name
                        self.files_content[relative_path] = content
                except Exception as e:
                    print(f"Error reading {file_path}: {e}")

        self.file_paths = list(self.files_content.keys())
        print(f"Files loaded: {len(self.file_paths)}")
        
        self._build_indices()
        return len(self.files_content)

    def _get_embedding(self, text):
        try:
            inputs = self.code_embedding_tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
            with torch.no_grad():
                outputs = self.code_embedding_model(**inputs)
                embedding = outputs.last_hidden_state[:, 0, :]
            return embedding
        except Exception as e:
            print(f"Error getting embedding: {e}")
            return torch.zeros((1, 768))

    def _extract_components(self, code_string, file_path):
        try:
            atok = asttokens.ASTTokens(code_string, parse=True)
            tree = atok.tree
            components = []
            for node in ast.walk(tree):
                if isinstance(node, (ast.FunctionDef, ast.ClassDef, ast.AsyncFunctionDef)):
                    start_line = node.lineno
                    end_line = max(n.lineno for n in ast.walk(node) if hasattr(n, 'lineno'))
                    component_code = atok.get_text(node)
                    components.append({
                        'type': type(node).__name__,
                        'name': node.name,
                        'code': component_code,
                        'start_line': start_line,
                        'end_line': end_line,
                        'file_path': file_path
                    })
            return components
        except SyntaxError:
            return []

    def _build_indices(self):
        self.code_file_idx.clear()
        self.code_file_embeddings.clear()
        self.filepath_cembeddings.clear()

        print(f"Building indices for {len(self.file_paths)} files")

        for filepath in self.file_paths:
            content = self.files_content[filepath]
            embedding = self._get_embedding(content)
            embedding_np = embedding.squeeze().cpu().numpy().astype('float32')
            if embedding_np.size == 0:
                print(f"Warning: Empty embedding for file {filepath}")
                continue

            self.code_file_embeddings.append(embedding_np)
            self.code_file_idx[len(self.code_file_embeddings) - 1] = filepath

            components = self._extract_components(content, filepath)
            index_to_component = {}
            all_component_embedding = []
            for idx, component in enumerate(components):
                component_embedding = self._get_embedding(component['code'])
                index_to_component[idx] = component
                all_component_embedding.append(component_embedding.squeeze().cpu().numpy().astype('float32'))

            if all_component_embedding:
                all_component_embedding = np.vstack(all_component_embedding)
                dimension = all_component_embedding.shape[1]
                component_index = faiss.IndexFlatIP(dimension)
                faiss.normalize_L2(all_component_embedding)
                component_index.add(all_component_embedding)
                self.filepath_cembeddings[filepath] = (component_index, index_to_component)

        if self.code_file_embeddings:
            try:
                stacked_code_embeddings = np.vstack(self.code_file_embeddings)
                dimension = stacked_code_embeddings.shape[1]
                self.code_file_index = faiss.IndexFlatIP(dimension)
                faiss.normalize_L2(stacked_code_embeddings)
                self.code_file_index.add(stacked_code_embeddings)
                print(f"Code file index built with {len(self.code_file_embeddings)} embeddings")
            except Exception as e:
                print(f"Error building code file index: {e}")

    def create_embeddings_for_folder(folder_path, drive ):
         # Create embeddings for all files in a folder and save directly to Google Drive
        indexer = CodeIndexer()
        
        num_files = indexer.load_folder(folder_path)
        print(f"Processed {num_files} files from {folder_path}")
        
        if folder_id is None:
            folder_metadata = {'title': 'code_embeddings', 'mimeType': 'application/vnd.google-apps.folder'}
            folder = drive.CreateFile(folder_metadata)
            folder.Upload()
            folder_id = folder['id']
            print(f"Created folder 'code_embeddings' in Google Drive with ID: {folder_id}")
        
        indexer.save_index_to_drive(drive, folder_id)
        
        return indexer, folder_id