import os
import re
import openai

openai.api_key = os.getenv("OPENAI_API_KEY")

preamble = """
I am a highly intelligent question answering bot. If you ask me a
question that is rooted in truth, I will give you the answer. If you
ask me a question that is nonsense, trickery, or has no clear answer,
I will respond with "Unknown".

"""

premise = """
cube1 is visible.
cube1 orientation is upright.
cube1 position is (200.5, 11.3).
cube2 is not visible.
cube2 orientation is upright.
cube2 position is (150.5, 34.5).
cube3 is visible.
cube3 orientation is sideways.
cube3 position is (-50.5, 138.7).
"""

query = "Which cubes are visible?"

print(premise)
print()

client = openai.OpenAI()

while True:
    query = input("Query? ")
    response = client.chat.completions.create(
        model="gpt-4-turbo-preview",
        messages = [
            {"role": "system", "content": preamble},
            {"role": "user", "content": premise},
            {"role": "user", "content": f"Q: {query}"}
        ]
    )
    for choice in response.choices:
        # remove LaTeX brackets from response
        cleaned_response = re.sub(r'\\[\[\]\(\)]', '', choice.message.content)
        print(cleaned_response)
    print()