import openai
import gradio as gr
import os

# 设置您的API密钥
openai.api_key = os.environ["chat_key"]

# 存储用户对话历史的字典
user_dialogue_histories = {}
max_tokens_per_user = 2000  # 您可以根据需要设置此值

def get_total_tokens(dialogue_history):
    total_tokens = 0
    for message in dialogue_history:
        total_tokens += len(message["content"])
    return total_tokens

def remove_earliest_messages(user_id, tokens_to_remove):
    while tokens_to_remove > 0 and user_dialogue_histories[user_id]:
        removed_message = user_dialogue_histories[user_id].pop(0)
        tokens_to_remove -= len(removed_message["content"])

def chat_with_chatgpt(user_id, user_message):
    # 如果用户不在字典中，添加一个新的对话历史列表
    if user_id not in user_dialogue_histories:
        user_dialogue_histories[user_id] = []

    # 将用户的消息添加到对话历史中
    user_dialogue_histories[user_id].append({"role": "user", "content": user_message})

    # 如果添加用户消息后对话历史中的token数超过限制，删除最早的消息，直到满足限制
    if get_total_tokens(user_dialogue_histories[user_id]) > max_tokens_per_user:
        remove_earliest_messages(user_id, get_total_tokens(user_dialogue_histories[user_id]) - max_tokens_per_user)

    # 调用ChatGPT API
    response = openai.Completion.create(
        engine="gpt-3.5-turbo",
        prompt=[{"role": "system", "content": "You are a helpful assistant."}] + user_dialogue_histories[user_id],
        max_tokens=150,
        n=1,
        stop=None,
        temperature=0.5,
    )

    chatgpt_response = response.choices[0].text.strip()

    # 将ChatGPT的回应添加到对话历史中
    user_dialogue_histories[user_id].append({"role": "assistant", "content": chatgpt_response})

    # 如果添加ChatGPT的回应后对话历史中的token数超过限制，删除最早的消息，直到满足限制
    if get_total_tokens(user_dialogue_histories[user_id]) > max_tokens_per_user:
        remove_earliest_messages(user_id, get_total_tokens(user_dialogue_histories[user_id]) - max_tokens_per_user)

    return chatgpt_response

# Gradio界面
def gradio_interface(user_id, user_message):
    response = chat_with_chatgpt(user_id, user_message)
    return response

inputs = [
    gr.inputs.Textbox(label="User ID", placeholder="Enter user ID here"),
    gr.inputs.Textbox(label="Message", placeholder="Enter your message here"),
]

output = gr.outputs.Textbox(label="ChatGPT Response")

gr.Interface(fn=gradio_interface, inputs=inputs, outputs=output, title="Chat with ChatGPT").launch()