Update app.py

app.py

@@ -2,13 +2,12 @@ import os
 import gradio as gr
 import google.generativeai as genai
 import asyncio
-import json
 
 ###############################################################################
 # 1. Настройка окружения и инициализация моделей
 ###############################################################################
 
-# Ваш ключ. В идеале: export GEMINI_API_KEY="..."
+# Подставьте свой ключ или берите из окружения
 GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
 if not GEMINI_API_KEY:
     print("Error: GEMINI_API_KEY is not set.")
@@ -16,14 +15,13 @@ if not GEMINI_API_KEY:
 
 genai.configure(api_key=GEMINI_API_KEY)
 
-# Список доступных моделей
+# Выберите доступные модели (пример)
 AVAILABLE_MODELS = [
-    "gemini-2.0-flash-exp",               # Обычная (не-thinking)
-    "gemini-exp-1206",                    # Обычная (не-thinking)
-    "gemini-2.0-flash-thinking-exp-1219", # Thinking
+    "gemini-2.0-flash-exp",
+    "gemini-exp-1206",
+    "gemini-2.0-flash-thinking-exp-1219",
 ]
 
-# Инициализация
 MODELS = {}
 for model_name in AVAILABLE_MODELS:
     try:
@@ -31,100 +29,84 @@ for model_name in AVAILABLE_MODELS:
     except Exception as e:
         print(f"[Предупреждение] Не удалось инициализировать модель {model_name}: {e}")
 
-
 ###############################################################################
-# 2. Словарь с «дефолтными» системными промптами для каждой модели
+# 2. Дефолтные промпты (developer role) для каждой модели
 ###############################################################################
+# Когда пользователь переключается на модель, мы добавляем это сообщение в историю.
 
-DEFAULT_SYSTEM_PROMPTS = {
+DEFAULT_DEVELOPER_PROMPTS = {
     "gemini-2.0-flash-exp": (
-        "You are a normal model. "
-        "Just respond directly to the user without any JSON or extra formatting."
+        "You are a normal model (developer role). "
+        "Provide direct answers with no JSON wrapping."
     ),
     "gemini-exp-1206": (
-        "You are also a normal model. "
-        "Please answer in plain text, with no JSON formatting."
+        "You are an experimental normal model (developer role). "
+        "Provide direct answers with no JSON wrapping."
     ),
     "gemini-2.0-flash-thinking-exp-1219": (
-        "You are a thinking model. "
-        "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."
+        "You are a thinking model (developer role). "
+        "Please provide your final answer in the format {output: ...}. "
+        "You may use internal thoughts but do not show them directly to the user."
     ),
 }
 
-
 ###############################################################################
-# 3. Определение ролей для ассистента в зависимости от модели
+# 3. Функция для определения роли ассистента (assistant vs model)
 ###############################################################################
 
 def _assistant_role(model_name: str) -> str:
     """
-    Пример: пусть для gemini-2.0-flash-thinking-exp-1219 требуется role="model",
-    а для остальных role="assistant". 
-    Настраивайте эту логику под нужды своих моделей.
+    Некоторые новые модели не принимают 'assistant', а требуют 'model'.
     """
-    if model_name == "gemini-2.0-flash-thinking-exp-1219":
+    # Допустим "gemini-exp-1206" и "gemini-2.0-flash-thinking-exp-1219" хотят "model"
+    if model_name in ["gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219"]:
         return "model"
     return "assistant"
 
-
 ###############################################################################
-# 4. Утилиты для преобразования истории между Gradio-форматом и Gemini
+# 4. Преобразование истории из Gradio в формат Generative AI
 ###############################################################################
 
-def _history_gradio_to_genai(history, model_name):
+def _history_to_genai(history, model_name):
     """
-    Gradio хранит историю как [(user_msg, bot_msg), (user_msg, bot_msg), ...].
-    Для моделей Gemini нужен формат [{'role': 'user'|'assistant'|'model', 'parts': ...}, ...].
-    Плюс добавляем в начало «дефолтный системный промпт» (если он ещё не добавлен).
+    Gradio хранит [(user_msg, bot_msg), ...].
+    Google GenAI ждёт [{"role": "...", "parts": "..."}].
+    Дополнительно учитываем developer-сообщения, которые мы вставляем.
     """
     genai_history = []
-    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})
-
-    # 2) Далее преобразуем основную историю
+    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": role_for_assistant, "parts": bot_text})
-
+            # Ответ ассистента (или model)
+            genai_history.append({"role": asst_role, "parts": bot_text})
     return genai_history
 
