import gradio as gr
import requests

# 设置第三方 API 基本 URL
API_BASE_URL = "http://key.aistory.uk/v1/chat/completions"  # 替换为你自己的API URL
API_KEY = "sk-HfD4NYIN6bq2DkSfIiUcciRvo9MkgMdFCsahP9NWEOUPHe8H"  # 替换为你自己的 API 密钥



# 定义 AI 响应函数，调用第三方 API
def ai_response(message, chat_history):
    # 定义系统提示词
    system_prompt = "You are a helpful assistant. Please assist the user with their inquiries."

    # 组合历史聊天记录和用户输入的信息
    conversation = [{"role": "system", "content": system_prompt}]
    for msg in chat_history:
        conversation.append({"role": msg[0], "content": msg[1]})
    conversation.append({"role": "user", "content": message})

    # 构建请求体
    payload = {
        "model": "gpt-4o",  # 使用 gpt-4o 模型（如果此模型为该 API 支持的模型）
        "messages": conversation,
        "max_tokens": 850
    }

    # 设置请求头，包括 API 密钥
    headers = {
        "Authorization": f"Bearer {API_KEY}",
        "Content-Type": "application/json"
    }

    # 发送请求到第三方 API
    try:
        response = requests.post(API_BASE_URL, json=payload, headers=headers)
        response.raise_for_status()  # 如果响应状态码不是 2xx，会抛出异常
        if response.status_code == 200:
            # 获取 API 响应内容
            response_data = response.json()
            assistant_message = response_data['choices'][0]['message']['content']

            # 返回新的聊天记录，转换为符合 gr.Chatbot 期望的元组格式
            chat_history.append(("user", message))
            chat_history.append(("assistant", assistant_message))

            return chat_history
        else:
            # 如果请求失败，输出错误信息
            return chat_history + [("assistant", f"API error: {response.status_code}, {response.text}")]
    except requests.exceptions.RequestException as e:
        # 捕获任何请求错误，并输出详细错误信息
        return chat_history + [("assistant", f"Request failed: {str(e)}")]

# 创建 Gradio 应用
def create_interface():
    with gr.Blocks() as demo:
        # 创建一个 Column 布局，用于将聊天记录和输入框放在同一列
        with gr.Column():
            # 创建一个聊天机器人输出组件，用于显示对话
            chat_output = gr.Chatbot()

            # 创建一个文本框用于输入消息
            message_input = gr.Textbox(label="请输入你的问题", placeholder="输入你的问题并按回车发送", lines=1)

            # 提交按钮，发送用户消息并获取AI回复
            message_input.submit(ai_response, inputs=[message_input, chat_output], outputs=[chat_output])

    return demo

# 启动 Gradio 应用
demo = create_interface()
demo.launch()