import gradio as gr
import os

from gradio import ChatMessage
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq
# 從環境變數中獲取 API 密鑰
api_key = os.getenv("groq_key")

# 初始化 groq 客戶端
client = Groq()
client.set_api_key(api_key)

# 定義回應函數
def respond(message, chat_history):
    # 構建聊天記錄
    messages = [{"role": "system", "content": "你是一個友善的 AI 助手，請幫助使用者解答問題。"}]
    for entry in chat_history:
        if entry['role'] == 'user':
            messages.append({"role": "user", "content": entry['content']})
        else:
            messages.append({"role": "assistant", "content": entry['content']})
    
    # 添加當前用戶消息
    messages.append({"role": "user", "content": message})

    # 獲取回應
    completion = client.chat.completions.create(
        model="llama-3.1-70b-versatile",
        messages=messages,
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=True,
        stop=None,
    )

    reply = ""
    for chunk in completion:
        reply += chunk.choices[0].delta.content or ""
    
    # 添加當前回應到聊天記錄
    chat_history.append(ChatMessage(role="user", content=message))
    chat_history.append(ChatMessage(role="assistant", content=reply))
    return chat_history

# 創建 Gradio Chatbot 界面
with gr.Blocks() as demo:
    chatbot = gr.Chatbot(type='messages')
    msg = gr.Textbox(label="輸入訊息")
    clear = gr.Button("清除聊天")

    def user(user_message, history):
        return "", history + [ChatMessage(role="user", content=user_message)]

    def bot(history):
        user_message = history[-1]['content']
        history = respond(user_message, history[:-1])
        return history

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

demo.launch()