import requests
import json
from time import sleep

API_KEY = "INSERT YOUR API KEY HERE"
AUTHORIZATION = "INSERT YOUR AUTHORIZATION TOKEN HERE"

url = "https://api.monsterapi.ai/apis/add-task"
fetch_url = "https://api.monsterapi.ai/apis/task-status"

headers = {
    "x-api-key": API_KEY,
    "Authorization": f"Bearer {AUTHORIZATION}",
    "Content-Type": "application/json",
}

data = {
    "model": "llama2-7b-chat",
    "data": {
        "prompt": "How many states in the United States of America?",
        "top_k": 10,
        "top_p": 0.9,
        "temp": 0.1,
        "max_length": 1000,
        "beam_size": 1
    }
}

response = requests.post(url, headers=headers, data=json.dumps(data))
process_id = response.json()["process_id"]

status = None

while True:
    response = requests.post(
        fetch_url,
        headers=headers,
        data=json.dumps({
            "process_id": process_id,
        })).json()

    status = response["response_data"]["status"]
    if status not in ("COMPLETED", "FAILED"):
        sleep(2)
    else:
        break

if status == "COMPLETED":
    print(response["response_data"]["result"]["text"])
else:
    print("Error:", response["response_data"]["result"]["errorMessage"])