Update app.py

app.py

@@ -1,119 +1,309 @@
+import os
 import gradio as gr
 import google.generativeai as genai
-import os
-import asyncio  # Import для асинхронности
+import asyncio
+
+###############################################################################
+# 1. Настройка окружения и инициализация модели
+###############################################################################
 
-# Безопасное получение API ключа
-GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
+# Задайте свой ключ через переменную окружения, например:
+# export GEMINI_API_KEY="ваш-ключ"
+GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc" 
 if not GEMINI_API_KEY:
-    print("Error: GEMINI_API_KEY environment variable not set.")
+    print("Error: GEMINI_API_KEY is not set.")
     exit()
 
 genai.configure(api_key=GEMINI_API_KEY)
 
-AVAILABLE_MODELS = ["gemini-1.5-flash", "gemini-1.5-pro", "gemini-2.0-flash-thinking-exp"]
+AVAILABLE_MODELS = [
+    "gemini-1.5-flash",
+    "gemini-1.5-pro",
+    "gemini-2.0-flash-thinking-exp"
+]
+
+# Создаём объекты-модели и храним их в словаре для быстрого доступа
+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}")
+
+###############################################################################
+# 2. Функции для генерации ответов
+###############################################################################
+
+def _convert_history_to_genai_format(history):
+    """
+    Gradio-хранит чат как список кортежей: [(user_msg, bot_msg), (user_msg, bot_msg), ...]
+    Для Gemini же нужна история вида [{'role': 'user', 'content': ...}, {'role': 'assistant', 'content': ...}, ...].
+
+    Возвращает список dict, где role = user/assistant, content = текст.
+    """
+    genai_history = []
+    for user_text, bot_text in history:
+        if user_text:  # сообщение от пользователя
+            genai_history.append({"role": "user", "content": user_text})
+        if bot_text:   # сообщение от ассистента
+            genai_history.append({"role": "assistant", "content": bot_text})
+    return genai_history
+
 
-# Инициализация моделей один раз при запуске приложения
-MODELS = {model_name: genai.GenerativeModel(model_name=model_name) for model_name in AVAILABLE_MODELS}
+def _convert_genai_to_chatbot_format(history):
+    """
+    Обратное преобразование:
+    [{'role': 'user', 'content': ...}, {'role': 'assistant', 'content': ...}]
+    -> [(user_msg, bot_msg), ...]
+    
+    Упрощённо полагаем, что роли будут идти поочерёдно user-assistant.
+    Если вдруг несколько подряд от одного и того же role, объединяем по очереди.
+    """
+    display_history = []
+    user_buffer = None
 
-async def respond(message, history, selected_model):
-    model = MODELS.get(selected_model)
-    if not model:
-        yield {"role": "assistant", "content": "Error: Selected model not available."}
+    for turn in history:
+        role = turn["role"]
+        content = turn["content"]
+        if role == "user":
+            # Запускаем новую пару
+            user_buffer = content
+        elif role == "assistant":
+            if user_buffer is not None:
+                display_history.append((user_buffer, content))
+                user_buffer = None
+            else:
+                # На случай, если идёт assistant без user
+                display_history.append((None, content))
+
+    # Если вдруг остался user_buffer без ответа (редко бывает)
+    if user_buffer is not None:
+        display_history.append((user_buffer, None))
+
+    return display_history
+
+
+async def _respond_stream(model_name, user_message, history):
+    """
+    Генерация ОТВЕТА С ПОСТЕПЕННЫМ стримингом (для моделей без "thinking").
+    Возвращает генератор, чтобы Gradio мог пошагово обновлять вывод.
+    """
+    if model_name not in MODELS:
+        yield "Ошибка: модель не найдена."
         return
 
+    model = MODELS[model_name]
+    genai_history = _convert_history_to_genai_format(history)
+    
     try:
-        chat = model.start_chat(history=history)
-        response_stream = chat.send_message(message, stream=True)
-        full_response = ""
+        chat = model.start_chat(history=genai_history)
+        # Запрашиваем стриминг-ответ
+        response_stream = chat.send_message(user_message, stream=True)
+
+        partial_text = ""
         for chunk in response_stream:
-            full_response += (chunk.text or "")
-            yield {"role": "assistant", "content": full_response}
+            chunk_text = chunk.text or ""
+            partial_text += chunk_text
+            # Возвращаем накопленный ответ
+            yield partial_text
     except Exception as e:
-        yield {"role": "assistant", "content": f"Error during API call: {e}"}
+        yield f"Error during API call: {e}"
 
