Update app.py

app.py

@@ -9,16 +9,15 @@ def load_system_role(role_name):
         roles = json.load(file)
     return roles.get(role_name, "Ты помощник по умолчанию.")
 
-# Функция для загрузки названий ролей из JSON файла
 def load_role_names():
     with open('system_roles.json', 'r', encoding='utf-8') as file:
         roles = json.load(file)
     return list(roles.keys())
 
-# Функция для генерации текста
 def generate(description, system_role_name, max_tokens):
     if not description:
-        return None, None
+        yield None, None
+        return
 
     system_role = load_system_role(system_role_name)
 
@@ -30,28 +29,36 @@ def generate(description, system_role_name, max_tokens):
     payload = {
         'messages': [{'role': 'system', 'content': system_role}, {'role': 'user', 'content': description}],
         'max_tokens': max_tokens,
-        'model': "gemini-1.5-pro-latest"
+        'model': "gemini-1.5-pro-latest",
+        'stream': True
     }
 
     try:
-        response = requests.post(os.getenv("BASE_URL"), headers=headers, json=payload, timeout=200)
-        response.raise_for_status()  # Проверяем на ошибки HTTP
-        data = response.json()  # Используем response.json() вместо json.loads()
+        response = requests.post(os.getenv("BASE_URL"), headers=headers, json=payload, stream=True, timeout=200)
+        response.raise_for_status()
 
-        if 'choices' in data and len(data['choices']) > 0:
-            command = data['choices'][0]['message']['content'].strip()
-            print("Генерация успешна")
-            return command, command
-        elif 'error' in data:
-            error_message = data['error']['message']
-            print(f'Ошибка: {error_message}')
-            return f"**Ошибка сервера!**\n\n```\n{error_message}\n```", f"Ошибка сервера!\n\n{error_message}"
-        else:
-            print(f'Не удалось сгенерировать текст. Ответ сервера: {data}')  # Выводим ответ сервера для отладки
-            return "**Неизвестная ошибка!**", "Неизвестная ошибка!"
+        full_text = ""
+        for chunk in response.iter_lines():
+            if chunk:
+                try:
+                    chunk = chunk.decode('utf-8').replace("data: ", "")
+                    if chunk == "[DONE]":
+                        break
+                    chunk_data = json.loads(chunk)
+                    if 'choices' in chunk_data and len(chunk_data['choices']) > 0:
+                        text_chunk = chunk_data['choices'][0]['delta'].get('content', "")
+                        full_text += text_chunk
+                        yield full_text, full_text
+                except json.JSONDecodeError:
+                    continue
+        if not full_text:
+            yield "**Не удалось получить ответ от сервера.**", "Не удалось получить ответ от сервера."
     except requests.exceptions.RequestException as e:
         print(f"Ошибка запроса: {e}")
-        return f"**Ошибка запроса!**\n\n```\n{e}\n```", f"Ошибка запроса!\n\n{e}"
+        yield f"**Ошибка запроса!**\n\n```\n{e}\n```", f"Ошибка запроса!\n\n{e}"
+    except Exception as e:
+        print(f"Ошибка: {str(e)}")
+        yield "Произошла ошибка при генерации", "Произошла ошибка при генерации"
 
 # Ссылка на файл CSS
 css_url = "https://neurixyufi-aihub.static.hf.space/style.css"