lumina-file-upload.py

import os
import dotenv
import base64
import mimetypes
import trieve_py_client
import trieve_py_client.api
import trieve_py_client.api.file_api


def to_base64(file):
    try:
        with open(file, "rb") as f:
            base64_file = base64.b64encode(f.read()).decode("utf-8")
    except Exception as error:
        print(f"An error occurred: {error}")
        return None

    base64_file = base64_file.replace("+", "-")  # Convert '+' to '-'
    base64_file = base64_file.replace("/", "_")  # Convert '/' to '_'
    base64_file = base64_file.rstrip("=")

    return base64_file


dotenv.load_dotenv()

dataset_id = os.getenv("DATASET_ID")
api_key = os.getenv("API_KEY")
api_url = os.getenv("API_URL")

configuration = trieve_py_client.Configuration(host=api_url)
configuration.api_key["ApiKey"] = api_key


def upload_files_to_trieve(local_folder_path, metadata):
    """
    Uploads all files in the specified local folder to the given S3 bucket.
    """
    try:
        for file in os.listdir(local_folder_path):
            local_file_path = os.path.join(local_folder_path, file)

            if os.path.isfile(local_file_path):
                print(f"Uploading {file}...")

                file_name = os.path.basename(local_file_path)
                file_mime_type, _ = mimetypes.guess_type(local_file_path)
                base64_file = to_base64(local_file_path)

                if base64_file:
                    print(f"Uploading {file} to Trieve...")

                    with trieve_py_client.ApiClient(configuration) as api_client:
                        api_instance = trieve_py_client.api.file_api.FileApi(api_client)
                        upload_file_data = trieve_py_client.UploadFileData(
                            file_name=file_name,
                            file_mime_type=file_mime_type,
                            base64_file=base64_file,
                            metadata=metadata,
                            create_chunks=True,
                        )
                        response = api_instance.upload_file_handler(
                            tr_dataset=dataset_id, upload_file_data=upload_file_data
                        )
                    print(f"Response: {response.to_dict()}")
                else:
                    print(f"Failed to convert {file} to base64")
    except Exception as e:
        print(f"An error occurred: {e}")


if __name__ == "__main__":
    file_path = "./files"
    metadata = {
        "todo": "add what seems useful",
    }
    upload_files_to_trieve(file_path, metadata)