-async def respond_thinking(message, history, selected_model):
-    if "thinking" not in selected_model:
-        yield {"role": "assistant", "content": "Thinking model не выбрана."}, ""
-        return
 
-    model = MODELS.get(selected_model)
-    if not model:
-        yield {"role": "assistant", "content": "Error: Selected model not available."}, ""
-        return
+async def _respond_thinking(model_name, user_message, history):
+    """
+    Генерация ответа для моделей с "thinking" (gemini-2.0-flash-thinking-exp).
+    1) Выдаём "Думаю..."
+    2) По завершении запроса возвращаем итоговый ответ
+    3) При желании можно вернуть размышления в отдельную переменную
+    """
+    if model_name not in MODELS:
+        return "Ошибка: модель не найдена.", ""
 
-    yield {"role": "assistant", "content": "Думаю..."}
+    model = MODELS[model_name]
+    genai_history = _convert_history_to_genai_format(history)
+
+    # 1) "Думаю..."
+    yield "Думаю...", ""  # (ответ в чат, размышления)
 
     try:
-        response = model.generate_content(message)
-        thinking_process_text = ""
-        model_response_text = ""
+        chat = model.start_chat(history=genai_history)
+        response = chat.send_message(user_message, stream=False)  # без stream
 
+        # Для thinking-моделей ответ приходит в виде parts
+        thinking_process_text = ""
+        final_text = ""
         if response.candidates:
             for part in response.candidates[0].content.parts:
-                if hasattr(part, 'thought') and part.thought == True:
-                    thinking_process_text += f"Model Thought:\n{part.text}\n\n"
+                # Если в part есть .thought == True, считаем это «размышлениями»
+                if getattr(part, "thought", False):
+                    thinking_process_text += part.text or ""
                 else:
-                    model_response_text += (part.text or "")
+                    final_text += part.text or ""
+
+        # Возвращаем итоговый ответ
+        yield final_text, thinking_process_text
 
-        yield {"role": "assistant", "content": model_response_text}, thinking_process_text
     except Exception as e:
-        yield {"role": "assistant", "content": f"Error during API call: {e}"}, f"Error during API call: {e}"
+        yield f"Error during API call: {e}", f"Error: {e}"
+
+
+###############################################################################
+# 3. Основная функция (колбэк для Gradio), обрабатывающая один «шаг» диалога
+###############################################################################
+
+async def user_send_message(user_message, history, model_name, thinking_store):
+    """
+    Эта функция вызывается при нажатии кнопки "Отправить" (или Enter).
+    Параметры:
+      user_message: текст от пользователя
+      history:    текущее состояние чата [(user_msg, bot_msg), ...]
+      model_name: выбранная модель
+      thinking_store: текущее значение поля "Размышления"
 
-async def process_message(message, history, model_name):
-    if "thinking" in model_name:
-        async for response, thinking in respond_thinking(message, history, model_name):
-            yield (response, thinking)
+    Возвращаем:
+      (updated_history, updated_thinking_store)
+    """
+    if not user_message.strip():
+        # Пустая строка – ничего не делаем
+        return history, thinking_store
+
+    # Добавляем в историю шаг пользователя
+    history.append((user_message, None))  # Пока ассистента нет
+
+    # Если модель thinking
+    if "thinking" in model_name.lower():
+        # Асинхронный генератор: сначала "Думаю...", потом финальный ответ
+        async for (assistant_text, thinking_text) in _respond_thinking(model_name, user_message, history):
+            # Обновляем последний элемент истории (None -> assistant_text)
+            history[-1] = (user_message, assistant_text)
+            # Обновляем thinking_output
+            yield history, thinking_text
     else:
-        async for response in respond(message, history, model_name):
-            yield (response, "")
+        # Стриминговый ответ
+        partial_answer = ""
+        async for partial_chunk in _respond_stream(model_name, user_message, history):
+            partial_answer = partial_chunk
+            history[-1] = (user_message, partial_answer)
+            # В режиме без thinking просто очищаем thinking_store
+            yield history, ""
+
+    # Когда генерация закончена, возвращаем итог
+    return history, thinking_store
+
+
+###############################################################################
+# 4. Колбэки для кнопки "Очистить"
+###############################################################################
 
-def clear_thinking():
-    return ""
+def clear_all():
+    """
+    Полная очистка истории чата и поля «Размышления».
+    Возвращает пустые значения для State и компонентов.
+    """
+    return [], ""  # history, thinking_store
+
+
+###############################################################################
+# 5. Определяем Gradio-интерфейс
+###############################################################################
 
 with gr.Blocks() as demo:
