Update app.py

app.py

@@ -20,7 +20,6 @@ def get_client(model_name):
     hf_token = os.getenv("HF_TOKEN")
     if not hf_token:
         raise ValueError("HuggingFace API 토큰이 필요합니다.")
-
     if model_name == "Cohere Command R+":
         model_id = COHERE_MODEL
     else:
@@ -42,12 +41,10 @@ def respond_cohere_qna(
         client = get_client(model_name)
     except ValueError as e:
         return f"오류: {str(e)}"
-
     messages = [
         {"role": "system", "content": system_message},
         {"role": "user", "content": question}
     ]
-
     try:
         response_full = client.chat_completion(
             messages,
@@ -73,14 +70,11 @@ def respond_chatgpt_qna(
     openai_token = os.getenv("OPENAI_TOKEN")
     if not openai_token:
         return "OpenAI API 토큰이 필요합니다."
-
     openai.api_key = openai_token
-
     messages = [
         {"role": "system", "content": system_message},
         {"role": "user", "content": question}
     ]
-
     try:
         response = openai.ChatCompletion.create(
             model="gpt-4o-mini",
@@ -108,15 +102,12 @@ def respond_deepseek_qna(
     deepseek_token = os.getenv("DEEPSEEK_TOKEN")
     if not deepseek_token:
         return "DeepSeek API 토큰이 필요합니다."
-
     openai.api_key = deepseek_token
     openai.api_base = "https://api.deepseek.com/v1"
-
     messages = [
         {"role": "system", "content": system_message},
         {"role": "user", "content": question}
     ]
-
     try:
         response = openai.ChatCompletion.create(
             model=model_name,  # 선택된 모델 사용
@@ -139,15 +130,13 @@ def respond_claude_qna(
     model_name: str  # 모델 이름 파라미터 추가
 ) -> str:
     """
-    Claude API를 사용한 개선된 응답 생성 함수
+    Claude API를 사용한 개선된 응답 생성 함수.
     """
     claude_api_key = os.getenv("CLAUDE_TOKEN")
     if not claude_api_key:
         return "Claude API 토큰이 필요합니다."
-
     try:
         client = anthropic.Anthropic(api_key=claude_api_key)
-        
         message = client.messages.create(
             model=model_name,
             max_tokens=max_tokens,
@@ -157,9 +146,7 @@ def respond_claude_qna(
                 {"role": "user", "content": question}
             ]
         )
-        
         return message.content[0].text
-            
     except anthropic.APIError as ae:
         return f"Claude API 오류: {str(ae)}"
     except anthropic.RateLimitError:
@@ -175,22 +162,15 @@ def respond_o1mini_qna(
 ):
     """
     o1-mini 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
-    o1-mini에서는 'system' 역할 메시지를 지원하지 않으므로,
-    system_message와 question을 하나의 'user' 메시지로 합쳐 전달합니다.
-    또한, o1-mini에서는 'max_tokens' 대신 'max_completion_tokens' 파라미터를 사용하며,
-    온도(temperature)는 고정값 1만 지원합니다.
+    o1-mini에서는 'system' 메시지를 지원하지 않으므로 system_message와 question을 하나의 'user' 메시지로 합쳐 전달합니다.
+    또한, o1-mini에서는 'max_tokens' 대신 'max_completion_tokens'를 사용하며, temperature는 고정값 1만 지원합니다.
     """
     openai_token = os.getenv("OPENAI_TOKEN")
     if not openai_token:
         return "OpenAI API 토큰이 필요합니다."
-
     openai.api_key = openai_token
-
     combined_message = f"{system_message}\n\n{question}"
-    messages = [
-        {"role": "user", "content": combined_message}
-    ]
-
+    messages = [{"role": "user", "content": combined_message}]
     try:
         response = openai.ChatCompletion.create(
             model="o1-mini",
@@ -208,29 +188,31 @@ def respond_gemini_qna(
     system_message: str,
     max_tokens: int,
     temperature: float,
-    top_p: float,
+    top_p: float,  # top_p는 Google API에서 지원하지 않으므로 사용하지 않습니다.
     model_id: str
 ):
     """
     Gemini 모델(예: gemini-2.0-flash, gemini-2.0-flash-lite-preview-02-05)을 이용해
     한 번의 질문(question)에 대한 답변을 반환하는 함수.
-    system_message와 question을 하나의 프롬프트로 합쳐서 사용합니다.
+    Google의 genai 라이브러리를 사용합니다.
     """
-    hf_token = os.getenv("HF_TOKEN")
-    if not hf_token:
-        return "HuggingFace API 토큰이 필요합니다."
-    
-    client = InferenceClient(model_id, token=hf_token)
+    from google import genai
+    from google.genai import types
+    gemini_api_key = os.getenv("GEMINI_API_KEY")
+    if not gemini_api_key:
+        return "Gemini API 토큰이 필요합니다."
+    client = genai.Client(api_key=gemini_api_key)
     prompt = f"{system_message}\n\n{question}"
-    
     try:
-        response = client.text_generation(
-            prompt,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p
+        response = client.models.generate_content(
+            model=model_id,
+            contents=[prompt],
+            config=types.GenerateContentConfig(
+                max_output_tokens=max_tokens,
+                temperature=temperature
+            )
         )
-        return response["generated_text"]
+        return response.text
     except Exception as e:
         return f"오류가 발생했습니다: {str(e)}"
 
@@ -250,7 +232,6 @@ with gr.Blocks() as demo:
             label="모델 선택",
             value="gpt-4o-mini"
         )
-        
         with gr.Column(visible=True) as chatgpt_ui:
             chatgpt_input1_o = gr.Textbox(label="입력1", lines=1)
             chatgpt_input2_o = gr.Textbox(label="입력2", lines=1)
@@ -409,28 +390,25 @@ with gr.Blocks() as demo:
             label="모델 선택",
             value="claude-3-5-sonnet-20241022"
         )
-    
         claude_input1 = gr.Textbox(label="입력1", lines=1)
         claude_input2 = gr.Textbox(label="입력2", lines=1)
         claude_input3 = gr.Textbox(label="입력3", lines=1)
         claude_input4 = gr.Textbox(label="입력4", lines=1)
         claude_input5 = gr.Textbox(label="입력5", lines=1)
-
         claude_answer_output = gr.Textbox(label="결과", interactive=False, lines=5)
-
         with gr.Accordion("고급 설정 (Claude)", open=False):
             claude_system_message = gr.Textbox(
                 label="System Message",
                 value="""반드시 한글로 답변할 것.
 너는 Anthropic에서 개발한 클로드이다.
-최대한 정확하고 친절하게 답변하라.""",
+최대한 정확하고 친절하게 답변하라.
+""",
                 lines=3
             )
             claude_max_tokens = gr.Slider(minimum=100, maximum=4000, 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_submit_button = gr.Button("전송")
-
         def merge_and_call_claude(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, model_radio):
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_claude_qna(
@@ -463,15 +441,12 @@ with gr.Blocks() as demo:
             label="모델 선택",
             value="V3 (deepseek-chat)"
         )
-
         deepseek_input1 = gr.Textbox(label="입력1", lines=1)
         deepseek_input2 = gr.Textbox(label="입력2", lines=1)
         deepseek_input3 = gr.Textbox(label="입력3", lines=1)
         deepseek_input4 = gr.Textbox(label="입력4", lines=1)
         deepseek_input5 = gr.Textbox(label="입력5", lines=1)
-
         deepseek_answer_output = gr.Textbox(label="결과", lines=5, interactive=False)
-
         with gr.Accordion("고급 설정 (DeepSeek)", open=False):
             deepseek_system_message = gr.Textbox(
                 value="""반드시 한글로 답변할 것.
@@ -485,7 +460,6 @@ with gr.Blocks() as demo:
             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_submit_button = gr.Button("전송")
-
         def merge_and_call_deepseek(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, model_radio):
             if model_radio == "V3 (deepseek-chat)":
                 model_name = "deepseek-chat"
@@ -522,9 +496,7 @@ with gr.Blocks() as demo:
         cohere_input3 = gr.Textbox(label="입력3", lines=1)
         cohere_input4 = gr.Textbox(label="입력4", lines=1)
         cohere_input5 = gr.Textbox(label="입력5", lines=1)
-
         cohere_answer_output = gr.Textbox(label="결과", lines=5, interactive=False)
-
         with gr.Accordion("고급 설정 (Cohere)", open=False):
             cohere_system_message = gr.Textbox(
                 value="""반드시 한글로 답변할 것.
@@ -538,7 +510,6 @@ with gr.Blocks() as demo:
             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_submit_button = gr.Button("전송")
-
         def merge_and_call_cohere(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_):
             question = " ".join([i1, i2, i3, i4, i5])
             return respond_cohere_qna(