import os
import gradio as gr

# Install the groq package if it is not installed
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq
    
# 從環境變數獲取 API 密鑰
groq_key = os.getenv('groq_key')

# 定義 Chatbot 回應函數
def chatbot_response(user_message):
    client = Groq(api_key=groq_key)
    completion = client.chat.completions.create(
        model="llama-3.1-70b-versatile",
        messages=[
            {
                "role": "system",
                "content": "brat"
            },
            {
                "role": "user",
                "content": user_message
            }
        ],
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=True,
        stop=None,
    )
    
    response = ""
    for chunk in completion:
        response += chunk.choices[0].delta.content or ""
    
    return response

# 創建 Gradio 介面
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.Button("Clear")

    def respond(message, chat_history):
        bot_message = chatbot_response(message)
        chat_history.append((message, bot_message))
        return chat_history, ""

    msg.submit(respond, [msg, chatbot], [chatbot, msg])
    clear.click(lambda: None, None, chatbot, queue=False)

# 啟動介面
demo.launch()