import gradio as gr
import pandas as pd
import os
from openai import OpenAI

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY)

def process_file(file):
    # 读取文件
    if file.name.endswith('.csv'):
        df = pd.read_csv(file)
    else:
        df = pd.read_excel(file)
    # 将 DataFrame 转换为字符串
    return df.to_string()

def respond(user_message, df_string, chat_history):
    sys_content = f"你是一个资料分析师，以下是数据分析资料:\n{df_string}"
    messages = chat_history if chat_history is not None else []
    messages.append({"role": "system", "content": sys_content})
    messages.append({"role": "user", "content": user_message})

    request_payload = {
        "model": "gpt-4-1106-preview",
        "messages": messages,
        "max_tokens": 2000
    }

    response = client.chat.completions.create(**request_payload)
    response_text = response.choices[0].message.content.strip()

    messages.append({"role": "assistant", "content": response_text})

    return "", messages

with gr.Blocks() as demo:
    with gr.Row():
        file_upload = gr.File(label="Upload your file")
        df_string_output = gr.Variable()

    with gr.Row():
        chatbot = gr.Chatbot()

    with gr.Row():
        msg = gr.Textbox(label="请输入对话内容")
        send_button = gr.Button("发送")

    file_upload.change(process_file, inputs=file_upload, outputs=df_string_output)

    send_button.click(
        respond, 
        inputs=[msg, df_string_output, chatbot], 
        outputs=[msg, chatbot]
    )

demo.launch()