Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import gradio as gr
 from gradio import ChatMessage
-from typing import Iterator
+from typing import Iterator, List, Tuple, Optional # Добавили типы
 import google.generativeai as genai
 import time # Import time module for potential debugging/delay
 
@@ -10,14 +10,20 @@ print("add API key")
 
 # get Gemini API Key from the environ variable
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+# --- Безопасность: Добавим проверку наличия ключа ---
+if not GEMINI_API_KEY:
+    raise ValueError("GEMINI_API_KEY environment variable not set!")
+# ----------------------------------------------------
 genai.configure(api_key=GEMINI_API_KEY)
 
+
 print("add API key complete ")
 print("add model")
 
-used_model = "gemma-3-27b-it"
-
-# we will be using the Gemini 2.0 Flash model with Thinking capabilities
+# --- Используем более подходящую модель для чата, если Gemma IT недоступна или неоптимальна ---
+# used_model = "gemma-3-27b-it" # Gemma может не поддерживать role='model' явно, Gemini Pro лучше для чата
+used_model = "gemini-1.5-flash-latest" # Рекомендуется для чата с историей
+# ------------------------------------------------------------------------------------------
 model = genai.GenerativeModel(used_model)
 
 print(f"add model {used_model} complete\n")
@@ -29,176 +35,624 @@ def format_chat_history(messages: list) -> list:
     """
     formatted_history = []
     for message in messages:
-        #print(f"t1 {message}")
-        # Skip thinking messages (messages with metadata)
-        #if not (message.get("role") == "assistant" and "metadata" in message):
-        #    print(f"t2 {message}")
-        #    formatted_history.append({
-        #        "role": "user" if message.get("role") == "user" else "assistant",
-        #        "parts": [message.get("content", "")]
-        #    })
-            
-        #print(f"t2 {message}")
-
-        if message.get("role") == "user" :
+        # Пропускаем пустые сообщения или сообщения без контента (на всякий случай)
+        content = message.get("content")
+        if not content:
+            continue
+
+        role = message.get("role")
+        if role == "user":
             formatted_history.append({
-            "role": "user",
-            "parts": [message.get("content", "")]
+                "role": "user",
+                "parts": [content]
             })
-        elif message.get("role") == "assistant" :
+        elif role == "assistant":
+             # Пропускаем "thinking" сообщения, если они есть (хотя мы их убрали)
+            if not message.get("metadata"):
+                 formatted_history.append({
+                    "role": "model", # Gemini API ожидает 'model' для ответов ассистента
+                    "parts": [content]
+                 })
+        # Игнорируем другие возможные роли или типы сообщений
+    print(f"Formatted history length: {len(formatted_history)}")
+    print("return formatted history")
+    return formatted_history
+
+# --- Новая функция для Undo ---
+def undo_last(history: List[ChatMessage]) -> Tuple[List[ChatMessage], List[ChatMessage]]:
+    """
+    Removes the last user message and the last assistant response.
+    Returns the updated history and the pair that was removed (for potential redo).
+    """
+    print("\nAttempting Undo")
+    undone_pair = []
+    if len(history) >= 2:
+        # Проверяем, что последние два сообщения - это user и assistant
+        if history[-2].role == "user" and history[-1].role == "assistant":
+            print("Found user/assistant pair to undo.")
+            undone_pair = history[-2:]
+            history = history[:-2]
+        else:
+            print("Last two messages are not a user/assistant pair. Cannot undo.")
+    else:
+        print("Not enough messages in history to undo.")
+    return history, undone_pair
+# -----------------------------
+
+# --- Новая функция для Redo ---
+def redo_last(history: List[ChatMessage], undone_pair: List[ChatMessage]) -> Tuple[List[ChatMessage], List[ChatMessage]]:
+    """
+    Restores the last undone user/assistant pair.
+    Returns the updated history and clears the undone pair state.
+    """
+    print("\nAttempting Redo")
+    if undone_pair:
+        print("Found undone pair to redo.")
+        history.extend(undone_pair)
+        undone_pair = [] # Очищаем состояние после восстановления
+    else:
+        print("No undone pair available to redo.")
+    return history, undone_pair
+# -----------------------------
+
+# --- Новая функция для Regenerate ---
+def regenerate_response(history: List[ChatMessage]) -> Tuple[List[ChatMessage], Optional[str], List[ChatMessage]]:
+    """
+    Removes the last assistant response and returns the history ending
+    with the last user message, and the content of that user message.
+    Also returns an empty list to clear the undone_state.
+    """
+    print("\nAttempting Regenerate")
+    last_user_message_content = None
+    history_for_regen = history[:] # Создаем копию
+
+    # Ищем последний user message с конца
+    last_user_index = -1
+    for i in range(len(history_for_regen) - 1, -1, -1):
+        if history_for_regen[i].role == "user":
+            last_user_index = i
+            break
+
+    if last_user_index != -1:
+        print(f"Found last user message at index {last_user_index}")
+        last_user_message_content = history_for_regen[last_user_index].content
+        # Обрезаем историю до этого user message включительно
+        history_for_regen = history_for_regen[:last_user_index + 1]
+        print("Prepared history for regeneration.")
+    else:
+        print("No user message found to regenerate from.")
+        history_for_regen = history # Возвращаем оригинал, если не нашли
+
+    # Регенерация инвалидирует предыдущее undo
+    return history_for_regen, last_user_message_content, [] # Возвращаем историю и контент пользователя, очищаем undone_state
+# ----------------------------------
+
+
+def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
+    print("start model response stream")
+    """
+    Streams response with conversation history support for text input only.
+    (Убрали логику "Thinking" для упрощения и совместимости)
+    """
+    # --- Добавили проверку на None или пустую строку ---
+    if not user_message or not user_message.strip():
+        print("Empty or None user message received. Skipping API call.")
+        # Можно добавить сообщение об ошибке или просто ничего не делать
+        # messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message."))
+        yield messages # Просто возвращаем текущее состояние
+        return
+    # -----------------------------------------------------
+
+    try:
+        print(f"\n=== New Request (Text) ===")
+        print(f"User message: {user_message}") # Может быть None при регенерации, но мы проверили выше
+
+        # Форматируем историю *перед* добавлением нового ответа ассистента
+        # Важно: история для API должна содержать только завершенные user/model пары
+        chat_history = format_chat_history(messages) # messages здесь уже содержит user message, с которого регенерируем, или новый
+
+        print("Chat parameter")
+        chat = model.start_chat(history=chat_history)
+        print("Start response")
+        # Отправляем *только* последнее сообщение пользователя (которое уже есть в messages)
+        # Gemini API берет историю из start_chat, а последнее сообщение из send_message
+        # Важно: user_message здесь - это то, на что мы генерируем ответ
+        response = chat.send_message(user_message, stream=True)
+
+        response_buffer = ""
+        # Добавляем пустое сообщение ассистента, которое будем обновлять
+        messages.append(ChatMessage(role="assistant", content=""))
+        yield messages # Показываем пустой контейнер для ответа
+
+        print("Streaming response...")
+        for chunk in response:
+            # --- Обработка возможных ошибок в чанке ---
+            try:
+                # Проверяем наличие candidates и parts
+                if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
+                    current_chunk_text = chunk.candidates[0].content.parts[0].text
+                    # print(f"Chunk: '{current_chunk_text}'") # Отладка чанков
+                    response_buffer += current_chunk_text
+                    messages[-1] = ChatMessage(
+                        role="assistant",
+                        content=response_buffer
+                    )
+                else:
+                    # Обработка случая, когда структура чанка неожиданная
+                    print(f"Warning: Unexpected chunk structure: {chunk}")
+                    # Можно добавить запасной вариант или пропустить чанк
+                    if hasattr(chunk, 'text'): # Иногда API возвращает просто text
+                         response_buffer += chunk.text
+                         messages[-1] = ChatMessage(role="assistant", content=response_buffer)
+
+            except (AttributeError, IndexError, ValueError) as chunk_err:
+                 # Ловим ошибки доступа к атрибутам/индексам или другие проблемы чанка
+                 print(f"Error processing chunk: {chunk_err}")
+                 print(f"Problematic chunk data: {chunk}")
+                 # Можно прервать или продолжить, добавив сообщение об ошибке
+                 messages[-1] = ChatMessage(
+                     role="assistant",
+                     content=response_buffer + f"\n\n[Ошибка обработки части ответа: {chunk_err}]"
+                 )
+                 yield messages
+                 return # Прерываем стриминг при ошибке в чанке
+            # -------------------------------------------
+
+            # time.sleep(0.05) # Раскомментируйте для замедления стриминга (отладка)
+            yield messages # Отправляем обновленное сообщение в интерфейс
+
+        print(f"\n=== Final Response ===\n{response_buffer}")
+        # Проверка на пустой финальный ответ (если модель ничего не вернула)
+        if not response_buffer.strip() and len(messages) > 0 and messages[-1].role == "assistant":
+             messages[-1] = ChatMessage(
+                 role="assistant",
+                 content="[Модель не дала ответа]"
+             )
+             yield messages
+
+
+    # --- Улучшенная обработка ошибок API ---
+    except Exception as e:
+        error_message = f"Произошла ошибка при обращении к Gemini API: {str(e)}"
+        print(f"\n=== Error ===\n{error_message}")
+        # Пытаемся добавить сообщение об ошибке в чат, если возможно
+        if messages and isinstance(messages[-1], ChatMessage) and messages[-1].role == "assistant" and messages[-1].content == "":
+            # Если последний элемент - пустое сообщение ассистента, заменяем его ошибкой
+             messages[-1] = ChatMessage(role="assistant", content=error_message)
+        else:
+            # Иначе добавляем новое сообщение с ошибкой
+             messages.append(ChatMessage(role="assistant", content=error_message))
+        yield messages
+    # -------------------------------------
+
+
+def user_message(msg: str, history: list) -> tuple[str, list, list]:
+    """Adds user message to chat history and clears the undone state."""
+    print(f"\nUser message added: '{msg}'")
+    history.append(ChatMessage(role="user", content=msg))
+    # Новое сообщение пользователя инвалидирует предыдущее undo
+    return "", history, [] # Возвращаем пустую строку, историю и пустой undone_state
+
+# Create the Gradio interface
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo:
+    gr.Markdown("# Chat with " + used_model)
+
+    gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fzelk12%2FGemini-2">
+               <img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fzelk12%2FGemini-2&countColor=%23263759" />
+               </a>""")
+
+    # --- Добавляем состояние для Undo/Redo ---
+    undone_state = gr.State([])
+    # --------------------------------------
+
+    chatbot = gr.Chatbot(
+        [], # Начинаем с пустого списка
+        type="messages",
+        label=used_model + " Chatbot (Streaming Output)",
+        render_markdown=True,
+        scale=1,
+        # editable="all", # editable="all" может мешать undo/redo, лучше отключить или user
+        editable=False,
+        avatar_images=(None,"https://lh3.googleusercontent.com/oxz0sUBF0iYoN4VvhqWTmux-cxfD1rxuYkuFEfm1SFaseXEsjjE4Je_C_V3UQPuJ87sImQK3HfQ3RXiaRnQetjaZbjJJUkiPL5jFJ1WRl5FKJZYibUA=w214-h214-n-nu"),
+        height=600 # Задаем высоту для лучшего вида
+    )
+
+    with gr.Row(equal_height=False): # equal_height=False может быть лучше для кнопок разной высоты
+        input_box = gr.Textbox(
+            lines=1,
+            label="Chat Message",
+            placeholder="Type your message here...",
+            scale=4
+        )
+
+        # --- Группируем кнопки управления ---
+        with gr.Column(scale=1, min_width=150): # Даем минимальную ширину колонке с кнопками
+             submit_button = gr.Button("Submit", scale=1, variant="primary") # Выделяем основную кнопку
+             with gr.Row():
+                 undo_button = gr.Button("Undo", scale=1)
+                 redo_button = gr.Button("Redo", scale=1)
+             regenerate_button = gr.Button("Regenerate", scale=1)
+             clear_button = gr.Button("Clear Chat", scale=1)
+
+
+    # --- Переименовали кнопки для ясности ---
+    # undo_button = test_button
+    # redo_button = test1_button
+    # regenerate_button = test2_button
+    # ----------------------------------------
+
+    # Add example prompts
+    example_prompts = [
+        ["Write a short poem about the sunset."],
+        ["Explain the theory of relativity in simple terms."],
+        ["If a train leaves Chicago at 6am traveling at 60mph, and another train leaves New York at 8am traveling at 80mph, at what time will they meet?"],
+        ["Summarize the plot of Hamlet."],
+        ["Write a haiku about a cat."]
+    ]
+
+    gr.Examples(
+        examples=example_prompts,
+        inputs=input_box,
+        label="Examples: Try these prompts!", # Убрали про thinking, т.к. убрали его вывод
+        examples_per_page=5
+    )
+
+    # --- Обновленные обработчики событий ---
+    msg_store = gr.State("") # Store for preserving user message during processing
+
+    # --- Обработчик для Enter в input_box ---
+    input_box.submit(
+        user_message,  # 1. Добавить сообщение пользователя в историю, очистить undone_state
+        inputs=[input_box, chatbot],
+        outputs=[input_box, chatbot, undone_state] # input_box очищается первым возвращаемым значением ""
+    ).then(
+        lambda x: (x, x), # 2. Скопировать сообщение из history[-1].content в msg_store для stream_gemini_response
+        inputs=[chatbot], # Берем историю, где последнее сообщение - пользовательское
+        outputs=[msg_store, msg_store], # Копируем контент последнего сообщения (пользователя)
+        fn=lambda history: (history[-1].content, history[-1].content) if history and history[-1].role == "user" else (None, None)
+    ).then(
+        stream_gemini_response,  # 3. Сгенерировать и стримить ответ
+        inputs=[msg_store, chatbot], # Используем сохраненное сообщение и обновленную историю
+        outputs=[chatbot]
+    )
+
+    # --- Обработчик для кнопки Submit ---
+    submit_button.click(
+        user_message, # 1. Добавить сообщение пользователя, очистить undone_state
+        inputs=[input_box, chatbot],
+        outputs=[input_box, chatbot, undone_state]
+    ).then(
+        lambda x: (x, x), # 2. Скопировать сообщение пользователя в msg_store
+        inputs=[chatbot],
+        outputs=[msg_store, msg_store],
+        fn=lambda history: (history[-1].content, history[-1].content) if history and history[-1].role == "user" else (None, None)
+    ).then(
+        stream_gemini_response, # 3. Сгенерировать ответ
+        inputs=[msg_store, chatbot],
+        outputs=[chatbot]
+    )
+
+    # --- Обработчик для кнопки Clear ---
+    clear_button.click(
+        lambda: ([], "", [], ""), # Очищаем chatbot, input_box, msg_store, undone_state
+        outputs=[chatbot, input_box, msg_store, undone_state],
+        queue=False # Очистка быстрая, очередь не нужна
+    )
+
+    # --- Обработчик для кнопки Undo ---
+    undo_button.click(
+        undo_last,
+        inputs=[chatbot],
+        outputs=[chatbot, undone_state], # Обновляем историю и сохраняем удаленную пару
+        queue=False # Действие быстрое
+    )
+
+    # --- Обработчик для кнопки Redo ---
+    redo_button.click(
+        redo_last,
+        inputs=[chatbot, undone_state], # Нужна история и сохраненная пара
+        outputs=[chatbot, undone_state], # Обновляем историю и очищаем сохраненную пару
+        queue=False # Действие быстрое
+    )
+
+    # --- Обработчик для кнопки Regenerate ---
+    regenerate_button.click(
+        regenerate_response, # 1. Подготовить историю и получить сообщение пользователя, очистить undone_state
+        inputs=[chatbot],
+        outputs=[chatbot, msg_store, undone_state] # Обновляем историю, помещаем user msg в msg_store, очищаем undone
+    ).then(
+        stream_gemini_response, # 2. Сгенерировать новый ответ на user msg из msg_store
+        inputs=[msg_store, chatbot], # Используем user msg из msg_store и подготовленную историю
+        outputs=[chatbot]
+    )
+    # -------------------------------------
+
+    gr.Markdown(
+        """
+        <br><br><br>
+        ---
+        ### About this Chatbot
+        **Try out the example prompts below!**
+        **Key Features:**
+        *   Powered by Google's **""" + used_model + """** model.
+        *   Supports **conversation history**, **undo**, **redo**, and **regenerate**.
+        *   Uses **streaming** for responses.
+        **Instructions:**
+        1.  Type your message or select an example.
+        2.  Press Enter or click Submit.
+        3.  Use Undo/Redo to correct mistakes or explore alternatives.
+        4.  Use Regenerate to get a different response to your last message.
+        5.  Use Clear Chat to start over.
+        """
+    )
+
+
+# Launch the interface
+if __name__ == "__main__":
+    # --- Добавим возможность запускать на публичном URL через share=True ---
+    # demo.launch(debug=True) # Для локальной отладки
+    demo.launch(share=True) # Для получения публичной ссылки (если нужно)
+    # demo.launch() # Стандартный локальный запуск
+    # -----------------------------------------------------------------import os
+import gradio as gr
+from gradio import ChatMessage
+from typing import Iterator, List, Tuple, Optional # Добавили типы
+import google.generativeai as genai
+import time # Import time module for potential debugging/delay
+
+print("import library complete")
+print("add API key")
+
+# get Gemini API Key from the environ variable
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+# --- Безопасность: Добавим проверку наличия ключа ---
+if not GEMINI_API_KEY:
+    raise ValueError("GEMINI_API_KEY environment variable not set!")
+# ----------------------------------------------------
+genai.configure(api_key=GEMINI_API_KEY)
+
+
+print("add API key complete ")
+print("add model")
+
+# --- Используем более подходящую модель для чата, если Gemma IT недоступна или неоптимальна ---
+# used_model = "gemma-3-27b-it" # Gemma может не поддерживать role='model' явно, Gemini Pro лучше для чата
+used_model = "gemini-1.5-flash-latest" # Рекомендуется для чата с историей
+# ------------------------------------------------------------------------------------------
+model = genai.GenerativeModel(used_model)
+
+print(f"add model {used_model} complete\n")
+
+def format_chat_history(messages: list) -> list:
+    print("\nstart format history")
+    """
+    Formats the chat history into a structure Gemini can understand
+    """
+    formatted_history = []
+    for message in messages:
+        # Пропускаем пустые сообщения или сообщения без контента (на всякий случай)
+        content = message.get("content")
+        if not content:
+            continue
+
+        role = message.get("role")
+        if role == "user":
             formatted_history.append({
-            "role": "model",
-            "parts": [message.get("content", "")]
+                "role": "user",
+                "parts": [content]
             })
