import gradio as gr
from together import Together

def send_message(api_key, user_message, history):
    if not api_key:
        return gr.update(value="Please enter your API Key.", visible=True), history
    history = history + [[user_message, None]]
    client = Together(api_key=api_key)
    messages = []
    for human, bot in history:
        if human:
            messages.append({"role": "user", "content": human})
        if bot:
            messages.append({"role": "assistant", "content": bot})
    try:
        response = client.chat.completions.create(
            model="meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
            messages=messages,
            max_tokens=512,
            temperature=0.7,
            top_p=0.7,
            top_k=50,
            repetition_penalty=1,
            stop=["<|eot_id|>", "<|eom_id|>"],
            stream=False
        )
        reply = response.choices[0].message.content
    except Exception as e:
        reply = f"Error: {e}"
    history[-1][1] = reply
    return "", history

with gr.Blocks() as demo:
    gr.Markdown("# Together AI Meta-Llama-3.1-405B-Instruct-Turbo Chatbot Demo")

    history = gr.State([])

    with gr.Column():
        api_key_input = gr.Textbox(label="Together AI API Key", type="password")
        chatbot = gr.Chatbot()
        msg = gr.Textbox(placeholder="Type your message here...")
        send_btn = gr.Button("Send")

    send_btn.click(send_message, [api_key_input, msg, history], [msg, chatbot])

demo.launch()