app.py

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

MODELS = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    "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": "microsoft/Phi-3-mini-4k-instruct",
    "Mixtral 8x7B": "mistralai/Mistral-7B-Instruct-v0.3",
    "Mixtral Nous-Hermes": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
    "Aya-23-35B": "CohereForAI/aya-23-35B"
}

COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

def get_client(model_name):
    hf_token = os.getenv("HF_TOKEN")
    if not hf_token:
        raise ValueError("HF_TOKEN 환경 변수가 필요합니다.")
    if model_name in MODELS:
        model_id = MODELS[model_name]
    elif model_name == "Cohere Command R+":
        model_id = COHERE_MODEL
    else:
        raise ValueError("유효하지 않은 모델 이름입니다.")
    return InferenceClient(model_id, token=hf_token)


def respond(
    message,
    chat_history,
    model_name,
    max_tokens,
    temperature,
    top_p,
    system_message,
):
    """
    기존 모델(Zephyr, Meta Llama 등) / Cohere Command R+ 모델에 대응하는 함수
    """
    try:
        client = get_client(model_name)
    except ValueError as e:
        chat_history.append((message, str(e)))
        return chat_history

    # Gradio에서 유지하던 형식
    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})

    try:
        # Cohere 모델(비스트리밍) vs 기타 모델(스트리밍) 구분
        if model_name == "Cohere Command R+":
            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))
            return 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"오류가 발생했습니다: {str(e)}"
        chat_history.append((message, error_message))
        yield chat_history


def cohere_respond(
    message,
    chat_history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    """
    Cohere Command R+ 전용 응답 함수
    """
    model_name = "Cohere Command R+"
    try:
        client = get_client(model_name)
    except ValueError as e:
        chat_history.append((message, str(e)))
        return chat_history

    messages = []
    # system_message는 messages 내 system 키로 추가하지 않고, 아래와 같이 적용해도 문제는 없음
    # 이 부분은 기존 구조를 유지
    messages.append({"role": "system", "content": system_message})

    for human, assistant in chat_history:
        if human:
            messages.append({"role": "user", "content": human})
        if assistant:
            messages.append({"role": "assistant", "content": assistant})

    messages.append({"role": "user", "content": message})

    try:
        response_full = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response_full.choices[0].message.content
        chat_history.append((message, assistant_message))
        return chat_history
    except Exception as e:
        error_message = f"오류가 발생했습니다: {str(e)}"
        chat_history.append((message, error_message))
        return chat_history


def chatgpt_respond(
    message,
    chat_history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    """
    ChatGPT (OpenAI) 전용 응답 함수
    """
    openai.api_key = os.getenv("OPENAI_API_KEY")
    if not openai.api_key:
        chat_history.append((message, "OPENAI_API_KEY 환경 변수가 필요합니다."))
        return chat_history

    # ChatGPT는 system role을 messages 내에서 허용하므로 현재 방식 OK
    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})

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4",  # 적절한 모델 ID를 사용
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response.choices[0].message['content']
        chat_history.append((message, assistant_message))
        return chat_history
    except Exception as e:
        error_message = f"오류가 발생했습니다: {str(e)}"
        chat_history.append((message, error_message))
        return chat_history


