Update app.py

app.py

@@ -12,6 +12,7 @@ from requests.exceptions import HTTPError
 import matplotlib.pyplot as plt
 from io import BytesIO
 import base64
+import time
 
 # 로깅 설정
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s:%(message)s', handlers=[logging.StreamHandler()])
@@ -24,7 +25,8 @@ intents.guilds = True
 intents.guild_messages = True
 
 # 추론 API 클라이언트 설정
-hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))
+hf_client_primary = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))
+hf_client_secondary = InferenceClient("CohereForAI/aya-23-35B", token=os.getenv("HF_TOKEN"))
 
 # 수학 전문 LLM 파이프라인 설정
 math_pipe = pipeline("text-generation", model="AI-MO/NuminaMath-7B-TIR")
@@ -35,20 +37,19 @@ SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))
 # 대화 히스토리를 저장할 전역 변수
 conversation_history = []
 
-
 def latex_to_image(latex_string):
     plt.figure(figsize=(10, 1))
     plt.axis('off')
-    plt.text(0.5, 0.5, f'${latex_string}$', size=20, ha='center', va='center', color='white')
+    plt.text(0.5, 0.5, latex_string, size=20, ha='center', va='center', color='white')
     
     buffer = BytesIO()
     plt.savefig(buffer, format='png', bbox_inches='tight', pad_inches=0.1, transparent=True, facecolor='black')
     buffer.seek(0)
     
+    image_base64 = base64.b64encode(buffer.getvalue()).decode()
     plt.close()
     
-    return buffer
-    
+    return image_base64
 
 def process_and_convert_latex(text):
     # 단일 $ 또는 이중 $$ 로 둘러싸인 LaTeX 수식을 찾습니다.
@@ -71,6 +72,7 @@ class MyClient(discord.Client):
         super().__init__(*args, **kwargs)
         self.is_processing = False
         self.math_pipe = math_pipe
+        self.hf_client = hf_client_primary  # 초기 클라이언트 설정
 
     async def on_ready(self):
         logging.info(f'{self.user}로 로그인되었습니다!')
@@ -87,11 +89,8 @@ class MyClient(discord.Client):
 
         self.is_processing = True
         try:
-            if isinstance(message.channel, discord.Thread):
-                thread = message.channel
-            else:
-                thread = await message.channel.create_thread(name=f"질문: {message.author.name}", message=message)
-
+            # 새로운 스레드 생성
+            thread = await message.channel.create_thread(name=f"질문: {message.author.name}", message=message)
             if self.is_math_question(message.content):
                 text_response = await self.handle_math_question(message.content)
                 await self.send_message_with_latex(thread, text_response)
@@ -100,7 +99,7 @@ class MyClient(discord.Client):
                 await self.send_message_with_latex(thread, response)
         finally:
             self.is_processing = False
-        
+
     def is_message_in_specific_channel(self, message):
         return message.channel.id == SPECIFIC_CHANNEL_ID or (
             isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
@@ -119,10 +118,9 @@ class MyClient(discord.Client):
         
         try:
             # Cohere 모델에게 AI-MO/NuminaMath-7B-TIR 모델의 결과를 번역하도록 요청
-            cohere_response_future = loop.run_in_executor(None, lambda: hf_client.chat_completion(
+            cohere_response = await self.retry_request(lambda: self.hf_client.chat_completion(
                 [{"role": "system", "content": "다음 텍스트를 한글로 번역하십시오: "}, {"role": "user", "content": math_result}], max_tokens=1000))
         
-            cohere_response = await cohere_response_future
             cohere_result = ''.join([part.choices[0].delta.content for part in cohere_response if part.choices and part.choices[0].delta and part.choices[0].delta.content])
 
             combined_response = f"수학 선생님 답변: ```{cohere_result}```"
@@ -150,7 +148,7 @@ class MyClient(discord.Client):
         messages = [{"role": "system", "content": f"{system_prefix}"}] + conversation_history
 
         try:
-            response = await asyncio.get_event_loop().run_in_executor(None, lambda: hf_client.chat_completion(
+            response = await self.retry_request(lambda: self.hf_client.chat_completion(
                 messages, max_tokens=1000, stream=True, temperature=0.7, top_p=0.85))
             full_response = ''.join([part.choices[0].delta.content for part in response if part.choices and part.choices[0].delta and part.choices[0].delta.content])
             conversation_history.append({"role": "assistant", "content": full_response})
@@ -162,24 +160,25 @@ class MyClient(discord.Client):
 
     async def send_message_with_latex(self, channel, message):
         try:
-        # 텍스트와 LaTeX 수식 분리
+            # 텍스트와 LaTeX 수식 분리
             text_parts = re.split(r'(\$\$.*?\$\$|\$.*?\$)', message, flags=re.DOTALL)
         
             for part in text_parts:
                 if part.startswith('$'):
-                # LaTeX 수식 처리 및 이미지로 출력
+                    # LaTeX 수식 처리 및 이미지로 출력
                     latex_content = part.strip('$')
-                    image_buffer = latex_to_image(latex_content)
-                    await channel.send(file=discord.File(fp=image_buffer, filename='equation.png'))
+                    image_base64 = latex_to_image(latex_content)
+                    image_binary = base64.b64decode(image_base64)
+                    await channel.send(file=discord.File(BytesIO(image_binary), 'equation.png'))
                 else:
-                # 텍스트 출력
+                    # 텍스트 출력
                     if part.strip():
                         await self.send_long_message(channel, part.strip())
     
         except Exception as e:
             logging.error(f"Error in send_message_with_latex: {str(e)}")
             await channel.send("An error occurred while processing the message.")
-     
+        
     async def send_long_message(self, channel, message):
         if len(message) <= 2000:
             await channel.send(message)
@@ -188,19 +187,22 @@ class MyClient(discord.Client):
             for part in parts:
                 await channel.send(part)
 
-
     async def retry_request(self, func, retries=5, delay=2):
         for i in range(retries):
             try:
                 return await func()
             except HTTPError as e:
-                if e.response.status_code == 503 and i < retries - 1:
-                    logging.warning(f"503 error encountered. Retrying in {delay} seconds...")
-                    await asyncio.sleep(delay)
+                if e.response.status_code == 503:
+                    if i < retries - 1:
+                        logging.warning(f"503 error encountered. Retrying in {delay} seconds...")
+                        await asyncio.sleep(delay)
+                    else:
+                        logging.warning("Switching to secondary model due to repeated 503 errors.")
+                        self.hf_client = hf_client_secondary
+                        await asyncio.sleep(delay)
                 else:
                     raise
 
-
 if __name__ == "__main__":
     discord_client = MyClient(intents=intents)
     discord_client.run(os.getenv('DISCORD_TOKEN'))