Update app.py

app.py

@@ -17,10 +17,15 @@ MODELS = {
 # Cohere Command R+ 모델 ID 정의
 COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"
 
-def get_client(model_name):
-    hf_token = os.getenv("HF_TOKEN")
+
+def get_client(model_name, hf_token):
+    """
+    모델 이름에 맞춰 InferenceClient 생성.
+    hf_token을 UI에서 입력받은 값으로 사용하도록 변경.
+    """
     if not hf_token:
-        raise ValueError("HF_TOKEN 환경 변수가 필요합니다.")
+        raise ValueError("HuggingFace API 토큰이 필요합니다.")
+
     if model_name in MODELS:
         model_id = MODELS[model_name]
     elif model_name == "Cohere Command R+":
@@ -29,6 +34,7 @@ def get_client(model_name):
         raise ValueError("유효하지 않은 모델 이름입니다.")
     return InferenceClient(model_id, token=hf_token)
 
+
 def respond(
     message,
     chat_history,
@@ -37,9 +43,10 @@ def respond(
     temperature,
     top_p,
     system_message,
+    hf_token,       # HF 토큰을 추가로 받음
 ):
     try:
-        client = get_client(model_name)
+        client = get_client(model_name, hf_token)
     except ValueError as e:
         chat_history.append((message, str(e)))
         return chat_history
@@ -85,6 +92,7 @@ def respond(
         chat_history.append((message, error_message))
         yield chat_history
 
+
 def cohere_respond(
     message,
     chat_history,
@@ -92,10 +100,11 @@ def cohere_respond(
     max_tokens,
     temperature,
     top_p,
+    hf_token,   # HF 토큰 추가
 ):
     model_name = "Cohere Command R+"
     try:
-        client = get_client(model_name)
+        client = get_client(model_name, hf_token)
     except ValueError as e:
         chat_history.append((message, str(e)))
         return chat_history
@@ -109,8 +118,6 @@ def cohere_respond(
 
     messages.append({"role": "user", "content": message})
 
-    response = ""
-    
     try:
         # Cohere Command R+ 모델을 위한 비스트리밍 처리
         response_full = client.chat_completion(
@@ -127,6 +134,7 @@ def cohere_respond(
         chat_history.append((message, error_message))
         return chat_history
 
+
 def chatgpt_respond(
     message,
     chat_history,
@@ -134,12 +142,18 @@ def chatgpt_respond(
     max_tokens,
     temperature,
     top_p,
+    openai_token,  # openai 토큰 추가
 ):
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-    if not openai.api_key:
-        chat_history.append((message, "OPENAI_API_KEY 환경 변수가 필요합니다."))
+    """
+    chatgpt용 응답. openai_token을 UI에서 입력받아 사용하도록 변경.
+    """
+    if not openai_token:
+        chat_history.append((message, "OpenAI API 토큰이 필요합니다."))
         return chat_history
 
+    # openai.api_key = os.getenv("OPENAI_API_KEY")  # 기존 코드 주석
+    openai.api_key = openai_token  # UI에서 받은 토큰 사용
+
     messages = [{"role": "system", "content": system_message}]
     for human, assistant in chat_history:
         messages.append({"role": "user", "content": human})
@@ -162,13 +176,28 @@ def chatgpt_respond(
         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("언어모델별 프롬프트 테스트 챗봇입니다.")
-    
+
+    # --- 토큰 입력 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 토큰을 입력하세요..."
+        )
+
     with gr.Tab("일반 모델"):
         with gr.Row():
             with gr.Column(scale=1):
@@ -196,8 +225,19 @@ with gr.Blocks() as demo:
                     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)
+        # 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)
     
     with gr.Tab("Cohere Command R+"):
@@ -226,16 +266,18 @@ with gr.Blocks() as demo:
             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_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_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)
     
     with gr.Tab("ChatGPT"):
@@ -264,16 +306,18 @@ with gr.Blocks() as demo:
             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_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.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)
 
 if __name__ == "__main__":