Update app.py

app.py

@@ -4,25 +4,24 @@ import google.generativeai as genai
 import asyncio
 
 ###############################################################################
-# 1. Настройка окружения и инициализация модели
+# 1. Настройка окружения и инициализация моделей
 ###############################################################################
 
-# Задайте свой ключ через переменную окружения, например:
-# export GEMINI_API_KEY="ваш-ключ"
-GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc" 
+# Ваш ключ. В идеале: os.environ.get("GEMINI_API_KEY")
+GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
 if not GEMINI_API_KEY:
     print("Error: GEMINI_API_KEY is not set.")
     exit()
+
 genai.configure(api_key=GEMINI_API_KEY)
 
-# Список доступных моделей
+# Список доступных моделей (добавьте/уберите при необходимости)
 AVAILABLE_MODELS = [
     "gemini-2.0-flash-exp",
     "gemini-exp-1206",
-    "gemini-2.0-flash-thinking-exp-1219"
+    "gemini-2.0-flash-thinking-exp-1219",
 ]
 
-# Инициализируем модели и храним их в словаре
 MODELS = {}
 for model_name in AVAILABLE_MODELS:
     try:
@@ -31,30 +30,54 @@ for model_name in AVAILABLE_MODELS:
         print(f"[Предупреждение] Не удалось инициализировать модель {model_name}: {e}")
 
 ###############################################################################
-# 2. Утилиты для преобразования истории между Gradio-форматом и Gemini
+# 2. Функция для определения роли ассистента
 ###############################################################################
 
-def _history_gradio_to_genai(history):
+def _assistant_role(model_name: str) -> str:
+    """
+    Для некоторых моделей ("gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219") 
+    требуется role="model", а не "assistant".
+    """
+    # Примерно так: если видим "flash-exp" без "1206" или явно gemini-2.0-flash-exp,
+    # то оставляем role="assistant". Иначе -> "model".
+    # Подстройте логику под ваши модели.
+
+    if model_name in ["gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219"]:
+        return "model"
+    # Все остальные — пусть будут "assistant"
+    return "assistant"
+
+###############################################################################
+# 3. Утилиты для преобразования истории между Gradio-форматом и Gemini
+###############################################################################
+
+def _history_gradio_to_genai(history, model_name):
     """
     Gradio хранит историю как [(user_msg, bot_msg), (user_msg, bot_msg), ...].
-    Библиотека google.generativeai ждёт [{'role': 'user', 'parts': ...}, {'role': 'assistant', 'parts': ...}, ...].
+    Модели требуют [{'role': 'user'|'model'|'assistant', 'parts': ...}, ...].
+    Для некоторых моделей вместо 'assistant' нужно 'model'.
     """
     genai_history = []
+    asst_role = _assistant_role(model_name)
+
     for user_text, bot_text in history:
+        # Сообщение от пользователя
         if user_text:
             genai_history.append({"role": "user", "parts": user_text})
+        # Сообщение от ассистента
         if bot_text:
-            genai_history.append({"role": "assistant", "parts": bot_text})
+            genai_history.append({"role": asst_role, "parts": bot_text})
+
     return genai_history
 
 
 ###############################################################################
-# 3. Асинхронные генераторы для получения ответа от моделей
+# 4. Асинхронные генераторы для получения ответа от моделей
 ###############################################################################
 
 async def _respond_stream(model_name, user_message, history):
     """
-    Генерация ответа (постепенного, stream=True) для обычных моделей (без thinking).
+    Генерация ответа (постепенного, stream=True) для обычных моделей.
     Возвращаем кусочки текста через yield.
     """
     if model_name not in MODELS:
@@ -62,21 +85,20 @@ async def _respond_stream(model_name, user_message, history):
         return
 
     model = MODELS[model_name]
-    genai_history = _history_gradio_to_genai(history)
+
+    # Конвертируем историю в нужный формат
+    genai_history = _history_gradio_to_genai(history, model_name)
 
     try:
-        # Создаём чат с уже накопленной историей
         chat = model.start_chat(history=genai_history)
-        # Запрашиваем стриминг (постепенный) ответа
         response_stream = chat.send_message(user_message, stream=True)
 
         partial_text = ""
         for chunk in response_stream:
-            chunk_text = chunk.text or ""
-            partial_text += chunk_text
+            partial_text += (chunk.text or "")
             yield partial_text
 
-        return  # Завершить генератор без значения
+        return
     except Exception as e:
         yield f"Ошибка при запросе к API: {e}"
         return
@@ -84,23 +106,22 @@ async def _respond_stream(model_name, user_message, history):
 
 async def _respond_thinking(model_name, user_message, history):
     """
-    Генерация ответа для thinking-модели (gemini-2.0-flash-thinking-exp).
-    1) Сначала отдаём "Думаю..."
-    2) Затем формируем итоговый ответ + собираем «размышления» part.thought == True.
+    Генерация ответа для «thinking»-модели (gemini-2.0-flash-thinking-exp и т.п.).
+    1) "Думаю..."
+    2) Затем итоговый ответ + размышления.
     """
     if model_name not in MODELS:
         yield "Ошибка: модель не найдена.", ""
         return
 
     model = MODELS[model_name]
