import gradio as gr
from openai import OpenAI
import os
import time

# Initialize the OpenAI Client with your API key and endpoint
api_key = os.environ.get("RUNPOD_API_KEY")  # Make sure your API key is correctly configured
client = OpenAI(
    api_key=api_key,
    base_url="https://api.runpod.ai/v2/vllm-k0g4c60zor9xuu/openai/v1",
)

def get_response(user_message, history):
    # Format the history for the OpenAI call
    history_openai_format = []
    for human, assistant in history:
        if human:
            history_openai_format.append({"role": "user", "content": human})
        if assistant:
            history_openai_format.append({"role": "assistant", "content": assistant})
    history_openai_format.append({"role": "user", "content": user_message})

    # Make the API call
    response = client.chat.completions.create(
        model='ambrosfitz/llama-3-history',
        messages=history_openai_format,
        temperature=0.5,
        max_tokens=1028
    )

    # Access the text response
    if response.choices:
        bot_message = response.choices[0].message.content.strip() if response.choices[0].message.content.strip() else "No response generated."
    else:
        bot_message = "No response generated."

    return bot_message

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def user(user_message, history):
        if not user_message.strip():
            return "", history
        bot_response = get_response(user_message, history)
        return "", history + [[user_message, bot_response]]

    def clear_chat():
        return "", []  # Clear the chat history

    msg.submit(user, inputs=[msg, chatbot], outputs=[msg, chatbot])
    clear.click(clear_chat, inputs=None, outputs=[msg, chatbot])

demo.launch()