initialize.py

import os
import tempfile
from pathlib import Path
from pydrive.auth import GoogleAuth
from pydrive.drive import GoogleDrive
from oauth2client.client import GoogleCredentials
from google.colab import auth
from code_searcher import CodeSearcher


def download_indices_from_gdrive(drive_folder_name="code_embeddings", local_folder="code_embeddings"):
    if local_folder is None:
        local_folder = tempfile.mkdtemp()
    else:
        os.makedirs(local_folder, exist_ok=True)

    auth.authenticate_user()
    gauth = GoogleAuth()
    gauth.credentials = GoogleCredentials.get_application_default()
    drive = GoogleDrive(gauth)

    folder_list = drive.ListFile({
        'q': f"title='{drive_folder_name}' and mimeType='application/vnd.google-apps.folder' and trashed=false"
    }).GetList()
    if not folder_list:
        print(f"Folder '{drive_folder_name}' not found.")
        return None

    folder_id = folder_list[0]['id']
    files = drive.ListFile({'q': f"'{folder_id}' in parents and trashed=false"}).GetList()
    if not files:
        print(f"No files in folder '{drive_folder_name}'.")
        return None

    for f in files:
        f.GetContentFile(os.path.join(local_folder, f['title']))
        print(f"Downloaded {f['title']}")

    return local_folder


def initialize_searcher(folder_name="code_embeddings", quantize=True):
    print("Initializing Code Searcher...")
    index_dir = download_indices_from_gdrive(folder_name)
    if not index_dir:
        return None
    try:
        return CodeSearcher(index_dir, quantize=quantize)
    except Exception:
        if quantize:
            print("Retrying without quantization...")
            return CodeSearcher(index_dir, quantize=False)
        return None