Update app.py

app.py

@@ -4,94 +4,24 @@ import openai
 import anthropic
 import os
 
-# 제거할 모델들을 MODELS 사전에서 제외
-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",
-    "DeepSeek-V3": "deepseek/deepseek-chat"
-}
-
 # Cohere Command R+ 모델 ID 정의
 COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
 
-def get_client(model_name, hf_token):
+def get_client(model_name: str):
     """
     모델 이름에 맞춰 InferenceClient 생성.
-    hf_token을 UI에서 입력받은 값으로 사용하도록 변경.
+    HuggingFace 토큰은 os.environ.get("HF_TOKEN")을 통해 환경변수로 가져온다.
     """
+    hf_token = os.environ.get("HF_TOKEN")
     if not hf_token:
-        raise ValueError("HuggingFace API 토큰이 필요합니다.")
+        raise ValueError("HuggingFace API 토큰이 필요합니다. (환경변수 HF_TOKEN 미설정)")
 
-    if model_name in MODELS:
-        model_id = MODELS[model_name]
-    elif model_name == "Cohere Command R+":
+    if 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,
-    hf_token,       # HF 토큰을 추가로 받음
-):
-    try:
-        client = get_client(model_name, hf_token)
-    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:
-        messages.append({"role": "user", "content": human})
-        messages.append({"role": "assistant", "content": assistant})
-    messages.append({"role": "user", "content": message})
-
-    try:
-        if model_name == "Cohere Command R+":
-            # 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,
@@ -99,11 +29,14 @@ def cohere_respond(
     max_tokens,
     temperature,
     top_p,
-    hf_token,   # HF 토큰 추가
 ):
+    """
+    Cohere Command R+ 모델 응답 함수.
+    HF 토큰은 함수 내부에서 os.environ을 통해 불러온다.
+    """
     model_name = "Cohere Command R+"
     try:
-        client = get_client(model_name, hf_token)
+        client = get_client(model_name)
     except ValueError as e:
         chat_history.append((message, str(e)))
         return chat_history