-
 ###############################################################################
-# 5. Стримовый ответ и «thinking» ответ
+# 5. Генераторы для стрима обычных моделей и "thinking" моделей
 ###############################################################################
 
 async def _respond_stream(model_name, user_message, history):
     """
-    Обычные модели: стримим ответ.
+    Стриминговый ответ для обычных моделей: 
+    - Кусочек за кусочком (partial_text).
     """
     if model_name not in MODELS:
         yield "Ошибка: модель не найдена."
         return
 
     model = MODELS[model_name]
-    genai_history = _history_gradio_to_genai(history, model_name)
+    genai_history = _history_to_genai(history, model_name)
 
     try:
         chat = model.start_chat(history=genai_history)
-        response_stream = chat.send_message(user_message, stream=True)
+        stream = chat.send_message(user_message, stream=True)
 
         partial_text = ""
-        for chunk in response_stream:
+        for chunk in stream:
             partial_text += (chunk.text or "")
             yield partial_text
 
@@ -136,19 +118,18 @@ async def _respond_stream(model_name, user_message, history):
 
 async def _respond_thinking(model_name, user_message, history):
     """
-    Thinking-модель:
-    - Сначала "Думаю..."
-    - Потом итог (обёрнутый в JSON, согласно системному промпту).
-    - Выделяем 'размышления' (if p.thought == True) для thinking_output.
+    Для thinking-моделей:
+    1) Выводим "Думаю..." 
+    2) После завершения — финальный ответ в формате {output: ...} + размышления.
     """
     if model_name not in MODELS:
         yield "Ошибка: модель не найдена.", ""
         return
 
     model = MODELS[model_name]
-    genai_history = _history_gradio_to_genai(history, model_name)
+    genai_history = _history_to_genai(history, model_name)
 
-    # Шаг 1: "Думаю..."
+    # Сначала "Думаю..."
     yield "Думаю...", ""
 
     try:
@@ -161,55 +142,51 @@ async def _respond_thinking(model_name, user_message, history):
         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.dumps({"output": final_text}, ensure_ascii=False)
+        # Для thinking-моделей просили итоговый ответ в {output: ...}
+        final_text_formatted = f"{{output: {final_text}}}"
 
-        yield final_text, thinking_process_text
+        yield final_text_formatted, thinking_process_text
         return
 
     except Exception as e:
         yield f"Ошибка при запросе к API: {e}", ""
         return
 
-
 ###############################################################################
-# 6. Основная функция: user_send_message
+# 6. Основная функция для ввода пользователя
 ###############################################################################
 
-async def user_send_message(user_message, history, model_name, thinking_text):
+async def user_send_message(
+    user_message: str,
+    history: list[tuple[str, str]],
+    model_name: str,
+    thinking_text: str
+):
     """
-    Gradio будет вызывать эту функцию при нажатии «Отправить» или Enter.
-    Параметры:
-      user_message : новая реплика пользователя
-      history      : [(user, assistant), ...]
-      model_name   : выбранная модель
-      thinking_text: текущее «размышление» (в отдельном Textbox)
+    Колбэк, когда пользователь отправляет запрос. 
+    Добавляем в history новый (user_msg, None), затем генерируем ответ.
     """
-    # Если пользователь ничего не ввёл
+    # Пустой ввод
     if not user_message.strip():
-        # Просто ничего не делаем
         yield history, thinking_text
         return
 
-    # Добавляем (user_message, None) в историю
+    # Добавляем (user_message, 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:
+        # Обычная модель
         partial_answer = ""
         async for chunk in _respond_stream(model_name, user_message, history):
             partial_answer = chunk
@@ -217,37 +194,179 @@ async def user_send_message(user_message, history, model_name, thinking_text):
             yield history, ""
         return
 
+###############################################################################
+# 7. Очистка диалога
+###############################################################################
+
+def clear_all():
+    """Сброс истории и размышлений."""
+    return [], ""
 
 ###############################################################################