-    genai_history = _history_gradio_to_genai(history)
+    genai_history = _history_gradio_to_genai(history, model_name)
 
-    # Шаг 1: сообщаем "Думаю..."
+    # Сначала "Думаю..."
     yield "Думаю...", ""
 
     try:
         chat = model.start_chat(history=genai_history)
-        # Без стриминга — дожидаемся полного ответа
         response = chat.send_message(user_message, stream=False)
 
         thinking_process_text = ""
@@ -122,7 +143,7 @@ async def _respond_thinking(model_name, user_message, history):
 
 
 ###############################################################################
-# 4. Основная функция для обработки нового пользовательского сообщения в Gradio
+# 5. Основная функция для Gradio
 ###############################################################################
 
 async def user_send_message(
@@ -133,31 +154,29 @@ async def user_send_message(
 ):
     """
     Параметры:
-      user_message : строка, введённая пользователем
-      history      : текущая история [(user, assistant), ...]
-      model_name   : выбранная модель
-      thinking_text: текущее содержимое поля «Размышления»
-    
-    Возвращаем (обновлённую историю, обновлённое thinking_text) пошагово через yield.
+      user_message : вход от пользователя
+      history      : [(user, assistant), ...]
+      model_name   : текущая выбранная модель
+      thinking_text: текущее «размышление»
+
+    Возвращаем (обновлённая_история, новое_thinking_text) пошагово через yield.
     """
+    # Если пользователь ничего не ввёл
     if not user_message.strip():
-        # Если пользователь не ввёл ничего, не делаем запрос
         yield history, thinking_text
         return
 
-    # Добавляем новый шаг в историю (ассистент пока None)
+    # Добавляем новую пару в историю
     history.append((user_message, None))
 
-    # Если модель — thinking
+    # Если в названии модели есть "thinking"
     if "thinking" in model_name.lower():
         async for (assistant_text, thought_text) in _respond_thinking(model_name, user_message, history):
-            # Обновляем последний элемент истории
             history[-1] = (user_message, assistant_text)
-            # Записываем «размышления» во второе поле
             yield history, thought_text
         return
     else:
-        # Обычная модель (stream)
+        # Обычная модель
         partial_answer = ""
         async for chunk in _respond_stream(model_name, user_message, history):
             partial_answer = chunk
@@ -167,7 +186,7 @@ async def user_send_message(
 
 
 ###############################################################################
-# 5. Колбэки для очистки
+# 6. Колбэки для очистки
 ###############################################################################
 
 def clear_all():
@@ -176,50 +195,49 @@ def clear_all():
 
 
 ###############################################################################
-# 6. Gradio-интерфейс
+# 7. Gradio-интерфейс
 ###############################################################################
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Chat с Gemini (поддержка разных моделей и «thinking»-режима)")
+    gr.Markdown("## Chat с Gemini (несколько моделей, «thinking»-режим)")
 
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=AVAILABLE_MODELS,
-            value="gemini-1.5-flash",
+            value="gemini-2.0-flash-exp",  # Можно выбрать по умолчанию
             label="Выберите модель"
         )
         clear_button = gr.Button("Очистить чат")
 
-    # Состояние для хранения истории и размышлений
-    history_state = gr.State([])   # [(user, assistant), ...]
-    thinking_store = gr.State("")  # текст размышлений
+    # Состояние для хранения [(user, assistant), ...]
+    history_state = gr.State([])
+    # Состояние для размышлений
+    thinking_store = gr.State("")
 
     chatbot = gr.Chatbot(label="Диалог с Gemini")
     user_input = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст...")
 
-    # Поле для размышлений
+    # Поле для размышлений (только для thinking-модели)
     thinking_output = gr.Textbox(
-        label="Размышления (только для gemini-2.0-flash-thinking-exp)",
+        label="Размышления",
         interactive=False
     )
 
     send_btn = gr.Button("Отправить")
 
-    # 1) При нажатии «Отправить»
+    # При нажатии «Отправить»
     send_chain = send_btn.click(
         fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
         outputs=[history_state, thinking_store],
         queue=True
     )
-    # Затем обновляем чат
     send_chain.then(
         fn=lambda h: h,
         inputs=[history_state],
         outputs=[chatbot],
         queue=True
     )
-    # И отображаем размышления
     send_chain.then(
         fn=lambda t: t,
         inputs=[thinking_store],
@@ -227,7 +245,7 @@ with gr.Blocks() as demo:
         queue=True
     )
 
-    # 2) При нажатии Enter в textbox
+    # При нажатии Enter
     submit_chain = user_input.submit(
         fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
@@ -247,7 +265,7 @@ with gr.Blocks() as demo:
         queue=True
     )
 
-    # 3) Кнопка «Очистить»
+    # Очистка
     clear_chain = clear_button.click(
         fn=clear_all,
         inputs=[],
@@ -265,6 +283,5 @@ with gr.Blocks() as demo:
         outputs=[thinking_output]
     )
 
-# Запуск приложения
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()