Update app.py

app.py

@@ -4,100 +4,123 @@ import google.generativeai as genai
 import asyncio
 
 ###############################################################################
-# 1. Настройка окружения и инициализация моделей
+# 1. Настройка окружения и моделей
 ###############################################################################
 
-# Ваш ключ. В идеале: os.environ.get("GEMINI_API_KEY")
-GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "ВАШ_КЛЮЧ_ЗДЕСЬ")
 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",
 ]
 
+# Инициализация моделей
 MODELS = {}
 for model_name in AVAILABLE_MODELS:
     try:
         MODELS[model_name] = genai.GenerativeModel(model_name=model_name)
     except Exception as e:
-        print(f"[Предупреждение] Не удалось инициализировать модель {model_name}: {e}")
+        print(f"[!] Не удалось инициализировать {model_name}: {e}")
 
 ###############################################################################
-# 2. Функция для определения роли ассистента
+# 2. Промпты по умолчанию (system message) для каждой модели
 ###############################################################################
+MODEL_SYSTEM_PROMPTS = {
+    # Пример: для обычной модели
+    "gemini-2.0-flash-exp": (
+        "You are a helpful assistant. "
+        "You respond in Markdown. "
+        "You do NOT wrap your final answer in {output: ...}."
+    ),
+    # Допустим, это тоже обычная модель
+    "gemini-exp-1206": (
+        "You are an experimental Gemini model. "
+        "You respond in Markdown. "
+        "You do NOT wrap your final answer in {output: ...}."
+    ),
+    # Думающая модель
+    "gemini-2.0-flash-thinking-exp-1219": (
+        "You are a thinking model. "
+        "You respond with thoughts internally, but produce the **final** answer in JSON format, like `{output: (final text)}`."
+    ),
+}
 
+###############################################################################
+# 3. Определяем, какую роль использовать: 'assistant' или 'model'
+###############################################################################
 def _assistant_role(model_name: str) -> str:
     """
-    Для некоторых моделей ("gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219") 
-    требуется role="model", а не "assistant".
+    Для некоторых моделей Gemini требуется role='model', а не 'assistant'.
+    Если столкнетесь с ошибкой "Please use a valid role: user, model", 
+    укажите здесь 'model'. 
     """
-    # Примерно так: если видим "flash-exp" без "1206" или явно gemini-2.0-flash-exp,
-    # то оставляем role="assistant". Иначе -> "model".
-    # Подстройте логику под ваши модели.
-
+    # Условие — подстройте под свои модели:
+    # Предположим, gemini-exp-1206 и gemini-2.0-flash-thinking-exp-1219
+    # требуют role='model', а остальные — 'assistant'
     if model_name in ["gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219"]:
         return "model"
-    # Все остальные — пусть будут "assistant"
     return "assistant"
 
 ###############################################################################
-# 3. Утилиты для преобразования истории между Gradio-форматом и Gemini
+# 4. Утилиты для конвертации истории
 ###############################################################################
-
 def _history_gradio_to_genai(history, model_name):
     """
-    Gradio хранит историю как [(user_msg, bot_msg), (user_msg, bot_msg), ...].
-    Модели требуют [{'role': 'user'|'model'|'assistant', 'parts': ...}, ...].
-    Для некоторых моделей вместо 'assistant' нужно 'model'.
+    Gradio: [(user_msg, bot_msg), (user_msg, bot_msg), ...]
+    → Gemini: [{'role': 'user'|'assistant'|'model'|'system', 'parts': ...}, ...]
+
+    Добавляем system-сообщение (промпт по умолчанию) в начало. 
+    Далее все user-реплики = {'role': 'user'}, 
+         все ответы ассистента = {'role': <assistant_role>}.
     """
     genai_history = []
