import os
import gradio as gr

# 確保 groq 套件可用
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq

# 設定環境變數
os.environ["GROQ_API_KEY"] = os.getenv("groq_key")  # 確保這個環境變數已設置

# 初始化 Groq 客戶端
client = Groq()

# 定義 chatbot 回應的函數
def chat_with_groq(user_input):
    completion = client.chat.completions.create(
        model="llama-3.1-70b-versatile",
        messages=[
            {
                "role": "system",
                "content": "我是國文老師，專門批改國小作文"
            },
            {
                "role": "user",
                "content": user_input
            }
        ],
        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 創建 chatbot
iface = gr.Interface(
    fn=chat_with_groq,
    inputs=gr.inputs.Textbox(label="請輸入您的問題："),
    outputs=gr.outputs.Textbox(label="回應："),
    title="國文老師 Chatbot",
    description="這是一個專門批改國小作文的 Chatbot"
)

# 啟動應用
iface.launch()