import os
import openai
import gradio as gr
from dotenv import load_dotenv

load_dotenv()

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

CSS = """
.contain { height: 100vh !important; display: flex; flex-direction: column; }
.gradio-container { height: 100vh !important; }
.main { height: 100vh !important; }
.wrap { height: 100vh !important; }
#component-0 { height: 100% !important; }
#component-1 { height: 100% !important; }
#chatbot { flex-grow: 1 !important; overflow: auto;}
footer {visibility: hidden}
"""


def predict(message, history):
    history_openai_format = []
    prompt = """
You are GPT-4, a large language model trained by OpenAI.
Carefully read the user's instructions. 
Respond using Markdown.
"""
    history_openai_format.append({"role": "system", "content": prompt})
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human})
        history_openai_format.append({"role": "assistant", "content": assistant})
    history_openai_format.append({"role": "user", "content": message})

    response = openai.ChatCompletion.create(
        model="gpt-4",
        messages=history_openai_format,
        temperature=0.2,
        stream=True,
    )

    partial_message = ""
    for chunk in response:
        if len(chunk["choices"][0]["delta"]) != 0:
            partial_message = partial_message + chunk["choices"][0]["delta"]["content"]
            yield partial_message


chatbot = gr.Chatbot(elem_id="chatbot")
gr.ChatInterface(
    predict,
    chatbot=chatbot,
    title="Tensora GPT-4",
    css=CSS,
).queue().launch()