-        
-    #print(f"t3 {formatted_history}")
+        elif role == "assistant":
+             # Пропускаем "thinking" сообщения, если они есть (хотя мы их убрали)
+            if not message.get("metadata"):
+                 formatted_history.append({
+                    "role": "model", # Gemini API ожидает 'model' для ответов ассистента
+                    "parts": [content]
+                 })
+        # Игнорируем другие возможные роли или типы сообщений
+    print(f"Formatted history length: {len(formatted_history)}")
     print("return formatted history")
     return formatted_history
 
+# --- Новая функция для Undo ---
+def undo_last(history: List[ChatMessage]) -> Tuple[List[ChatMessage], List[ChatMessage]]:
+    """
+    Removes the last user message and the last assistant response.
+    Returns the updated history and the pair that was removed (for potential redo).
+    """
+    print("\nAttempting Undo")
+    undone_pair = []
+    if len(history) >= 2:
+        # Проверяем, что последние два сообщения - это user и assistant
+        if history[-2].role == "user" and history[-1].role == "assistant":
+            print("Found user/assistant pair to undo.")
+            undone_pair = history[-2:]
+            history = history[:-2]
+        else:
+            print("Last two messages are not a user/assistant pair. Cannot undo.")
+    else:
+        print("Not enough messages in history to undo.")
+    return history, undone_pair
+# -----------------------------
+
+# --- Новая функция для Redo ---
+def redo_last(history: List[ChatMessage], undone_pair: List[ChatMessage]) -> Tuple[List[ChatMessage], List[ChatMessage]]:
+    """
+    Restores the last undone user/assistant pair.
+    Returns the updated history and clears the undone pair state.
+    """
+    print("\nAttempting Redo")
+    if undone_pair:
+        print("Found undone pair to redo.")
+        history.extend(undone_pair)
+        undone_pair = [] # Очищаем состояние после восстановления
+    else:
+        print("No undone pair available to redo.")
+    return history, undone_pair
+# -----------------------------
+
+# --- Новая функция для Regenerate ---
+def regenerate_response(history: List[ChatMessage]) -> Tuple[List[ChatMessage], Optional[str], List[ChatMessage]]:
+    """
+    Removes the last assistant response and returns the history ending
+    with the last user message, and the content of that user message.
+    Also returns an empty list to clear the undone_state.
+    """
+    print("\nAttempting Regenerate")
+    last_user_message_content = None
+    history_for_regen = history[:] # Создаем копию
+
+    # Ищем последний user message с конца
+    last_user_index = -1
+    for i in range(len(history_for_regen) - 1, -1, -1):
+        if history_for_regen[i].role == "user":
+            last_user_index = i
+            break
+
+    if last_user_index != -1:
+        print(f"Found last user message at index {last_user_index}")
+        last_user_message_content = history_for_regen[last_user_index].content
+        # Обрезаем историю до этого user message включительно
+        history_for_regen = history_for_regen[:last_user_index + 1]
+        print("Prepared history for regeneration.")
+    else:
+        print("No user message found to regenerate from.")
+        history_for_regen = history # Возвращаем оригинал, если не нашли
+
+    # Регенерация инвалидирует предыдущее undo
+    return history_for_regen, last_user_message_content, [] # Возвращаем историю и контент пользователя, очищаем undone_state
+# ----------------------------------
+
+
 def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
     print("start model response stream")
     """
-    Streams thoughts and response with conversation history support for text input only.
+    Streams response with conversation history support for text input only.
+    (Убрали логику "Thinking" для упрощения и совместимости)
     """
