import openai
import gradio as gr

# 初始化 OpenAI API 客户端
client = openai.OpenAI(
    api_key="sk-hm35RR1E0dfB26C8873BT3BlBKFJE681B3d87a6c4B3e8C44",
    base_url="https://aigptx.top/"
)

# 默认系统消息（可以通过UI定制）
system_prompt = "您好！欢迎来到聊天机器人。请随时提出您的问题或想法。"

# 添加用户消息到聊天记录
def add_text(history, text, custom_prompt):
    if not custom_prompt:
        custom_prompt = system_prompt  # 使用默认系统提示
    if text:
        if not history:
            history = [("System Message", custom_prompt)]
        history.append((text, None))
    else:
        history.append(("System Message", "您发送了一个空消息，请输入您的问题或想法。"))
    return history, ""

# 调用 OpenAI GPT 生成回复
def predict(history, custom_prompt):
    if not custom_prompt:
        custom_prompt = system_prompt  # 使用默认系统提示
    if not history or (len(history) == 1 and history[0][0] == "System Message"):
        history = [("System Message", custom_prompt)]
        yield history
        return history, ""

    history_ = history[1:] if history[0][0] == "System Message" else history

    history_openai_format = [{"role": "user", "content": human} for human, assistant in history_ if human]

    if not history_openai_format:
        yield history
        return history

    response = client.chat.completions.create(
        model='gpt-4-1106-preview',
        messages=[{"role": "system", "content": custom_prompt}] + history_openai_format,
        temperature=0.9,
    )

    history[-1][1] = response.choices[0].message.content
    yield history

# 构建 Gradio 界面
with gr.Blocks() as demo:
    gr.Markdown("<h1><center>聊天机器人</center></h1>")

    with gr.Row():
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                value=[("System Message", system_prompt)],
            )
            with gr.Row():
                textbox = gr.Textbox(placeholder="请输入您的问题或想法", show_label=False, scale=2)
                submit = gr.Button("💬 发送", scale=1)

                # 自定义系统消息输入框
                custom_system_prompt = gr.Textbox(
                    value=system_prompt,
                    placeholder="自定义系统消息，例如：'我是一个乐于助人的助手，请告诉我如何帮您'",
                    label="自定义系统提示",
                    lines=2
                )

                # 在用户输入消息时，传递自定义系统提示给 GPT 模型
                textbox_msg = textbox.submit(
                    add_text, 
                    [chatbot, textbox, custom_system_prompt], 
                    [chatbot, textbox], 
                    queue=False
                ).then(
                    predict, 
                    [chatbot, custom_system_prompt], 
                    chatbot
                )

                submit.click(
                    add_text, 
                    inputs=[chatbot, textbox, custom_system_prompt], 
                    outputs=[chatbot, textbox], 
                    queue=False
                ).then(
                    predict, 
                    [chatbot, custom_system_prompt], 
                    chatbot
                )

demo.queue()
demo.launch()