app.py

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

# Cohere Command R+ 모델 ID 정의
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

def get_client(model_name: str):
    """
    모델 이름에 맞춰 InferenceClient 생성.
    HuggingFace 토큰은 os.environ.get("HF_TOKEN")을 통해 환경변수로 가져온다.
    """
    hf_token = os.environ.get("HF_TOKEN")
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다. (환경변수 HF_TOKEN 미설정)")

    if model_name == "Cohere Command R+":
        model_id = COHERE_MODEL
    else:
        raise ValueError("유효하지 않은 모델 이름입니다.")
    return InferenceClient(model_id, token=hf_token)

def cohere_respond(
    message,
    chat_history,
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    """
    Cohere Command R+ 모델 응답 함수.
    HF 토큰은 함수 내부에서 os.environ을 통해 불러온다.
    """
    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 = [{"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 토큰은 함수 내부에서 os.environ을 통해 불러온다.
    """
    openai_token = os.environ.get("OPENAI_TOKEN")
    if not openai_token:
        chat_history.append((message, "OpenAI API 토큰이 필요합니다. (환경변수 OPENAI_TOKEN 미설정)"))
        return chat_history

    openai.api_key = openai_token  # 환경변수에서 받은 토큰 사용

    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-4o-mini",  # 또는 다른 모델 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,
):
    """
    Claude 모델 응답 함수.
    Claude 토큰은 함수 내부에서 os.environ을 통해 불러온다.
    """
    claude_token = os.environ.get("CLAUDE_TOKEN")
    if not claude_token:
        chat_history.append((message, "Claude API 토큰이 필요합니다. (환경변수 CLAUDE_TOKEN 미설정)"))
        return chat_history

    try:
        client = anthropic.Anthropic(api_key=claude_token)
        
        response = client.messages.create(
            model="claude-3-haiku-20240307",
            max_tokens=max_tokens,
            temperature=temperature,
            system=system_message,
            messages=[
                {
                    "role": "user",
                    "content": message
                }
            ]
        )
        
        assistant_message = response.content[0].text
        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 deepseek_respond(
    message,
    chat_history,
    system_message,
    deepseek_model_choice,
    max_tokens,
    temperature,
    top_p,
):
    """
    DeepSeek 모델 응답 함수.
    DeepSeek 토큰은 함수 내부에서 os.environ을 통해 불러온다.
    deepseek_model_choice에 따라 deepseek-chat 또는 deepseek-reasoner를 선택하며,
    스트리밍 방식으로 응답을 받아옵니다.
    """
    deepseek_token = os.environ.get("DEEPSEEK_TOKEN")
    if not deepseek_token:
        chat_history.append((message, "DeepSeek API 토큰이 필요합니다. (환경변수 DEEPSEEK_TOKEN 미설정)"))
        yield chat_history
        return

    openai.api_key = deepseek_token
    openai.api_base = "https://api.deepseek.com/v1"

    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})

    # 모델 선택: 기본은 deepseek-chat
    if deepseek_model_choice == "R1(deepseek-reasoner)":
        model = "deepseek-reasoner"
    else:
        model = "deepseek-chat"

    try:
        response = openai.ChatCompletion.create(
            model=model,
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=True
        )
        assistant_message = ""
        # 새로운 대화 항목을 추가하고 초기값을 스트리밍하면서 갱신
        chat_history.append((message, assistant_message))
        yield chat_history
        for chunk in response:
            # "content"가 None인 경우 빈 문자열로 처리하여 오류 방지
            delta = chunk.choices[0].delta.get("content") or ""
            assistant_message += delta
            chat_history[-1] = (message, assistant_message)
            yield chat_history
        return
    except Exception as e:
        error_message = f"오류가 발생했습니다: {str(e)}"
        chat_history.append((message, error_message))
        yield chat_history
        return

def clear_conversation():
    return []

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

    # --------------------------------------------------
    # 일반 모델 관련 UI/기능 제거 (요청 사항에 따라 삭제)
    # --------------------------------------------------

    # Cohere Command R+
    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=100, maximum=8000, value=2000, step=100, label="Max new tokens")
            cohere_temperature = gr.Slider(minimum=0.1, maximum=2.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("대화 내역 지우기")
        
        inputs_for_cohere = [
            cohere_msg, 
            cohere_chatbot, 
            cohere_system_message, 
            cohere_max_tokens, 
            cohere_temperature, 
            cohere_top_p
        ]
        cohere_msg.submit(cohere_respond, inputs_for_cohere, cohere_chatbot)
        cohere_submit_button.click(cohere_respond, inputs_for_cohere, cohere_chatbot)
        cohere_clear_button.click(clear_conversation, outputs=cohere_chatbot, queue=False)

    # ChatGPT
    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=100, maximum=5000, value=2000, step=100, 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("대화 내역 지우기")
        
        inputs_for_chatgpt = [
            chatgpt_msg,
            chatgpt_chatbot,
            chatgpt_system_message,
            chatgpt_max_tokens,
            chatgpt_temperature,
            chatgpt_top_p
        ]
        chatgpt_msg.submit(chatgpt_respond, inputs_for_chatgpt, chatgpt_chatbot)
        chatgpt_submit_button.click(chatgpt_respond, inputs_for_chatgpt, chatgpt_chatbot)
        chatgpt_clear_button.click(clear_conversation, outputs=chatgpt_chatbot, queue=False)

    # Claude
    with gr.Tab("Claude"):
        with gr.Row():
            claude_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 Anthropic에서 개발한 클로드이다.
최대한 정확하고 친절하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            claude_max_tokens = gr.Slider(minimum=100, maximum=8000, value=2000, step=100, 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("대화 내역 지우기")
        
        inputs_for_claude = [
            claude_msg, 
            claude_chatbot, 
            claude_system_message, 
            claude_max_tokens, 
            claude_temperature, 
            claude_top_p
        ]
        claude_msg.submit(claude_respond, inputs_for_claude, claude_chatbot)
        claude_submit_button.click(claude_respond, inputs_for_claude, claude_chatbot)
        claude_clear_button.click(clear_conversation, outputs=claude_chatbot, queue=False)

    # DeepSeek
    with gr.Tab("DeepSeek"):
        with gr.Row():
            deepseek_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 DeepSeek-V3, 최고의 언어 모델이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            deepseek_model_choice = gr.Radio(
                choices=["V3(deepseek-chat)", "R1(deepseek-reasoner)"],
                value="V3(deepseek-chat)",
                label="모델 선택"
            )
            deepseek_max_tokens = gr.Slider(minimum=100, maximum=8000, value=2000, step=100, label="Max Tokens")
            deepseek_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            deepseek_top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-P",
            )
        
        deepseek_chatbot = gr.Chatbot(height=600)
        deepseek_msg = gr.Textbox(label="메세지를 입력하세요")
        with gr.Row():
            deepseek_submit_button = gr.Button("전송")
            deepseek_clear_button = gr.Button("대화 내역 지우기")
        
        inputs_for_deepseek = [
            deepseek_msg,
            deepseek_chatbot,
            deepseek_system_message,
            deepseek_model_choice,
            deepseek_max_tokens,
            deepseek_temperature,
            deepseek_top_p
        ]
        # Textbox.submit에서는 stream 인자를 제거합니다.
        deepseek_msg.submit(deepseek_respond, inputs_for_deepseek, deepseek_chatbot)
        deepseek_submit_button.click(deepseek_respond, inputs_for_deepseek, deepseek_chatbot)
        deepseek_clear_button.click(clear_conversation, outputs=deepseek_chatbot, queue=False)

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