Update app.py

app.py

@@ -19,7 +19,7 @@ MODELS = {model_name: genai.GenerativeModel(model_name=model_name) for model_nam
 async def respond(message, history, selected_model):
     model = MODELS.get(selected_model)
     if not model:
-        yield "Error: Selected model not available.", ""
+        yield {"role": "assistant", "content": "Error: Selected model not available."}, ""
         return
 
     try:
@@ -28,21 +28,21 @@ async def respond(message, history, selected_model):
         full_response = ""
         for chunk in response_stream:
             full_response += (chunk.text or "")
-            yield full_response, ""  # Пустая строка для thinking output
+            yield {"role": "assistant", "content": full_response}, ""  # Формат messages
     except Exception as e:
-        yield f"Error during API call: {e}", ""
+        yield {"role": "assistant", "content": f"Error during API call: {e}"}, ""
 
 async def respond_thinking(message, history, selected_model):
     if "thinking" not in selected_model:
-        yield "Thinking model не выбрана.", ""
+        yield {"role": "assistant", "content": "Thinking model не выбрана."}, ""
         return
 
     model = MODELS.get(selected_model)
     if not model:
-        yield "Error: Selected model not available.", ""
+        yield {"role": "assistant", "content": "Error: Selected model not available."}, ""
         return
 
-    yield "", "Думаю..." # Сообщение о начале размышлений
+    yield {"role": "assistant", "content": "Думаю..."}, "" # Сообщение о начале размышлений
 
     try:
         response = model.generate_content(message)
@@ -56,9 +56,9 @@ async def respond_thinking(message, history, selected_model):
                 else:
                     model_response_text += (part.text or "")
 
-        yield model_response_text, thinking_process_text
+        yield {"role": "assistant", "content": model_response_text}, thinking_process_text
     except Exception as e:
-        yield f"Error during API call: {e}", f"Error during API call: {e}"
+        yield {"role": "assistant", "content": f"Error during API call: {e}"}, f"Error during API call: {e}"
 
 def update_chatbot_function(model_name):
     if "thinking" in model_name:
@@ -66,6 +66,17 @@ def update_chatbot_function(model_name):
     else:
         return respond
 
+async def process_message(message, history, model_name):
+    if "thinking" in model_name:
+        generator = respond_thinking(message, history, model_name)
+        thinking_output_content = ""
+        async for response, thinking in generator:
+            yield response, thinking
+        #  Здесь можно было бы дополнительно обработать thinking_output_content, если нужно
+    else:
+        async for response, _ in respond(message, history, model_name):
+            yield response, ""
+
 with gr.Blocks() as demo:
     gr.Markdown("# Gemini Chatbot с режимом размышления")
 
@@ -75,45 +86,35 @@ with gr.Blocks() as demo:
         )
 
     chatbot = gr.ChatInterface(
-        respond,  # Изначально используем асинхронную функцию respond
+        process_message,  # Функция обработки сообщений передается сюда
         additional_inputs=[model_selection],
         title="Gemini Chat",
         description="Общайтесь с моделями Gemini от Google.",
+        type="messages"
     )
 
     thinking_output = gr.Code(label="Процесс размышления (для моделей с размышлением)")
 
     def change_function(model_name):
-        return update_chatbot_function(model_name)
+        # Обновляем функцию обработки сообщений при смене модели
+        if "thinking" in model_name:
+            return respond_thinking
+        else:
+            return respond
+
+    def change_chatbot(model_name):
+        return gr.ChatInterface.update(fn=process_message)
 
     model_selection.change(
-        change_function,
+        change_chatbot,
         inputs=[model_selection],
         outputs=[chatbot],
     )
 
-    async def process_message(message, history, model_name):
-        if "thinking" in model_name:
-            generator = respond_thinking(message, history, model_name)
-            response, thinking = await generator.__anext__() # Получаем первое значение (пустое сообщение и "Думаю...")
-            yield response, thinking
-            final_response, final_thinking = await generator.__anext__() # Получаем окончательный ответ и размышления
-            yield final_response, final_thinking
-        else:
-            async for response, _ in respond(message, history, model_name):
-                yield response, ""
+    def clear_thinking():
+        return ""
 
-    chatbot.on(
-        process_message,
-        inputs=[chatbot, chatbot.chat_memory, model_selection],
-        outputs=[chatbot, thinking_output],
-    )
-
-    chatbot.change(
-        lambda: "",
-        inputs=[],
-        outputs=[thinking_output]
-    )
+    chatbot.clear(inputs=[], outputs=[thinking_output])
 
 if __name__ == "__main__":
     demo.launch()