-    if not user_message.strip(): # Robust check: if text message is empty or whitespace
-        messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message. Empty input is not allowed.")) # More specific message
-        yield messages
-        print("Empty text message")
+    # --- Добавили проверку на None или пустую строку ---
+    if not user_message or not user_message.strip():
+        print("Empty or None user message received. Skipping API call.")
+        # Можно добавить сообщение об ошибке или просто ничего не делать
+        # messages.append(ChatMessage(role="assistant", content="Please provide a non-empty text message."))
+        yield messages # Просто возвращаем текущее состояние
         return
+    # -----------------------------------------------------
 
     try:
         print(f"\n=== New Request (Text) ===")
-        print(f"User message: {user_message}")
+        print(f"User message: {user_message}") # Может быть None при регенерации, но мы проверили выше
 
-        # Format chat history for Gemini
-        chat_history = format_chat_history(messages)
+        # Форматируем историю *перед* добавлением нового ответа ассистента
+        # Важно: история для API должна содержать только завершенные user/model пары
+        chat_history = format_chat_history(messages) # messages здесь уже содержит user message, с которого регенерируем, или новый
 
-        #print(f"hist {chat_history}")
-        
-        # Initialize Gemini chat
         print("Chat parameter")
         chat = model.start_chat(history=chat_history)
         print("Start response")
