import gradio as gr
import os

# 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

# Set up the Groq client with the secret key
groq_key = os.getenv('groq_key')
if not groq_key:
    raise ValueError("groq_key environment variable is not set")

client = Groq(api_key=groq_key)

# 定义聊天机器人的响应函数
def chatbot_response(messages):
    try:
        completion = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {
                    "role": "system",
                    "content": "You are a corporate secretary who is skilled at drafting business emails. The prompt will feed you addressee, main message, and final greetings."
                }
            ] + messages,
            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
    except Exception as e:
        print(f"Error in chatbot_response: {e}")
        return "Error: Unable to get response from the API."

# 使用 gradio 创建聊天机器人界面
def respond(message, history):
    try:
        history.append(("user", message))
        bot_response = chatbot_response([{"role": "user", "content": message}])
        history.append(("assistant", bot_response))
        return "", history
    except Exception as e:
        print(f"Error in respond function: {e}")
        return "", history

# 初始化 Gradio 应用
with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    
    msg.submit(respond, [msg, chatbot], [msg, chatbot])

# 运行应用程序
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)