def claude_respond(
    message,
    chat_history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    """
    Anthropic Claude 전용 응답 함수 (Messages API)
    - "system" role은 messages 목록에 넣지 않고,
      최상위 system 인자로 넘겨야 한다.
    """
    anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
    if not anthropic_api_key:
        chat_history.append((message, "ANTHROPIC_API_KEY 환경 변수가 필요합니다."))
        return chat_history

    client = anthropic.Anthropic(api_key=anthropic_api_key)

    # Anthropic Messages API에 맞게 "system" role은 top-level 파라미터로,
    # 나머지는 user/assistant만 messages에 넣음
    # (messages 내 "system" role이 있으면 오류 발생)
    anthro_messages = []
    for human, assistant in chat_history:
        if human:
            anthro_messages.append({"role": "user", "content": human})
        if assistant:
            anthro_messages.append({"role": "assistant", "content": assistant})
    anthro_messages.append({"role": "user", "content": message})

    try:
        response = client.messages.create(
            model="claude-3-haiku-20240307",  # Claude 모델 예시 (사용 가능 모델 확인)
            system=system_message,         # ← 여기서 system_message를 전달
            messages=anthro_messages,      # user/assistant만 포함
            max_tokens_to_sample=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response["completion"]
        chat_history.append((message, assistant_message))
        return chat_history
    except Exception as e:
        error_message = f"오류가 발생했습니다: {str(e)}"
        chat_history.append((message, error_message))
        return chat_history


def clear_conversation():
    return []


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

    with gr.Tab("일반 모델"):
        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)

    with gr.Tab("Cohere Command R+"):
        with gr.Row():
            cohere_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 최고의 비서이다.
내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            cohere_max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
            cohere_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
            cohere_top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-P",
            )

        cohere_chatbot = gr.Chatbot(height=600)
        cohere_msg = gr.Textbox(label="메세지를 입력하세요")
        with gr.Row():
            cohere_submit_button = gr.Button("전송")
            cohere_clear_button = gr.Button("대화 내역 지우기")
        
        cohere_msg.submit(
            cohere_respond,
            [cohere_msg, cohere_chatbot, cohere_system_message, cohere_max_tokens, cohere_temperature, cohere_top_p],
            cohere_chatbot
        )
        cohere_submit_button.click(
            cohere_respond,
            [cohere_msg, cohere_chatbot, cohere_system_message, cohere_max_tokens, cohere_temperature, cohere_top_p],
            cohere_chatbot
        )
        cohere_clear_button.click(clear_conversation, outputs=cohere_chatbot, queue=False)

    with gr.Tab("ChatGPT"):
        with gr.Row():
            chatgpt_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 ChatGPT, OpenAI에서 개발한 언어 모델이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            chatgpt_max_tokens = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max Tokens")
            chatgpt_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            chatgpt_top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-P",
            )

        chatgpt_chatbot = gr.Chatbot(height=600)
        chatgpt_msg = gr.Textbox(label="메세지를 입력하세요")
        with gr.Row():
            chatgpt_submit_button = gr.Button("전송")
            chatgpt_clear_button = gr.Button("대화 내역 지우기")
        
        chatgpt_msg.submit(
            chatgpt_respond,
            [chatgpt_msg, chatgpt_chatbot, chatgpt_system_message, chatgpt_max_tokens, chatgpt_temperature, chatgpt_top_p],
            chatgpt_chatbot
        )
        chatgpt_submit_button.click(
            chatgpt_respond,
            [chatgpt_msg, chatgpt_chatbot, chatgpt_system_message, chatgpt_max_tokens, chatgpt_temperature, chatgpt_top_p],
            chatgpt_chatbot
        )
        chatgpt_clear_button.click(clear_conversation, outputs=chatgpt_chatbot, queue=False)

    with gr.Tab("Claude"):
        with gr.Row():
            claude_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 Claude, Anthropic에서 개발한 언어 모델이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            claude_max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Tokens")
            claude_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            claude_top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-P",
            )

        claude_chatbot = gr.Chatbot(height=600)
        claude_msg = gr.Textbox(label="메세지를 입력하세요")
        with gr.Row():
            claude_submit_button = gr.Button("전송")
            claude_clear_button = gr.Button("대화 내역 지우기")

        # Claude 전용 함수 호출
        claude_msg.submit(
            claude_respond,
            [claude_msg, claude_chatbot, claude_system_message, claude_max_tokens, claude_temperature, claude_top_p],
            claude_chatbot
        )
        claude_submit_button.click(
            claude_respond,
            [claude_msg, claude_chatbot, claude_system_message, claude_max_tokens, claude_temperature, claude_top_p],
            claude_chatbot
        )
        claude_clear_button.click(clear_conversation, outputs=claude_chatbot, queue=False)

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