# Simple Python example to do TextSynth completion requests
#
# The API secret key must be in the TEXTSYNTH_API_SECRET_KEY environment
# variable.
import os
import sys
import requests

api_url = "https://api.textsynth.com"
# read from environment variable TEXTSYNTH_API_SECRET_KEY
api_key = os.environ["TEXTSYNTH_API_SECRET_KEY"]
# engine to use
api_engine = "gptj_6B"

def textsynth_completion(prompt, max_tokens):
    response = requests.post(api_url + "/v1/engines/" + api_engine + "/completions", headers = { "Authorization": "Bearer " + api_key }, json = { "prompt": prompt, "max_tokens": max_tokens })
    resp = response.json()
    if "text" in resp: 
        return resp["text"]
    else:
        print("ERROR", resp)
        assert False


def main():
    if len(sys.argv) <= 1:
        print("usage: {:s} prompt".format(sys.argv[0]))
        sys.exit(1)
    prompt = sys.argv[1]
    res = textsynth_completion(prompt, 30)
    print(prompt + res)

main()