import gradio as gr
from huggingface_hub import InferenceClient
import os
import openai

MODELS = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
    "Meta Llama 3.1 8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
    "Meta-Llama 3.1 70B-Instruct": "meta-llama/Meta-Llama-3.1-70B-Instruct",
    "Microsoft Phi-3-mini-4k": "microsoft/Phi-3-mini-4k-instruct",
    "Mixtral 8x7B": "mistralai/Mistral-7B-Instruct-v0.3",
    "Mixtral Nous-Hermes": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
    "Cohere Command R+": "CohereForAI/c4ai-command-r-plus",
    "Cohere Aya-23-35B": "CohereForAI/aya-23-35B",
    "OpenAI GPT-4o Mini": "openai/gpt-4o-mini"  # 새로운 모델 추가
}

def get_client(model_name):
    if model_name == "OpenAI GPT-4o Mini":
        return None  # OpenAI API를 사용할 것이므로 Hugging Face 클라이언트는 필요 없음
    model_id = MODELS[model_name]
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("HF_TOKEN environment variable is required")
    return InferenceClient(model_id, token=hf_token)

def call_openai_api(content, system_message, max_tokens, temperature, top_p):
    openai.api_key = os.getenv("OPENAI_API_KEY")
    response = openai.ChatCompletion.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_message},
            {"role": "user", "content": content},
        ],
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    )
    return response.choices[0].message['content']

def respond(
    message,
    chat_history,
    model_name,
    max_tokens,
    temperature,
    top_p,
    system_message,
):
    try:
        if model_name == "OpenAI GPT-4o Mini":
            assistant_message = call_openai_api(message, system_message, max_tokens, temperature, top_p)
            chat_history.append((message, assistant_message))
            yield chat_history
        else:
            client = get_client(model_name)

            messages = [{"role": "system", "content": system_message}]
            for human, assistant in chat_history:
                messages.append({"role": "user", "content": human})
                messages.append({"role": "assistant", "content": assistant})
            messages.append({"role": "user", "content": message})

            if "Cohere" in model_name:
                # Cohere 모델을 위한 비스트리밍 처리
                response = client.chat_completion(
                    messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                )
                assistant_message = response.choices[0].message.content
                chat_history.append((message, assistant_message))
                yield chat_history
            else:
                # 다른 모델들을 위한 스트리밍 처리
                stream = client.chat_completion(
                    messages,
                    max_tokens=max_tokens,
                    temperature=temperature,
                    top_p=top_p,
                    stream=True,
                )
                partial_message = ""
                for response in stream:
                    if response.choices[0].delta.content is not None:
                        partial_message += response.choices[0].delta.content
                        if len(chat_history) > 0 and chat_history[-1][0] == message:
                            chat_history[-1] = (message, partial_message)
                        else:
                            chat_history.append((message, partial_message))
                        yield chat_history
    except Exception as e:
        error_message = f"An error occurred: {str(e)}"
        chat_history.append((message, error_message))
        yield chat_history

def clear_conversation():
    return []

with gr.Blocks() as demo:
    gr.Markdown("# Prompting AI Chatbot")
    gr.Markdown("언어모델별 프롬프트 테스트 챗봇입니다.")

    with gr.Row():
        with gr.Column(scale=1):
            model_name = gr.Radio(
                choices=list(MODELS.keys()),
                label="Language Model",
                value="Zephyr 7B Beta"
            )
            max_tokens = gr.Slider(minimum=0, maximum=2000, value=500, step=100, label="Max Tokens")
            temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
            system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 최고의 비서이다.
내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )

        with gr.Column(scale=2):
            chatbot = gr.Chatbot()
            msg = gr.Textbox(label="메세지를 입력하세요")
            with gr.Row():
                submit_button = gr.Button("전송")
                clear_button = gr.Button("대화 내역 지우기")

    msg.submit(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
    submit_button.click(respond, [msg, chatbot, model_name, max_tokens, temperature, top_p, system_message], chatbot)
    clear_button.click(clear_conversation, outputs=chatbot, queue=False)

if __name__ == "__main__":
    demo.launch()