CodeRAG/save_to_drive.py at main · RohanDisa/CodeRAG · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
import pickle
import tempfile
import numpy as np
import faiss
from pathlib import Path

def save_index_to_drive(indexer, drive, folder_id):
    print("Saving indices directly to Google Drive...")
    with tempfile.TemporaryDirectory() as temp_dir:
        temp_dir_path = Path(temp_dir)
        file_uploads = []

        code_index_path = temp_dir_path / "code_file_index.faiss"
        faiss.write_index(indexer.code_file_index, str(code_index_path))
        file_uploads.append(("code_file_index.faiss", code_index_path))

        for file_path, (component_index, _) in indexer.filepath_cembeddings.items():
            safe_name = file_path.replace("/", "_").replace("\\", "_")
            component_path = temp_dir_path / f"component_{safe_name}.faiss"
            faiss.write_index(component_index, str(component_path))
            file_uploads.append((f"component_{safe_name}.faiss", component_path))

        np.save(str(temp_dir_path / "file_paths.npy"), indexer.file_paths)
        file_uploads.append(("file_paths.npy", temp_dir_path / "file_paths.npy"))

        np.save(str(temp_dir_path / "code_file_idx.npy"), indexer.code_file_idx)
        file_uploads.append(("code_file_idx.npy", temp_dir_path / "code_file_idx.npy"))

        with open(str(temp_dir_path / "files_content.pkl"), 'wb') as f:
            pickle.dump(indexer.files_content, f)
        file_uploads.append(("files_content.pkl", temp_dir_path / "files_content.pkl"))

        component_mappings = {file_path: mapping for file_path, (_, mapping) in indexer.filepath_cembeddings.items()}
        with open(str(temp_dir_path / "component_mappings.pkl"), 'wb') as f:
            pickle.dump(component_mappings, f)
        file_uploads.append(("component_mappings.pkl", temp_dir_path / "component_mappings.pkl"))

        for filename, filepath in file_uploads:
            try:
                file_metadata = {'title': filename, 'parents': [{'id': folder_id}]}
                gfile = drive.CreateFile(file_metadata)
                gfile.SetContentFile(str(filepath))
                gfile.Upload()
                print(f"Uploaded {filename} to Google Drive")
            except Exception as e:
                print(f"Error uploading {filename}: {e}")

    print("All indices and data saved to Google Drive")