Upload a file to Trieve with Python

trieve-file-upload.py

import json
import os
import dotenv
import base64
import requests
import mimetypes


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")


def upload_files_to_trieve(local_folder_path, metadata):
    """
    Uploads all files in the specified local folder to the given S3 bucket.
    """

    metadata_json = json.dumps(metadata)

    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...")
                    # Upload the file to Trieve
                    response = requests.post(
                        f"{api_url}/file",
                        headers={
                            "Authorization": f"{api_key}",
                            "TR-Dataset": f"{dataset_id}",
                            "Content-Type": "application/json",
                        },
                        data=json.dumps({
                            "file_name": file_name,
                            "file_mime_type": file_mime_type,
                            "base64_file": base64_file,
                            "metadata": metadata,
                            "create_chunks": True,
                        }),
                    )
                    print(f"Response: {response.status_code}")
                    print(f"Response: {response.json()}")
                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)