+
+    # Системное сообщение (промпт по умолчанию)
+    system_prompt = MODEL_SYSTEM_PROMPTS.get(model_name, "")
+    if system_prompt:
+        genai_history.append({"role": "system", "parts": system_prompt})
+
     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": asst_role, "parts": bot_text})
 
     return genai_history
 
-
 ###############################################################################
-# 4. Асинхронные генераторы для получения ответа от моделей
+# 5. Функции-генераторы для ответа
 ###############################################################################
 
 async def _respond_stream(model_name, user_message, history):
     """
-    Генерация ответа (постепенного, stream=True) для обычных моделей.
-    Возвращаем кусочки текста через yield.
+    Обычная модель, stream=True: выдаём ответ порциями.
     """
     if model_name not in MODELS:
-        yield "Ошибка: модель не найдена."
+        yield "Error: model not found."
         return
 
     model = MODELS[model_name]
-
-    # Конвертируем историю в нужный формат
     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:
             partial_text += (chunk.text or "")
             yield partial_text
-
         return
     except Exception as e:
         yield f"Ошибка при запросе к API: {e}"
@@ -106,18 +129,20 @@ 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) Затем итоговый ответ + размышления.
+    2) По завершении:
+       - собираем "мысли" (part.thought)
+       - итоговый ответ оборачиваем в {output: ...} (согласно ТЗ)
     """
     if model_name not in MODELS:
-        yield "Ошибка: модель не найдена.", ""
+        yield "Error: model not found.", ""
         return
 
     model = MODELS[model_name]
     genai_history = _history_gradio_to_genai(history, model_name)
 
-    # Сначала "Думаю..."
+    # Шаг 1: "Думаю..."
     yield "Думаю...", ""
 
     try:
@@ -126,7 +151,6 @@ async def _respond_thinking(model_name, user_message, history):
 
         thinking_process_text = ""
         final_text = ""
-
         if response.candidates:
             parts = response.candidates[0].content.parts
             for p in parts:
@@ -135,44 +159,37 @@ async def _respond_thinking(model_name, user_message, history):
                 else:
                     final_text += p.text or ""
 
-        yield final_text, thinking_process_text
+        # По ТЗ: оборачиваем финальный ответ в {output: ...}
+        final_text_json = f"{{output: ({final_text})}}"
+
+        yield final_text_json, thinking_process_text
         return
     except Exception as e:
         yield f"Ошибка при запросе к API: {e}", ""
         return
 
-
 ###############################################################################
-# 5. Основная функция для Gradio
+# 6. Основная функция для одного шага диалога
 ###############################################################################
-
-async def user_send_message(
-    user_message: str,
-    history: list[tuple[str, str]],
-    model_name: str,
-    thinking_text: str
-):
+async def user_send_message(user_message, history, model_name, thinking_text):
     """
-    Параметры:
-      user_message : вход от пользователя
-      history      : [(user, assistant), ...]
-      model_name   : текущая выбранная модель
-      thinking_text: текущее «размышление»
-
-    Возвращаем (обновлённая_история, новое_thinking_text) пошагово через yield.
+    Пользователь ввёл user_message; на основе history + model_name генерируем ответ.
+    Возвращаем (history, thinking_text) через yield на ��аждом шаге.
     """
-    # Если пользователь ничего не ввёл
+    # Если пустая строка — ничего не делаем
     if not user_message.strip():
         yield history, thinking_text
         return
 
-    # Добавляем новую пару в историю
+    # Добавляем ход в history, пока ответ ассистента None
     history.append((user_message, None))
 
-    # Если в названии модели есть "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:
@@ -184,68 +201,106 @@ async def user_send_message(
             yield history, ""
         return
 
-
 ###############################################################################
-# 6. Колбэки для очистки
+# 7. Логика переключения модели, чтобы «сообщать» о переключении
 ###############################################################################
+def switch_model(old_model, new_model, history):
+    """
+    При смене модели добавляем в историю user-сообщение:
+      "I’m switching from {old_model} to {new_model}."
+    Сохраняем новое имя модели в state (prev_model_state).
+    """
+    if old_model and new_model and old_model != new_model:
+        switch_message = f"I’m switching from {old_model} to {new_model}."
+        history.append((switch_message, None))  # user_msg, None
+    return new_model, history
 
+###############################################################################
+# 8. Функция очистки
+###############################################################################
 def clear_all():
-    """ Очистить историю и размышления. """
+    """Сбросить историю и размышления."""
     return [], ""
 
-
 ###############################################################################
-# 7. Gradio-интерфейс
+# 9. Сборка Gradio-интерфейса
 ###############################################################################
-
 with gr.Blocks() as demo:
-    gr.Markdown("## Chat с Gemini (несколько моделей, «thinking»-режим)")
+    gr.Markdown("## Chat with Gemini (Thinking & Non-Thinking Models)")
+
+    # Храним «предыдущую модель», чтобы отследить переключение
+    prev_model_state = gr.State("")  
+    # Храним историю [(user, assistant), ...]
+    history_state = gr.State([])   
+    # Храним размышления (только для думающей модели)
+    thinking_store = gr.State("")
 
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=AVAILABLE_MODELS,
-            value="gemini-2.0-flash-exp",  # Можно выбрать по умолчанию
-            label="Выберите модель"
+            value="gemini-2.0-flash-exp",
+            label="Choose a model"
         )
         clear_button = gr.Button("Очистить чат")
 
-    # Состояние для хранения [(user, assistant), ...]
-    history_state = gr.State([])
-    # Состояние для размышлений
-    thinking_store = gr.State("")
-
-    chatbot = gr.Chatbot(label="Диалог с Gemini")
-    user_input = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст...")
-
-    # Поле для размышлений (только для thinking-модели)
+    chatbot = gr.Chatbot(
+        label="Диалог с Gemini",
+        markdown=True  # Включаем поддержку Markdown
+    )
+    user_input = gr.Textbox(
+        label="Ваш вопрос",
+        placeholder="Введите текст...",
+    )
     thinking_output = gr.Textbox(
-        label="Размышления",
+        label="Размышления (если модель думающая)",
         interactive=False
     )
-
     send_btn = gr.Button("Отправить")
 
-    # При нажатии «Отправить»
+    # --- Обработка переключения модели ---
+    model_dropdown.change(
+        fn=switch_model,
+        inputs=[prev_model_state, model_dropdown, history_state],
+        outputs=[prev_model_state, history_state],
+        queue=False
+    ).then(
+        # После смены модели обновим чат (чтобы увидеть user-сообщение о переключении)
+        fn=lambda h: h,
+        inputs=[history_state],
+        outputs=[chatbot],
+        queue=False
+    )
+
+    # --- При нажатии кнопки «Отправить» ---
     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],
         outputs=[thinking_output],
         queue=True
     )
+    # Очищаем поле ввода
+    send_chain.then(
+        fn=lambda: "",
+        inputs=[],
+        outputs=[user_input],
+        queue=True
+    )
 
-    # При нажатии Enter
+    # --- При нажатии Enter в поле ввода ---
     submit_chain = user_input.submit(
         fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
@@ -264,8 +319,15 @@ with gr.Blocks() as demo:
         outputs=[thinking_output],
         queue=True
     )
+    # Очищаем поле
+    submit_chain.then(
+        fn=lambda: "",
+        inputs=[],
+        outputs=[user_input],
+        queue=True
+    )
 
-    # Очистка
+    # --- Кнопка «Очистить» ---
     clear_chain = clear_button.click(
         fn=clear_all,
         inputs=[],
@@ -282,6 +344,11 @@ with gr.Blocks() as demo:
         inputs=[],
         outputs=[thinking_output]
     )
+    clear_chain.then(
+        fn=lambda: "",
+        inputs=[],
+        outputs=[user_input]
+    )
 
 if __name__ == "__main__":
     demo.launch()