import os
import gradio as gr

# 尝试导入 groq 库
try:
    from groq import Groq
except ImportError:
    os.system('pip install groq')
    from groq import Groq

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

# 获取 API 密钥
api_key = os.getenv('groq_key')

# 定义聊天函数
def chatbot_response(message):
    completion = client.chat.completions.create(
        model="llama-3.1-70b-versatile",
        messages=[
            {
                "role": "system",
                "content": "你是一位國中地理老師，使用的語言為繁體中文(zh-tw)。你的興趣是繪製等高線圖，熱愛講冷笑話。無論學生問你什麼問題，你都會把話題引導到地理相關的討論上。"
            },
            {"role": "user", "content": 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

# 定义加载函数
def load():
    return [
        ("Here's an audio", gr.Audio("https://github.com/gradio-app/gradio/raw/main/test/test_files/audio_sample.wav")),
        ("Here's a video", gr.Video("https://github.com/gradio-app/gradio/raw/main/demo/video_component/files/world.mp4"))
    ]

# 创建 Gradio 界面
with gr.Blocks() as demo:
    chatbot_ui = gr.Chatbot()
    button = gr.Button("Load audio and video")
    
    # 按钮点击事件
    button.click(load, None, chatbot_ui)
    
    # 文本输入和响应
    message_input = gr.Textbox(placeholder="在此輸入您的問題...")
    message_output = gr.Textbox(label="Chatbot Response")

    message_input.submit(chatbot_response, message_input, message_output)

# 启动 Gradio 应用
demo.launch()