import os
import gradio as gr
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq

# Initialize Groq client
client = Groq(api_key=os.getenv("groq_key"))

# Chat history to store conversation
chat_history = []

def chat(message, history):
    messages = [
        {
            "role": "system", 
            "content": "你是一個美術老師，請用繁體中文稱讚學生的問題，問什麼問題都會引導到藝術"
        }
    ]
    
    # Add chat history
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    
    # Add current message
    messages.append({"role": "user", "content": message})

    # Get response from Groq
    completion = client.chat.completions.create(
        model="llama3-8b-8192",
        messages=messages,
        temperature=1,
        max_tokens=1024,
        top_p=1,
        stream=True,
        stop=None
    )

    # Build response text from stream
    response = ""
    for chunk in completion:
        if chunk.choices[0].delta.content is not None:
            response += chunk.choices[0].delta.content
            yield response

# Create Gradio interface
demo = gr.ChatInterface(
    fn=chat,
    title="藝術老師聊天機器人",
    description="我是一位熱愛藝術的老師，讓我們來聊聊藝術吧！",
    examples=["你好", "我想學畫畫", "什麼是藝術"],
    retry_btn=None,
    undo_btn=None,
    clear_btn="清除對話",
)

if __name__ == "__main__":
    demo.launch()