Update app.py

app.py

@@ -19,56 +19,11 @@ import io
 from transformers import AutoTokenizer, AutoModel
 from utils import best_text_choice
 
-import httpx
 import asyncio
-import ssl
 
 # Авторизация в GigaChat Pro
 gc_key = os.getenv('GC_KEY')
 
-class AsyncGigaChat:
-    def __init__(self, credentials, model, base_url, max_tokens, temperature, verify_ssl_certs=True):
-        self.credentials = credentials
-        self.model = model
-        self.base_url = base_url
-        self.max_tokens = max_tokens
-        self.temperature = temperature
-        self.verify_ssl_certs = verify_ssl_certs
-        self.client = httpx.AsyncClient(verify=self.verify_ssl_certs)
-
-    async def __call__(self, messages):
-        headers = {
-            "Authorization": f"Bearer {self.credentials}",
-            "Content-Type": "application/json"
-        }
-        payload = {
-            "model": self.model,
-            "messages": [msg.content for msg in messages],
-            "max_tokens": self.max_tokens,
-            "temperature": self.temperature
-        }
-        response = await self.client.post(f"{self.base_url}/generate", headers=headers, json=payload)
-        response.raise_for_status()
-        return response.json()
-        
-chat_pro_async = AsyncGigaChat(
-    credentials=gc_key,
-    model='GigaChat-Pro-preview',
-    base_url='https://gigachat-preview.devices.sberbank.ru/api/v1/',
-    max_tokens=68,
-    temperature=1.15,
-    verify_ssl_certs=False
-)
-
-chat_pro_check_async = AsyncGigaChat(
-    credentials=gc_key,
-    model='GigaChat-Pro-preview',
-    base_url='https://gigachat-preview.devices.sberbank.ru/api/v1/',
-    max_tokens=3000,
-    temperature=0.8,
-    verify_ssl_certs=False
-)
-
 
 tokenizer = AutoTokenizer.from_pretrained("Sergeyzh/rubert-tiny-turbo")
 model = AutoModel.from_pretrained("Sergeyzh/rubert-tiny-turbo")
@@ -86,25 +41,23 @@ client_mistral_check = Mistral(api_key=MISTRAL_API_KEY)
 
 morph = pymorphy3.MorphAnalyzer()
 
+chat_pro = GigaChat(
+   credentials=gc_key,
+   model='GigaChat-Pro-preview',
+   base_url='https://gigachat-preview.devices.sberbank.ru/api/v1/',
+   max_tokens=68,
+   temperature=1.15,
+   verify_ssl_certs=False
+)
 
-#chat_pro = GigaChat(credentials=gc_key, model='GigaChat', max_tokens=68, temperature=1.15, verify_ssl_certs=False)
-#chat_pro = GigaChat(
-#   credentials=gc_key,
-#   model='GigaChat-Pro-preview',
-#   base_url='https://gigachat-preview.devices.sberbank.ru/api/v1/',
-#   max_tokens=68,
-#   temperature=1.15,
-#   verify_ssl_certs=False
-#)
-
-#chat_pro_check = GigaChat(
-#    credentials=gc_key,
-#    model='GigaChat-Pro-preview',
-#    base_url='https://gigachat-preview.devices.sberbank.ru/api/v1/',
-#    max_tokens=3000,
-#    temperature=0.8,
-#    verify_ssl_certs=False
-#)
+chat_pro_check = GigaChat(
+    credentials=gc_key,
+    model='GigaChat-Pro-preview',
+    base_url='https://gigachat-preview.devices.sberbank.ru/api/v1/',
+    max_tokens=3000,
+    temperature=0.8,
+    verify_ssl_certs=False
+)
 
 approach_stats = {
     "Начни сообщение с призыва к действию с продуктом. Не начинай с вопроса.": {"failed_checks": defaultdict(int), "total_attempts": 0},
@@ -242,7 +195,7 @@ def clean_message(message):
 async def generate_message_mistral_generate(prompt):
     try:
         messages = [SystemMessage(content=prompt)]
-        res = await chat_pro_async(messages)  # Предполагается, что chat_pro_async использует httpx
+        res = await chat_pro.ainvoke(messages)
         cleaned_message = clean_message(res.content.strip())
         return cleaned_message
     except Exception as e:
@@ -280,11 +233,11 @@ async def generate_message_mistral_generate(prompt):
 #    return None
 
 #def generate_message_mistral_check(prompt, max_retries=5):
-async def generate_message_mistral_check(prompt):
+async def generate_message_mistral_check(prompt):    
     try:
         messages = [SystemMessage(content=prompt)]
-        res = await chat_pro_check_async(messages)  # Предполагается, что chat_pro_async использует httpx
-        cleaned_message = clean_message(res.content.strip())
+        res2 = await chat_pro_check.ainvoke(messages)
+        cleaned_message = clean_message(res2.content.strip())
         return cleaned_message
     except Exception as e:
         return f"Ошибка при обращении к GigaChat-Pro: {e}"