lumina-trieve-reference.py

lumina-trieve-reference.py

import requests
import time
import os
import json
import random
from dotenv import load_dotenv

load_dotenv()

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


def send_to_trieve(chunks_data: object, chunk_group_data: object):
    headers = {
        "Authorization": api_key,
        "TR-Dataset": dataset_id,
        "Content-Type": "application/json",
    }
    results = send_chunk_group(headers, chunks_data, chunk_group_data)
    print(results)


def send_chunk_group(headers, chunks_data, chunk_group_data):
    start_time = time.time()

    chunk_group_url = f"{api_url}/api/chunk_group"
    chunk_url = f"{api_url}/api/chunk"

    response_group = requests.post(
        chunk_group_url, headers=headers, data=chunk_group_data
    )
    if not response_group.ok:
        print("chunk_group error", response_group.status_code, response_group.json())

    response_chunk = requests.post(
        chunk_url, headers=headers, data=chunks_data
    )
    if not response_chunk.ok:
        print("chunk error", response_chunk.status_code, response_chunk.json())

    end_time = time.time()
    results = {
        "time_taken": end_time - start_time,
        "status_code_group": response_group.status_code,
        "status_code_data": response_chunk.status_code,
    }

    return results

random_group_id = random.randint(0, 1000000)

example_group_data = {
    "name": "example_group",
    "tracking_id": f"example_group_tracking_id_{random_group_id}",
}

example_chunks_data = [{
    "chunk_html": "Hello World",
    "group_tracking_ids": [f"example_group_tracking_id_{random_group_id}"],
    "tracking_id": f"example_chunk_tracking_id_{random_group_id}",
}]

send_to_trieve(json.dumps(example_chunks_data), json.dumps(example_group_data))