from openai import OpenAI
import gradio as gr
import os

# OpenAI APIキーの取得
api_key = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=api_key)

def chat_with_openai(user_input, history):
    if history is None:
        history = []

    # 最初のメッセージの場合、システムメッセージの代わりにユーザーメッセージとして追加
    if len(history) == 0:
        history.append({"role": "user", "content": "You are a helpful assistant."})

    # ユーザーの入力を履歴に追加
    history.append({"role": "user", "content": user_input})

    # OpenAI APIへのリクエスト
    completion = client.chat.completions.create(
        model="o1-preview",
        messages=history
    )

    # アシスタントの応答を取得
    assistant_message = completion.choices[0].message.content

    # 履歴にアシスタントの応答を追加
    history.append({"role": "assistant", "content": assistant_message})

    # チャットボットに表示するメッセージを整形
    messages_to_display = []
    for i in range(1, len(history)-1, 2):
        user_msg = history[i]["content"]
        assistant_msg = history[i+1]["content"]
        messages_to_display.append((user_msg, assistant_msg))

    return messages_to_display, history

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    state = gr.State([])

    with gr.Row():
        txt = gr.Textbox(
            show_label=False,
            placeholder="メッセージを入力してください..."
        )
        btn = gr.Button("送信")

    def on_submit(user_input, history):
        messages, updated_history = chat_with_openai(user_input, history)
        return messages, updated_history, ""

    txt.submit(on_submit, [txt, state], [chatbot, state, txt])
    btn.click(on_submit, [txt, state], [chatbot, state, txt])

demo.launch()