shared_data/vector_storage.py

import os

from shared_data.supabase_client import supabase, BUCKET

VECTOR_FOLDER = os.getenv("DATABASE_PATH", "vectordatabase")


def upload_vector_db(db_path=VECTOR_FOLDER):

    files = [
        "index.faiss",
        "index.pkl"
    ]

    for filename in files:

        local_path = os.path.join(
            db_path,
            filename
        )

        if not os.path.exists(local_path):
            print(f"{filename} not found")
            continue

        remote_path = f"vectors/{filename}"

        with open(local_path, "rb") as f:
            file_data = f.read()

        try:

            supabase.storage.from_(BUCKET).upload(
                path=remote_path,
                file=file_data,
                file_options={
                    "upsert": "true"
                }
            )

            print(f"Uploaded {filename}")

        except Exception as e:
            print(e)


def download_vector_db(db_path=VECTOR_FOLDER):

    os.makedirs(db_path, exist_ok=True)

    files = [
        "index.faiss",
        "index.pkl"
    ]

    for filename in files:

        remote_path = f"vectors/{filename}"

        local_path = os.path.join(
            db_path,
            filename
        )

        try:

            data = supabase.storage.from_(BUCKET).download(
                remote_path
            )

            with open(local_path, "wb") as f:
                f.write(data)

            print(f"Downloaded {filename}")

        except Exception as e:
            import traceback
            print(traceback.format_exc())
            raise


def delete_vector_db():

    supabase.storage.from_(BUCKET).remove(
        [
            "vectors/index.faiss",
            "vectors/index.pkl"
        ]
    )