import os
import openai
import json
from pprint import pprint
from tenacity import retry, wait_random_exponential, stop_after_attempt

openai.api_key = os.environ.get("OPENAI_API_KEY")
MODEL = 'gpt-4'

'''
basic underlying method to get completions. retries built in
'''
@retry(wait=wait_random_exponential(multiplier=2, max=10), stop=stop_after_attempt(3))
def get_completion(messages, temperature=1):

    try:
        chat_completion = openai.ChatCompletion.create(
            model=MODEL,
            temperature=temperature,
            messages=messages
        )

        return chat_completion.choices[0].message.content

    except Exception as e:
        print("Unable to generate ChatCompletion response")
        print(f"Exception: {e}")
        return e

'''
basic text completion
'''
def answer(system_message, user_message, temperature=1):

    messages = [{
      "role": "system",
      "content": system_message
    },{
      "role": "user",
      "content": user_message
    }]

    completion = get_completion(
      messages=messages,
      temperature = temperature
    )

    return completion