@@ -118,7 +51,6 @@ def cohere_respond(
     messages.append({"role": "user", "content": message})
 
     try:
-        # Cohere Command R+ 모델을 위한 비스트리밍 처리
         response_full = client.chat_completion(
             messages,
             max_tokens=max_tokens,
@@ -140,16 +72,17 @@ def chatgpt_respond(
     max_tokens,
     temperature,
     top_p,
-    openai_token,  # openai 토큰 추가
 ):
     """
-    chatgpt용 응답. openai_token을 UI에서 입력받아 사용하도록 변경.
+    ChatGPT 모델 응답 함수.
+    OpenAI 토큰은 함수 내부에서 os.environ을 통해 불러온다.
     """
+    openai_token = os.environ.get("OPENAI_TOKEN")
     if not openai_token:
-        chat_history.append((message, "OpenAI API 토큰이 필요합니다."))
+        chat_history.append((message, "OpenAI API 토큰이 필요합니다. (환경변수 OPENAI_TOKEN 미설정)"))
         return chat_history
 
-    openai.api_key = openai_token  # UI에서 받은 토큰 사용
+    openai.api_key = openai_token  # 환경변수에서 받은 토큰 사용
 
     messages = [{"role": "system", "content": system_message}]
     for human, assistant in chat_history:
@@ -180,19 +113,19 @@ def claude_respond(
     max_tokens,
     temperature,
     top_p,
-    claude_token,  # Claude 토큰 추가
 ):
     """
-    Claude용 응답. claude_token을 UI에서 입력받아 사용하도록 변경.
+    Claude 모델 응답 함수.
+    Claude 토큰은 함수 내부에서 os.environ을 통해 불러온다.
     """
+    claude_token = os.environ.get("CLAUDE_TOKEN")
     if not claude_token:
-        chat_history.append((message, "Claude API 토큰이 필요합니다."))
+        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,
@@ -221,13 +154,14 @@ def deepseek_respond(
     max_tokens,
     temperature,
     top_p,
-    deepseek_token,  # DeepSeek 토큰 추가
 ):
     """
-    DeepSeek용 응답. deepseek_token을 UI에서 입력받아 사용하도록 변경.
+    DeepSeek 모델 응답 함수.
+    DeepSeek 토큰은 함수 내부에서 os.environ을 통해 불러온다.
     """
+    deepseek_token = os.environ.get("DEEPSEEK_TOKEN")
     if not deepseek_token:
-        chat_history.append((message, "DeepSeek API 토큰이 필요합니다."))
+        chat_history.append((message, "DeepSeek API 토큰이 필요합니다. (환경변수 DEEPSEEK_TOKEN 미설정)"))
         return chat_history
 
     openai.api_key = deepseek_token
@@ -258,75 +192,18 @@ def deepseek_respond(
 def clear_conversation():
     return []
 
+# --------------------------------------------
+# Gradio 앱 시작
+# --------------------------------------------
 with gr.Blocks() as demo:
     gr.Markdown("# Prompting AI Chatbot")
     gr.Markdown("언어모델별 프롬프트 테스트 챗봇입니다.")
 
-    # --- 토큰 입력 UI 추가 ---
-    with gr.Row():
-        hf_token_box = gr.Textbox(
-            label="HuggingFace 토큰",
-            type="password",
-            placeholder="HuggingFace API 토큰을 입력하세요..."
-        )
-        openai_token_box = gr.Textbox(
-            label="OpenAI 토큰",
-            type="password",
-            placeholder="OpenAI API 토큰을 입력하세요..."
-        )
-        claude_token_box = gr.Textbox(
-            label="Claude 토큰",
-            type="password",
-            placeholder="Claude API 토큰을 입력하세요..."
-        )
-        deepseek_token_box = gr.Textbox(
-            label="DeepSeek 토큰",
-            type="password",
-            placeholder="DeepSeek API 토큰을 입력하세요..."
-        )
+    # --------------------------------------------------
+    # 일반 모델 관련 UI/기능 제거 (요청 사항에 따라 삭제)
+    # --------------------------------------------------
 
-    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=100, maximum=8000, value=2000, 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("대화 내역 지우기")
-    
-        # respond 함수에 hf_token 인자를 전달하도록 수정
-        inputs_for_normal = [
-            msg, 
-            chatbot, 
-            model_name, 
-            max_tokens, 
-            temperature, 
-            top_p, 
-            system_message,
-            hf_token_box
-        ]
-        msg.submit(respond, inputs_for_normal, chatbot)
-        submit_button.click(respond, inputs_for_normal, chatbot)
-        clear_button.click(clear_conversation, outputs=chatbot, queue=False)
-    
+    # Cohere Command R+
     with gr.Tab("Cohere Command R+"):
         with gr.Row():
             cohere_system_message = gr.Textbox(
@@ -353,20 +230,19 @@ with gr.Blocks() as demo:
             cohere_submit_button = gr.Button("전송")
             cohere_clear_button = gr.Button("대화 내역 지우기")
         
-        # cohere_respond 함수에 hf_token 인자를 전달하도록 수정
         inputs_for_cohere = [
             cohere_msg, 
             cohere_chatbot, 
             cohere_system_message, 
             cohere_max_tokens, 
             cohere_temperature, 
-            cohere_top_p,
-            hf_token_box
+            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(
@@ -393,20 +269,19 @@ with gr.Blocks() as demo:
             chatgpt_submit_button = gr.Button("전송")
             chatgpt_clear_button = gr.Button("대화 내역 지우기")
         
-        # chatgpt_respond 함수에 openai_token 인자를 전달하도록 수정
         inputs_for_chatgpt = [
-            chatgpt_msg, 
-            chatgpt_chatbot, 
-            chatgpt_system_message, 
-            chatgpt_max_tokens, 
-            chatgpt_temperature, 
-            chatgpt_top_p,
-            openai_token_box
+            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(
@@ -433,20 +308,19 @@ with gr.Blocks() as demo:
             claude_submit_button = gr.Button("전송")
             claude_clear_button = gr.Button("대화 내역 지우기")
         
-        # claude_respond 함수에 claude_token 인자를 전달하도록 수정
         inputs_for_claude = [
             claude_msg, 
             claude_chatbot, 
             claude_system_message, 
             claude_max_tokens, 
             claude_temperature, 
-            claude_top_p,
-            claude_token_box
+            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(
@@ -473,19 +347,17 @@ with gr.Blocks() as demo:
             deepseek_submit_button = gr.Button("전송")
             deepseek_clear_button = gr.Button("대화 내역 지우기")
         
-        # deepseek_respond 함수에 deepseek_token 인자를 전달하도록 수정
         inputs_for_deepseek = [
-            deepseek_msg, 
-            deepseek_chatbot, 
-            deepseek_system_message, 
-            deepseek_max_tokens, 
-            deepseek_temperature, 
-            deepseek_top_p,
-            deepseek_token_box
+            deepseek_msg,
+            deepseek_chatbot,
+            deepseek_system_message,
+            deepseek_max_tokens,
+            deepseek_temperature,
+            deepseek_top_p
         ]
         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()
+    demo.launch()