Update app.py

app.py

@@ -20,10 +20,53 @@ from transformers import AutoTokenizer, AutoModel
 from utils import best_text_choice
 
 import httpx
-import aiohttp
 import asyncio
 import ssl
 
+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()
+
+    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")
 # Load the DataFrame used in the best_text_choice function
@@ -196,30 +239,13 @@ def clean_message(message):
 #def generate_message_mistral_generate(prompt, max_retries=5):
 async def generate_message_mistral_generate(prompt):
     try:
-        # Создаём SSL-контекст для отключения проверки сертификатов
-        # В продакшн-среде НЕ РЕКОМЕНДУЕТСЯ использовать verify=False
-        ssl_context = ssl.create_default_context()
-        ssl_context.check_hostname = False
-        ssl_context.verify_mode = ssl.CERT_NONE
-
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                "https://gigachat-preview.devices.sberbank.ru/api/v1/",
-                headers={
-                    "Authorization": f"Bearer {gc_key}",
-                    "Content-Type": "application/json"
-                },
-                json={"messages": [{"role": "system", "content": prompt}]},
-                ssl=ssl_context  # Используем ssl=ssl_context для отключения проверки
-            ) as response:
-                response.raise_for_status()
-                res = await response.json()
-                cleaned_message = clean_message(res.get("content", "").strip())
-                return cleaned_message
-    except aiohttp.ClientError as e:
-        return f"HTTP ошибка при обращении к GigaChat-Pro: {e}"
+        messages = [SystemMessage(content=prompt)]
+        res = await chat_pro_async(messages)  # Предполагается, что chat_pro_async использует httpx
+        cleaned_message = clean_message(res.content.strip())
+        return cleaned_message
     except Exception as e:
-        return f"Неизвестная ошибка при обращении к GigaChat-Pro: {e}"
+        return f"Ошибка при обращении к GigaChat-Pro: {e}"
+
     
 #    retries = 0
 #    while retries < max_retries:
@@ -254,28 +280,12 @@ async def generate_message_mistral_generate(prompt):
 #def generate_message_mistral_check(prompt, max_retries=5):
 async def generate_message_mistral_check(prompt):
     try:
-        ssl_context = ssl.create_default_context()
-        ssl_context.check_hostname = False
-        ssl_context.verify_mode = ssl.CERT_NONE
-
-        async with aiohttp.ClientSession() as session:
-            async with session.post(
-                "https://gigachat-preview.devices.sberbank.ru/api/v1/",
-                headers={
-                    "Authorization": f"Bearer {gc_key}",
-                    "Content-Type": "application/json"
-                },
-                json={"messages": [{"role": "system", "content": prompt}]},
-                ssl=ssl_context
-            ) as response:
-                response.raise_for_status()
-                res2 = await response.json()
-                cleaned_message = clean_message(res2.get("content", "").strip())
-                return cleaned_message
-    except aiohttp.ClientError as e:
-        return f"HTTP ошибка при обращении к GigaChat-Pro: {e}"
+        messages = [SystemMessage(content=prompt)]
+        res = await chat_pro_check_async(messages)  # Предполагается, что chat_pro_async и��пользует httpx
+        cleaned_message = clean_message(res.content.strip())
+        return cleaned_message
     except Exception as e:
-        return f"Неизвестная ошибка при обращении к GigaChat-Pro: {e}"
+        return f"Ошибка при обращении к GigaChat-Pro: {e}"
 #    retries = 0
 #    while retries < max_retries:
 #        try: