Update app.py

app.py

@@ -2,113 +2,118 @@ import os
 import gradio as gr
 import google.generativeai as genai
 import asyncio
+import json
 
 ###############################################################################
-# 1. Настройка окружения и моделей
+# 1. Настройка окружения и инициализация моделей
 ###############################################################################
 
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "ВАШ_КЛЮЧ_ЗДЕСЬ")
+# Ваш ключ. В идеале: export 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-exp",               # Обычная (не-thinking)
+    "gemini-exp-1206",                    # Обычная (не-thinking)
+    "gemini-2.0-flash-thinking-exp-1219", # Thinking
 ]
 
-# Инициализация моделей
+# Инициализация
 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. Промпты по умолчанию (system message) для каждой модели
+# 2. Словарь с «дефолтными» системными промптами для каждой модели
 ###############################################################################
-MODEL_SYSTEM_PROMPTS = {
-    # Пример: для обычной модели
+
+DEFAULT_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: ...}."
+        "You are a normal model. "
+        "Just respond directly to the user without any JSON or extra formatting."
     ),
-    # Допустим, это тоже обычная модель
     "gemini-exp-1206": (
-        "You are an experimental Gemini model. "
-        "You respond in Markdown. "
-        "You do NOT wrap your final answer in {output: ...}."
+        "You are also a normal model. "
+        "Please answer in plain text, with no JSON formatting."
     ),
-    # Думающая модель
     "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)}`."
+        "Generate your final answer in the JSON format: {\"output\": \"...\"}. "
+        "You may include your chain-of-thought internally, but the user should only see the final JSON."
     ),
 }
 
+
 ###############################################################################
-# 3. Определяем, какую роль использовать: 'assistant' или 'model'
+# 3. Определение ролей для ассистента в зависимости от модели
 ###############################################################################
+
 def _assistant_role(model_name: str) -> str:
     """
-    Для некоторых моделей Gemini требуется role='model', а не 'assistant'.
-    Если столкнетесь с ошибкой "Please use a valid role: user, model", 
-    укажите здесь 'model'. 
+    Пример: пусть для gemini-2.0-flash-thinking-exp-1219 требуется role="model",
+    а для остальных role="assistant". 
+    Настраивайте эту логику под нужды своих моделей.
     """
-    # Условие — подстройте под свои модели:
-    # Предположим, 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"]:
+    if model_name == "gemini-2.0-flash-thinking-exp-1219":
         return "model"
     return "assistant"
 
+
 ###############################################################################
-# 4. Утилиты для конвертации истории
+# 4. Утилиты для преобразования истории между Gradio-форматом и Gemini
 ###############################################################################
+
 def _history_gradio_to_genai(history, model_name):
     """
-    Gradio: [(user_msg, bot_msg), (user_msg, bot_msg), ...]
-    → Gemini: [{'role': 'user'|'assistant'|'model'|'system', 'parts': ...}, ...]
-
-    Добавляем system-сообщение (промпт по умолчанию) в начало. 
-    Далее все user-реплики = {'role': 'user'}, 
-         все ответы ассистента = {'role': <assistant_role>}.
+    Gradio хранит историю как [(user_msg, bot_msg), (user_msg, bot_msg), ...].
+    Для моделей Gemini нужен формат [{'role': 'user'|'assistant'|'model', 'parts': ...}, ...].
+    Плюс добавляем в начало «дефолтный системный промпт» (если он ещё не добавлен).
     """
     genai_history = []
-
-    # Системное сообщение (промпт по умолчанию)
-    system_prompt = MODEL_SYSTEM_PROMPTS.get(model_name, "")
-    if system_prompt:
+    role_for_assistant = _assistant_role(model_name)
+
+    # 1) Добавляем «системное» сообщение, если его ещё нет
+    #    (Положим, он будет первым в истории, со специальной role="system").
+    #    Некоторые модели могут не поддерживать role="system". 
+    #    Если так, придётся хранить это как user.
+    system_prompt = DEFAULT_SYSTEM_PROMPTS.get(model_name, "")
+    # Проверим, добавляли ли мы уже системный промпт (для простоты — ищем по role="system")
+    has_system = any(h.get("role") == "system" for h in genai_history)
+
+    # Если нет в genai_history, добавим
+    if system_prompt and not has_system:
+        # Вставляем первым сообщением
         genai_history.append({"role": "system", "parts": system_prompt})
 
-    asst_role = _assistant_role(model_name)
-
+    # 2) Далее преобразуем основную историю
     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})
+            genai_history.append({"role": role_for_assistant, "parts": bot_text})
 
     return genai_history
 
+
 ###############################################################################
-# 5. Функции-генераторы для ответа
+# 5. Стримовый ответ и «thinking» ответ
 ###############################################################################
 
 async def _respond_stream(model_name, user_message, history):
     """
-    Обычная модель, stream=True: выдаём ответ порциями.
+    Обычные модели: стримим ответ.
     """
     if model_name not in MODELS:
-        yield "Error: model not found."
+        yield "Ошибка: модель не найдена."
         return
 
     model = MODELS[model_name]
@@ -117,10 +122,12 @@ async def _respond_stream(model_name, user_message, history):
     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}"
@@ -129,14 +136,13 @@ async def _respond_stream(model_name, user_message, history):
 
 async def _respond_thinking(model_name, user_message, history):
     """
-    Думающая модель:
-    1) "Думаю..."
-    2) По завершении:
-       - собираем "мысли" (part.thought)
-       - итоговый ответ оборачиваем в {output: ...} (согласно ТЗ)
+    Thinking-модель:
+    - Сначала "Думаю..."
+    - Потом итог (обёрнутый в JSON, согласно системному промпту).
+    - Выделяем 'размышления' (if p.thought == True) для thinking_output.
     """
     if model_name not in MODELS:
-        yield "Error: model not found.", ""
+        yield "Ошибка: модель не найдена.", ""
         return
 
     model = MODELS[model_name]
@@ -151,49 +157,59 @@ 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:
+                # Если это "размышления"
                 if getattr(p, "thought", False):
                     thinking_process_text += p.text or ""
                 else:
                     final_text += p.text or ""
 
-        # По ТЗ: оборачиваем финальный ответ в {output: ...}
-        final_text_json = f"{{output: ({final_text})}}"
+        # Если системный промпт прописывает формат {"output": "..."}, 
+        # то (скорее всего) модель сама уже это подставит.
+        # Но вдруг нужно «подстраховаться» вручную. 
+        # В зависимости от задачи:
+        # final_text = json.dumps({"output": final_text}, ensure_ascii=False)
 
-        yield final_text_json, thinking_process_text
+        yield final_text, thinking_process_text
         return
+
     except Exception as e:
         yield f"Ошибка при запросе к API: {e}", ""
         return
 
+
 ###############################################################################
-# 6. Основная функция для одного шага диалога
+# 6. Основная функция: user_send_message
 ###############################################################################
+
 async def user_send_message(user_message, history, model_name, thinking_text):
     """
-    Пользователь ввёл user_message; на основе history + model_name генерируем ответ.
-    Возвращаем (history, thinking_text) через yield на каждом шаге.
+    Gradio будет вызывать эту функцию при нажатии «Отправить» или Enter.
+    Параметры:
+      user_message : новая реплика пользователя
+      history      : [(user, assistant), ...]
+      model_name   : выбранная модель
+      thinking_text: текущее «размышление» (в отдельном Textbox)
     """
-    # Если пустая строка — ничего не делаем
+    # Если пользователь ничего не ввёл
     if not user_message.strip():
+        # Просто ничего не делаем
         yield history, thinking_text
         return
 
-    # Добавляем ход в history, пока ответ ассистента None
+    # Добавляем (user_message, 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:
-        # Обычная модель
         partial_answer = ""
         async for chunk in _respond_stream(model_name, user_message, history):
             partial_answer = chunk
@@ -201,77 +217,62 @@ async def user_send_message(user_message, history, model_name, thinking_text):
             yield history, ""
         return
 
+
 ###############################################################################
-# 7. Логика переключения модели, чтобы «сообщать» о переключении
+# 7. Доп. колбэк для смены модели
 ###############################################################################