-# 7. Доп. колбэк для смены модели
+# 8. Когда меняем модель в Dropdown
 ###############################################################################
 
-def switch_model(new_model_name, history):
+def on_model_change(selected_model, history, thinking_text):
     """
-    Когда пользователь меняет выпадающий список (Dropdown) с моделью,
-    мы добавляем в историю специальное «системное» сообщение,
-    чтобы модель знала о переключении.
+    При переключении модели добавляем в историю developer-сообщение, 
+    + добавляем дефолтный промпт этой модели (тоже developer).
     """
-    history.append(
-        ("System notice", f"User switched to a different model: {new_model_name}.")
+    new_history = history.copy()
+
+    # Cообщаем модели, что переключились (developer role)
+    new_history.append((
+        # developer-сообщение выглядит как "role=developer" → user_text
+        #   => для удобства считаем, что (user_text, None) => потом в _history_to_genai 
+        #   оно будет "role=developer", а не user. 
+        # Но проще завести отдельное поле? 
+        # Для упрощения – сделаем pseudo-user c developer-меткой.
+        # Или сделаем костыль: user_msg="[developer]Switched to X" → 
+        #   и в _history_to_genai, если строка начинается с "[developer]", 
+        #   ставим role="developer"? 
+        # 
+        # Более аккуратно — используем 3-элементный кортеж? 
+        # Но Gradio Chatbot ждёт (str, str). 
+        # 
+        # Проще при конвертации _history_to_genai:
+        #   если user_text.startswith("<developer>"), 
+        #   ставим role="developer". 
+        # 
+        # По заданию: "замени роль system на developer". 
+        # Ок, будем распознавать префикс "<developer>: ". 
+        # 
+        "<developer>: Switched to model: " + selected_model,
+        None
+    ))
+
+    # Добавляем дефолтный промпт (developer role)
+    default_prompt = DEFAULT_DEVELOPER_PROMPTS.get(
+        selected_model, 
+        "No default prompt for this model."
     )
-    return history
+    new_history.append((
+        "<developer>: " + default_prompt,
+        None
+    ))
 
+    return new_history, thinking_text
 
 ###############################################################################
-# 8. Очистка
+# 9. Функция конвертации истории с учётом developer role
 ###############################################################################
 
-def clear_all():
-    return [], ""
+def _history_to_genai_enhanced(history, model_name):
+    """
+    Улучшенная версия, отличающая developer-сообщения 
+    (префикс "<developer>: ") от user-сообщений.
+    """
+    asst_role = _assistant_role(model_name)
+    genai_history = []
+
+    for user_text, bot_text in history:
+        if user_text:
+            if user_text.startswith("<developer>: "):
+                # Считаем это developer role
+                dev_content = user_text.replace("<developer>: ", "", 1)
+                genai_history.append({"role": "developer", "parts": dev_content})
+            else:
+                # Обычный пользователь
+                genai_history.append({"role": "user", "parts": user_text})
+
+        if bot_text:
+            # Ответ ассистента / модель
+            genai_history.append({"role": asst_role, "parts": bot_text})
+
+    return genai_history
+
+# Обновим наши генераторы, чтобы использовать _history_to_genai_enhanced:
+async def _respond_stream_enh(model_name, user_message, history):
+    if model_name not in MODELS:
+        yield "Ошибка: модель не найдена."
+        return
+
+    model = MODELS[model_name]
+    genai_history = _history_to_genai_enhanced(history, model_name)
+
+    try:
+        chat = model.start_chat(history=genai_history)
+        stream = chat.send_message(user_message, stream=True)
+
+        partial_text = ""
+        for chunk in stream:
+            partial_text += (chunk.text or "")
+            yield partial_text
+
+        return
+    except Exception as e:
+        yield f"Ошибка при запросе к API: {e}"
+        return
+
+async def _respond_thinking_enh(model_name, user_message, history):
+    if model_name not in MODELS:
+        yield "Ошибка: модель не найдена.", ""
+        return
+
+    model = MODELS[model_name]
+    genai_history = _history_to_genai_enhanced(history, model_name)
+
+    # "Думаю..."
+    yield "Думаю...", ""
+
+    try:
+        chat = model.start_chat(history=genai_history)
+        response = chat.send_message(user_message, stream=False)
+
+        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 ""
 
+        # JSON-обёртка
+        final_text_formatted = f"{{output: {final_text}}}"
+
+        yield final_text_formatted, thinking_process_text
+        return
+    except Exception as e:
+        yield f"Ошибка при запросе к API: {e}", ""
+        return
+
+async def user_send_message_enh(
+    user_message: str,
+    history: list[tuple[str, str]],
+    model_name: str,
+    thinking_text: str
+):
+    if not user_message.strip():
+        yield history, thinking_text
+        return
+
+    history.append((user_message, None))
+
+    if "thinking" in model_name.lower():
+        async for (assistant_text, thought_text) in _respond_thinking_enh(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_enh(model_name, user_message, history):
+            partial_answer = chunk
+            history[-1] = (user_message, partial_answer)
+            yield history, ""
+        return
 
 ###############################################################################
-# 9. Gradio-интерфейс
+# 10. Построение интерфейса Gradio
 ###############################################################################
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Chat с Gemini (несколько моделей, «thinking»-режим, JSON-ответы)")
+    gr.Markdown("## Chat с Gemini. Поддержка developer role, переключения моделей, JSON-ответа для thinking")
 
     with gr.Row():
         model_dropdown = gr.Dropdown(
@@ -257,43 +376,57 @@ with gr.Blocks() as demo:
         )
         clear_button = gr.Button("Очистить чат")
 
-    # Состояние для хранения истории
     history_state = gr.State([])
-    # Состояние для «размышлений»
     thinking_store = gr.State("")
 
     chatbot = gr.Chatbot(label="Диалог с Gemini")
     user_input = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст...")
+    thinking_output = gr.Textbox(label="Размышления", interactive=False)
+    send_btn = gr.Button("Отправить")
 
-    thinking_output = gr.Textbox(
-        label="Размышления (для thinking-моделей)",
-        interactive=False
+    ################################################
+    # (A) Обработка переключения модели
+    ################################################
+    def handle_model_change(selected_model, history, thinking):
+        new_history, new_thinking = on_model_change(selected_model, history, thinking)
+        return new_history, new_thinking
+
+    # Когда пользователь меняет модель:
+    model_change = model_dropdown.change(
+        fn=handle_model_change,
+        inputs=[model_dropdown, history_state, thinking_store],
+        outputs=[history_state, thinking_store],
+        queue=False
+    ).then(
+        # После добавления developer-сообщений в историю → обновляем чат
+        fn=lambda h: h,
+        inputs=[history_state],
+        outputs=[chatbot],
+        queue=False
     )
 
-    send_btn = gr.Button("Отправить")
-
-    # 1) Логика при нажатии «Отправить»
+    ################################################
+    # (B) При нажатии «Отправить»
+    ################################################
     send_chain = send_btn.click(
-        fn=user_send_message,
+        fn=user_send_message_enh,
         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
     )
-    # Очищаем поле ввода (user_input)
+    # Очистка поля ввода
     send_chain.then(
         fn=lambda: "",
         inputs=[],
@@ -301,9 +434,11 @@ with gr.Blocks() as demo:
         queue=False
     )
 
-    # 2) При нажатии Enter в поле ввода
+    ################################################
+    # (C) При нажатии Enter в textbox
+    ################################################
     submit_chain = user_input.submit(
-        fn=user_send_message,
+        fn=user_send_message_enh,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
         outputs=[history_state, thinking_store],
         queue=True
@@ -320,6 +455,7 @@ with gr.Blocks() as demo:
         outputs=[thinking_output],
         queue=True
     )
+    # Очистка поля ввода
     submit_chain.then(
         fn=lambda: "",
         inputs=[],
@@ -327,19 +463,9 @@ with gr.Blocks() as demo:
         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) Кнопка «Очистить чат»
+    ################################################
+    # (D) Кнопка «Очистить»
+    ################################################
     clear_chain = clear_button.click(
         fn=clear_all,
         inputs=[],
@@ -355,11 +481,13 @@ with gr.Blocks() as demo:
         fn=lambda _: "",
         inputs=[],
         outputs=[thinking_output]
-    ).then(
+    )
+    clear_chain.then(
         fn=lambda: "",
         inputs=[],
         outputs=[user_input]
     )
 
+# Запуск
 if __name__ == "__main__":
     demo.launch()