Create New Trieve Dataset With Python

create_trieve_dataset.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()

organization_id = os.getenv("ORGANIZATION_ID")
api_key = os.getenv("API_KEY")


def create_trieve_dataset(dataset_name):
    """
    Create a new dataset in your Trieve organization
    """
    response = requests.post(
        f"https://api.trieve.ai/api/dataset",
        headers={
            "Authorization": f"{api_key}",
            "TR-Organization": f"{organization_id}",
            "Content-Type": "application/json",
        },
        data=json.dumps({
            "organization_id": organization_id,
            "dataset_name": dataset_name,
            "server_configuration": {},
            "client_configuration": {},
        }),
    )
    print(f"Response: {response.status_code}")

    new_dataset = response.json()
    dataset_id = new_dataset["id"]
    print(f"Dataset ID: {dataset_id}")
    print(f"Dataset: {new_dataset}")
    
if __name__ == "__main__":
    create_trieve_dataset("marketer persona")