-def switch_model(old_model, new_model, history):
+
+def switch_model(new_model_name, history):
     """
-    При смене модели добавляем в историю user-сообщение:
-      "I’m switching from {old_model} to {new_model}."
-    Сохраняем новое имя модели в state (prev_model_state).
+    Когда пользователь меняет выпадающий список (Dropdown) с моделью,
+    мы добавляем в историю специальное «системное» сообщение,
+    чтобы модель знала о переключении.
     """
-    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
+    history.append(
+        ("System notice", f"User switched to a different model: {new_model_name}.")
+    )
+    return history
+
 
 ###############################################################################
-# 8. Функция очистки
+# 8. Очистка
 ###############################################################################
+
 def clear_all():
-    """Сбросить историю и размышления."""
     return [], ""
 
+
 ###############################################################################
-# 9. Сборка Gradio-интерфейса
+# 9. Gradio-интерфейс
 ###############################################################################
-with gr.Blocks() as demo:
-    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.Blocks() as demo:
+    gr.Markdown("## Chat с Gemini (несколько моделей, «thinking»-режим, JSON-ответы)")
 
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=AVAILABLE_MODELS,
             value="gemini-2.0-flash-exp",
-            label="Choose a model"
+            label="Выберите модель"
         )
         clear_button = gr.Button("Очистить чат")
 
-    chatbot = gr.Chatbot(
-        label="Диалог с Gemini",
-        markdown=True  # Включаем поддержку Markdown
-    )
-    user_input = gr.Textbox(
-        label="Ваш вопрос",
-        placeholder="Введите текст...",
-    )
+    # Состояние для хранения истории
+    history_state = gr.State([])
+    # Состояние для «размышлений»
+    thinking_store = gr.State("")
+
+    chatbot = gr.Chatbot(label="Диалог с Gemini")
+    user_input = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст...")
+
     thinking_output = gr.Textbox(
-        label="Размышления (если модель думающая)",
+        label="Размышления (для thinking-моделей)",
         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_btn = gr.Button("Отправить")
 
-    # --- При нажатии кнопки «Отправить» ---
+    # 1) Логика при нажатии «Отправить»
     send_chain = send_btn.click(
         fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
@@ -285,22 +286,22 @@ with gr.Blocks() as demo:
         outputs=[chatbot],
         queue=True
     )
-    # Обновляем размышления
+    # Обновляем поле «Размышления»
     send_chain.then(
         fn=lambda t: t,
         inputs=[thinking_store],
         outputs=[thinking_output],
         queue=True
     )
-    # Очищаем поле ввода
+    # Очищаем поле ввода (user_input)
     send_chain.then(
         fn=lambda: "",
         inputs=[],
         outputs=[user_input],
-        queue=True
+        queue=False
     )
 
-    # --- При нажатии Enter в поле ввода ---
+    # 2) При нажатии Enter в поле ввода
     submit_chain = user_input.submit(
         fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
@@ -319,15 +320,26 @@ with gr.Blocks() as demo:
         outputs=[thinking_output],
         queue=True
     )
-    # Очищаем поле
     submit_chain.then(
         fn=lambda: "",
         inputs=[],
         outputs=[user_input],
-        queue=True
+        queue=False
+    )
+
+    # 3) Колбэк при смене модели: добавляем системное сообщение в историю
+    model_dropdown.change(
+        fn=switch_model,
+        inputs=[model_dropdown, history_state],
+        outputs=[history_state],
+        queue=False
+    ).then(
+        fn=lambda h: h,  # обновим Chatbot
+        inputs=[history_state],
+        outputs=[chatbot]
     )
 
-    # --- Кнопка «Очистить» ---
+    # 4) Кнопка «Очистить чат»
     clear_chain = clear_button.click(
         fn=clear_all,
         inputs=[],
@@ -343,8 +355,7 @@ with gr.Blocks() as demo:
         fn=lambda _: "",
         inputs=[],
         outputs=[thinking_output]
-    )
-    clear_chain.then(
+    ).then(
         fn=lambda: "",
         inputs=[],
         outputs=[user_input]