+        # Отправляем *только* последнее сообщение пользователя (которое уже есть в messages)
+        # Gemini API берет историю из start_chat, а последнее сообщение из send_message
+        # Важно: user_message здесь - это то, на что мы генерируем ответ
         response = chat.send_message(user_message, stream=True)
 
-        # Initialize buffers and flags
-        thought_buffer = ""
         response_buffer = ""
-        #thinking_complete = False
-
-        # Add initial thinking message
-        #messages.append(
-        #    ChatMessage(
-        #        role="assistant",
-        #        content="",
-        #        metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
-        #    )
-        #)
-        
-        messages.append(
-            ChatMessage(
-                role="assistant",
-                content=response_buffer
-            )
-        )
-        #print(f"mes {messages} \n\nhis {chat_history}")
-        
-        thinking_complete = True
+        # Добавляем пустое сообщение ассистента, которое будем обновлять
+        messages.append(ChatMessage(role="assistant", content=""))
+        yield messages # Показываем пустой контейнер для ответа
 
+        print("Streaming response...")
         for chunk in response:
-            print("chunk start")
-            parts = chunk.candidates[0].content.parts
-            print("parts")
-            current_chunk = parts[0].text
-            print("current_chunk")
-            
-            print(f"\n=========\nparts len: {len(parts)}\n\nparts: {parts}\n\ncurrent chunk: {current_chunk}\n=========\n")
-            
-            if len(parts) == 2 and not thinking_complete:
-                # Complete thought and start response
-                thought_buffer += current_chunk
-                print(f"\n=== Complete Thought ===\n{thought_buffer}")
-
-                messages[-1] = ChatMessage(
-                    role="assistant",
-                    content=thought_buffer,
-                    metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
-                )
-                yield messages
-
-                # Start response
-                response_buffer = parts[1].text
-                print(f"\n=== Starting Response ===\n{response_buffer}")
-
-                messages.append(
-                    ChatMessage(
+            # --- Обработка возможных ошибок в чанке ---
+            try:
+                # Проверяем наличие candidates и parts
+                if chunk.candidates and chunk.candidates[0].content and chunk.candidates[0].content.parts:
+                    current_chunk_text = chunk.candidates[0].content.parts[0].text
+                    # print(f"Chunk: '{current_chunk_text}'") # Отладка чанков
+                    response_buffer += current_chunk_text
+                    messages[-1] = ChatMessage(
                         role="assistant",
                         content=response_buffer
                     )
-                )
-                thinking_complete = True
-
-            elif thinking_complete:
-                # Stream response
-                response_buffer += current_chunk
-                print(f"\n=== Response Chunk ===\n{current_chunk}")
-
-                messages[-1] = ChatMessage(
-                    role="assistant",
-                    content=response_buffer
-                )
-
-            else:
-                # Stream thinking
-                thought_buffer += current_chunk
-                print(f"\n=== Thinking Chunk ===\n{current_chunk}")
-
-                messages[-1] = ChatMessage(
-                    role="assistant",
-                    content=thought_buffer,
-                    metadata={"title": "⚙️ Thinking: *The thoughts produced by the model are experimental"}
-                )
-            #time.sleep(0.05) #Optional: Uncomment this line to add a slight delay for debugging/visualization of streaming. Remove for final version
-            print("Response end")
-            yield messages
+                else:
+                    # Обработка случая, когда структура чанка неожиданная
+                    print(f"Warning: Unexpected chunk structure: {chunk}")
+                    # Можно добавить запасной вариант или пропустить чанк
+                    if hasattr(chunk, 'text'): # Иногда API возвращает просто text
+                         response_buffer += chunk.text
+                         messages[-1] = ChatMessage(role="assistant", content=response_buffer)
+
+            except (AttributeError, IndexError, ValueError) as chunk_err:
+                 # Ловим ошибки доступа к атрибутам/индексам или другие проблемы чанка
+                 print(f"Error processing chunk: {chunk_err}")
+                 print(f"Problematic chunk data: {chunk}")
+                 # Можно прервать или продол��ить, добавив сообщение об ошибке
+                 messages[-1] = ChatMessage(
+                     role="assistant",
+                     content=response_buffer + f"\n\n[Ошибка обработки части ответа: {chunk_err}]"
+                 )
+                 yield messages
+                 return # Прерываем стриминг при ошибке в чанке
+            # -------------------------------------------
+
+            # time.sleep(0.05) # Раскомментируйте для замедления стриминга (отладка)
+            yield messages # Отправляем обновленное сообщение в интерфейс
 
         print(f"\n=== Final Response ===\n{response_buffer}")
+        # Проверка на пустой финальный ответ (если модель ничего не вернула)
+        if not response_buffer.strip() and len(messages) > 0 and messages[-1].role == "assistant":
+             messages[-1] = ChatMessage(
+                 role="assistant",
+                 content="[Модель не дала ответа]"
+             )
+             yield messages
+
 
+    # --- Улучшенная обработка ошибок API ---
     except Exception as e:
-        print(f"\n=== Error ===\n{str(e)}")
-        messages.append(
-            ChatMessage(
-                role="assistant",
-                content=f"I apologize, but I encountered an error: {str(e)}"
-            )
-        )
+        error_message = f"Произошла ошибка при обращении к Gemini API: {str(e)}"
+        print(f"\n=== Error ===\n{error_message}")
+        # Пытаемся добавить сообщение об ошибке в чат, если возможно
+        if messages and isinstance(messages[-1], ChatMessage) and messages[-1].role == "assistant" and messages[-1].content == "":
+            # Если последний элемент - пустое сообщение ассистента, заменяем его ошибкой
+             messages[-1] = ChatMessage(role="assistant", content=error_message)
+        else:
+            # Иначе добавляем новое сообщение с ошибкой
+             messages.append(ChatMessage(role="assistant", content=error_message))
         yield messages
+    # -------------------------------------
+
 
-def user_message(msg: str, history: list) -> tuple[str, list]:
-    """Adds user message to chat history"""
+def user_message(msg: str, history: list) -> tuple[str, list, list]:
+    """Adds user message to chat history and clears the undone state."""
+    print(f"\nUser message added: '{msg}'")
     history.append(ChatMessage(role="user", content=msg))
-    return "", history
+    # Новое сообщение пользователя инвалидирует предыдущее undo
+    return "", history, [] # Возвращаем пустую строку, историю и пустой undone_state
 
 # Create the Gradio interface
-with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo: # Using Soft theme with adjusted hues for a refined look
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", neutral_hue="neutral")) as demo:
     gr.Markdown("# Chat with " + used_model)
 
-    
     gr.HTML("""<a href="https://visitorbadge.io/status?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fzelk12%2FGemini-2">
                <img src="https://api.visitorbadge.io/api/combined?path=https%3A%2F%2Fhuggingface.co%2Fspaces%2Fzelk12%2FGemini-2&countColor=%23263759" />
                </a>""")
 
-    
+    # --- Добавляем состояние для Undo/Redo ---
+    undone_state = gr.State([])
+    # --------------------------------------
+
     chatbot = gr.Chatbot(
+        [], # Начинаем с пустого списка
         type="messages",
-        label=used_model + " Chatbot (Streaming Output)", #Label now indicates streaming
+        label=used_model + " Chatbot (Streaming Output)",
         render_markdown=True,
         scale=1,
-        editable="all",
-        avatar_images=(None,"https://lh3.googleusercontent.com/oxz0sUBF0iYoN4VvhqWTmux-cxfD1rxuYkuFEfm1SFaseXEsjjE4Je_C_V3UQPuJ87sImQK3HfQ3RXiaRnQetjaZbjJJUkiPL5jFJ1WRl5FKJZYibUA=w214-h214-n-nu")
+        # editable="all", # editable="all" может мешать undo/redo, лучше отключить или user
+        editable=False,
+        avatar_images=(None,"https://lh3.googleusercontent.com/oxz0sUBF0iYoN4VvhqWTmux-cxfD1rxuYkuFEfm1SFaseXEsjjE4Je_C_V3UQPuJ87sImQK3HfQ3RXiaRnQetjaZbjJJUkiPL5jFJ1WRl5FKJZYibUA=w214-h214-n-nu"),
+        height=600 # Задаем высоту для лучшего вида
     )
 
-    with gr.Row(equal_height=True):
+    with gr.Row(equal_height=False): # equal_height=False может быть лучше для кнопок разной высоты
         input_box = gr.Textbox(
             lines=1,
             label="Chat Message",
@@ -206,16 +660,23 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", n
             scale=4
         )
 
-        with gr.Column(scale=1):
-            submit_button = gr.Button("Submit", scale=1)
-            clear_button = gr.Button("Clear Chat", scale=1)
+        # --- Группируем кнопки управления ---
+        with gr.Column(scale=1, min_width=150): # Даем минимальную ширину колонке с кнопками
+             submit_button = gr.Button("Submit", scale=1, variant="primary") # Выделяем основную кнопку
+             with gr.Row():
+                 undo_button = gr.Button("Undo", scale=1)
+                 redo_button = gr.Button("Redo", scale=1)
+             regenerate_button = gr.Button("Regenerate", scale=1)
+             clear_button = gr.Button("Clear Chat", scale=1)
 
-    with gr.Row(equal_height=True):
-        test_button = gr.Button("undo", scale=1)
-        test1_button = gr.Button("redo", scale=1)
-        test2_button = gr.Button("regenerate", scale=1)
 
-    # Add example prompts - removed file upload examples. Kept text focused examples.
+    # --- Переименовали кнопки для ясности ---
+    # undo_button = test_button
+    # redo_button = test1_button
+    # regenerate_button = test2_button
+    # ----------------------------------------
+
+    # Add example prompts
     example_prompts = [
         ["Write a short poem about the sunset."],
         ["Explain the theory of relativity in simple terms."],
@@ -227,70 +688,104 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue="teal", secondary_hue="slate", n
     gr.Examples(
         examples=example_prompts,
         inputs=input_box,
-        label="Examples: Try these prompts to see Gemini's thinking!",
-        examples_per_page=5 # Adjust as needed
+        label="Examples: Try these prompts!", # Убрали про thinking, т.к. убрали его вывод
+        examples_per_page=5
     )
 
-    # Set up event handlers
-    msg_store = gr.State("")  # Store for preserving user message
-    
+    # --- Обновленные обработчики событий ---
+    msg_store = gr.State("") # Store for preserving user message during processing
+
+    # --- Обработчик для Enter в input_box ---
     input_box.submit(
-        lambda msg: (msg, msg, ""),  # Store message and clear input
-        inputs=[input_box],
-        outputs=[msg_store, input_box, input_box],
-        queue=False
+        user_message,  # 1. Добавить сообщение пользователя в историю, очистить undone_state
+        inputs=[input_box, chatbot],
+        outputs=[input_box, chatbot, undone_state] # input_box очищается первым возвращаемым значением ""
     ).then(
-        user_message,  # Add user message to chat
-        inputs=[msg_store, chatbot],
-        outputs=[input_box, chatbot],
-        queue=False
+        lambda x: (x, x), # 2. Скопировать сообщение из history[-1].content в msg_store для stream_gemini_response
+        inputs=[chatbot], # Берем историю, где последнее сообщение - пользовательское
+        outputs=[msg_store, msg_store], # Копируем контент последнего сообщения (пользователя)
+        fn=lambda history: (history[-1].content, history[-1].content) if history and history[-1].role == "user" else (None, None)
     ).then(
-        stream_gemini_response,  # Generate and stream response
-        inputs=[msg_store, chatbot],
-        outputs=chatbot
+        stream_gemini_response,  # 3. Сгенерировать и стримить ответ
+        inputs=[msg_store, chatbot], # Используем сохраненное сообщение и обновленную историю
+        outputs=[chatbot]
     )
 
+    # --- Обработчик для кнопки Submit ---
     submit_button.click(
-        lambda msg: (msg, msg, ""),  # Store message and clear input
-        inputs=[input_box],
-        outputs=[msg_store, input_box, input_box],
-        queue=False
+        user_message, # 1. Добавить сообщение пользователя, очистить undone_state
+        inputs=[input_box, chatbot],
+        outputs=[input_box, chatbot, undone_state]
     ).then(
-        user_message,  # Add user message to chat
-        inputs=[msg_store, chatbot],
-        outputs=[input_box, chatbot],
-        queue=False
+        lambda x: (x, x), # 2. Скопировать сообщение пользователя в msg_store
+        inputs=[chatbot],
+        outputs=[msg_store, msg_store],
+        fn=lambda history: (history[-1].content, history[-1].content) if history and history[-1].role == "user" else (None, None)
     ).then(
-        stream_gemini_response,  # Generate and stream response
+        stream_gemini_response, # 3. Сгенер��ровать ответ
         inputs=[msg_store, chatbot],
-        outputs=chatbot
+        outputs=[chatbot]
     )
 
+    # --- Обработчик для кнопки Clear ---
     clear_button.click(
-        lambda: ([], "", ""),
-        outputs=[chatbot, input_box, msg_store],
-        queue=False
+        lambda: ([], "", [], ""), # Очищаем chatbot, input_box, msg_store, undone_state
+        outputs=[chatbot, input_box, msg_store, undone_state],
+        queue=False # Очистка быстрая, очередь не нужна
+    )
+
+    # --- Обработчик для кнопки Undo ---
+    undo_button.click(
+        undo_last,
+        inputs=[chatbot],
+        outputs=[chatbot, undone_state], # Обновляем историю и сохраняем удаленную пару
+        queue=False # Действие быстрое
+    )
+
+    # --- Обработчик для кнопки Redo ---
+    redo_button.click(
+        redo_last,
+        inputs=[chatbot, undone_state], # Нужна история и сохраненная пара
+        outputs=[chatbot, undone_state], # Обновляем историю и очищаем сохраненную пару
+        queue=False # Действие быстрое
+    )
+
+    # --- Обработчик для кнопки Regenerate ---
+    regenerate_button.click(
+        regenerate_response, # 1. Подготовить историю и получить сообщение пользователя, очистить undone_state
+        inputs=[chatbot],
+        outputs=[chatbot, msg_store, undone_state] # Обновляем историю, помещаем user msg в msg_store, очищаем undone
+    ).then(
+        stream_gemini_response, # 2. Сгенерировать новый ответ на user msg из msg_store
+        inputs=[msg_store, chatbot], # Используем user msg из msg_store и подготовленную историю
+        outputs=[chatbot]
     )
+    # -------------------------------------
 
-    gr.Markdown(  # Description moved to the bottom - updated for text-only
+    gr.Markdown(
         """
-        <br><br><br>  <!-- Add some vertical space -->
+        <br><br><br>
         ---
         ### About this Chatbot
-        **Try out the example prompts below to see Gemini in action!**
+        **Try out the example prompts below!**
         **Key Features:**
-        *   Powered by Google's **Gemini 2.0 Flash** model.
-        *   Supports **conversation history** for multi-turn chats.
-        *   Uses **streaming** for a more interactive experience.
+        *   Powered by Google's **""" + used_model + """** model.
+        *   Supports **conversation history**, **undo**, **redo**, and **regenerate**.
+        *   Uses **streaming** for responses.
         **Instructions:**
-        1.  Type your message in the input box below or select an example.
-        2.  Press Enter or click Submit to send.
-        3.  Observe the chatbot's "Thinking" process followed by the final response.
-        4.  Use the "Clear Chat" button to start a new conversation.
+        1.  Type your message or select an example.
+        2.  Press Enter or click Submit.
+        3.  Use Undo/Redo to correct mistakes or explore alternatives.
+        4.  Use Regenerate to get a different response to your last message.
+        5.  Use Clear Chat to start over.
         """
     )
 
 
 # Launch the interface
 if __name__ == "__main__":
-    demo.launch(debug=True)
+    # --- Добавим возможность запускать на публичном URL через share=True ---
+    # demo.launch(debug=True) # Для локальной отладки
+    demo.launch(share=True) # Для получения публичной ссылки (если нужно)
+    # demo.launch() # Стандартный локальный запуск
+    # -----------------------------------------------------------------