-    gr.Markdown("# Gemini Chatbot с режимом размышления")
+    gr.Markdown("## Gemini Chatbot с сохранением истории и «thinking»-режимом")
+
+    # Храним историю диалога (список кортежей) в gr.State
+    history_state = gr.State([])
+
+    # Храним «размышления» (thinking) в отдельном тексте
+    thinking_store = gr.State("")
 
     with gr.Row():
-        model_selection = gr.Dropdown(
-            AVAILABLE_MODELS, value="gemini-1.5-flash", label="Выберите модель Gemini"
+        model_dropdown = gr.Dropdown(
+            choices=AVAILABLE_MODELS,
+            value="gemini-1.5-flash",
+            label="Выберите модель"
         )
+        clear_button = gr.Button("Очистить чат")
 
-    thinking_output = gr.Code(label="Процесс размышления (для моделей с размышлением)")
+    # Сам чат: компонент Chatbot для отображения диалога
+    chatbot = gr.Chatbot(label="Диалог с Gemini")
 
-    chatbot = gr.ChatInterface(
-        process_message,  # Функция обработки сообщений передается сюда
-        additional_inputs=[model_selection],
-        title="Gemini Chat",
-        description="Общайтесь с моделями Gemini от Google.",
-        type="messages"
+    # Поле ввода сообщения
+    user_input = gr.Textbox(
+        label="Ваш вопрос/сообщение",
+        placeholder="Введите вопрос...",
     )
 
-    with gr.Row():
-        clear_button = gr.Button("Очистить")
-
-    def change_chatbot(model_name):
-        return gr.ChatInterface.update()
+    # Поле для показа «Размышлений» (если модель thinking)
+    thinking_output = gr.Textbox(
+        label="Размышления модели (только для gemini-2.0-flash-thinking-exp)",
+        interactive=False
+    )
 
-    async def update_thinking(history):
-        if history and history[-1][0]["role"] == "assistant" and len(history[-1]) > 1:
-            return history[-1][1]
-        return ""
+    # Кнопка «Отправить»
+    send_btn = gr.Button("Отправить")
 
-    model_selection.change(
-        change_chatbot,
-        inputs=[model_selection],
+    ############################################################################
+    # Связываем действия пользователя с колбэками
+    ############################################################################
+    # 1. При нажатии «Отправить» запускаем user_send_message
+    #    - обновляем историю и thinking_output
+    #    - одновременно обновляем отображение чата
+    send_btn.click(
+        fn=user_send_message,
+        inputs=[user_input, history_state, model_dropdown, thinking_store],
+        outputs=[history_state, thinking_store],
+        queue=True,
+    ).then(
+        # После обновления history_state, пересчитываем отображение чата
+        fn=lambda h: h,
+        inputs=[history_state],
         outputs=[chatbot],
+        preprocess=False,
+        postprocess=True
+    ).then(
+        # После этого выводим текущее состояние thinking_store
+        fn=lambda t: t,
+        inputs=[thinking_store],
+        outputs=[thinking_output],
     )
 
-    chatbot.change(
-        update_thinking,
-        inputs=[chatbot],
-        outputs=[thinking_output]
+    # 2. При нажатии Enter в текстовом поле — та же логика, что и на кнопке
+    user_input.submit(
+        fn=user_send_message,
+        inputs=[user_input, history_state, model_dropdown, thinking_store],
+        outputs=[history_state, thinking_store],
+        queue=True,
+    ).then(
+        fn=lambda h: h,
+        inputs=[history_state],
+        outputs=[chatbot],
+        preprocess=False,
+        postprocess=True
+    ).then(
+        fn=lambda t: t,
+        inputs=[thinking_store],
+        outputs=[thinking_output],
     )
 
-    clear_button.click(lambda: None, None, chatbot, queue=False)
-    clear_button.click(clear_thinking, outputs=[thinking_output], queue=False)
+    # 3. Кнопка «Очистить»
+    clear_button.click(
+        fn=clear_all,
+        inputs=[],
+        outputs=[history_state, thinking_store],
+        queue=False
+    ).then(
+        fn=lambda h: h,
+        inputs=[history_state],
+        outputs=chatbot,
+    ).then(
+        fn=lambda _: "",
+        inputs=[],
+        outputs=thinking_output
+    )
 
+###############################################################################
+# 6. Запуск
+###############################################################################
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()