utils.py

import gdown
from typing import Optional
import zipfile
import os


def download_from_google_drive(file_id: str, destination: Optional[str] = None) -> None:
    """
    Download a file from Google Drive using its file ID.

    :param file_id: The file ID on Google Drive.
    :param destination: The path to save the file. If None, the file will be saved in the current directory with its original name.
    """
    url = f'https://drive.google.com/uc?id={file_id}'
    output = destination if destination else f"{file_id}.file"
    output = "gdown_chroma_db.zip"
    gdown.download(url, output, quiet=False)

def unzip_file(zip_path: str, extract_to: Optional[str] = None) -> None:
    """
    Unzip a ZIP file to a specified directory.

    :param zip_path: The path to the ZIP file.
    :param extract_to: The directory to extract the files. If None, extracts to the same directory as the ZIP file.
    """
    if extract_to is None:
        extract_to = os.path.dirname(zip_path)

    # Ensure the extraction directory exists
    # os.makedirs(extract_to, exist_ok=True)

    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_to)
        print(f"Files extracted to {extract_to}")


if __name__ == '__main__':
    # Example usage
    file_id = os.getenv(gdown_file_id)  # Replace with your file ID
    download_from_google_drive(file_id)
    zip_file_path = "/content/chroma_db-complete.zip"  # Replace with your zip file path
    extract_path = "/chroma_db-complete"
    unzip_file(zip_file_path,extract_path)