Update app.py

app.py

@@ -8,18 +8,17 @@ import asyncio
 ###############################################################################
 
 # Подставьте свой ключ или берите из окружения
-GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc"
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
 if not GEMINI_API_KEY:
-    print("Error: GEMINI_API_KEY is not set.")
+    print("Error: GEMINI_API_KEY is not set. Please set the environment variable.")
     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-1.5-flash",
+    # Add other available models if needed
 ]
 
 MODELS = {}
@@ -35,19 +34,11 @@ for model_name in AVAILABLE_MODELS:
 # Когда пользователь переключается на модель, мы добавляем это сообщение в историю.
 
 DEFAULT_DEVELOPER_PROMPTS = {
-    "gemini-2.0-flash-exp": (
-        "You are a normal model (developer role). "
-        "Provide direct answers with no JSON wrapping."
-    ),
-    "gemini-exp-1206": (
-        "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 (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."
+    "gemini-1.5-flash": (
+        "You are a helpful assistant (developer role). "
+        "Provide direct answers."
     ),
+    # Add default prompts for other models if needed
 }
 
 ###############################################################################
@@ -58,9 +49,7 @@ def _assistant_role(model_name: str) -> str:
     """
     Некоторые новые модели не принимают 'assistant', а требуют 'model'.
     """
-    # Допустим "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"
+    # Adjust based on model requirements
     return "assistant"
 
 ###############################################################################
@@ -78,8 +67,11 @@ def _history_to_genai(history, model_name):
 
     for user_text, bot_text in history:
         if user_text:
-            # Пользовательское сообщение
-            genai_history.append({"role": "user", "parts": user_text})
+            if user_text.startswith("<developer>: "):
+                genai_history.append({"role": "system", "parts": user_text.replace("<developer>: ", "", 1)})
+            else:
+                # Пользовательское сообщение
+                genai_history.append({"role": "user", "parts": user_text})
         if bot_text:
             # Ответ ассистента (или model)
             genai_history.append({"role": asst_role, "parts": bot_text})
@@ -91,7 +83,7 @@ def _history_to_genai(history, model_name):
 
 async def _respond_stream(model_name, user_message, history):
     """
-    Стриминговый ответ для обычных моделей: 
+    Стриминговый ответ для обычных моделей:
     - Кусочек за кусочком (partial_text).
     """
     if model_name not in MODELS:
@@ -103,10 +95,10 @@ async def _respond_stream(model_name, user_message, history):
 
     try:
         chat = model.start_chat(history=genai_history)
-        stream = chat.send_message(user_message, stream=True)
+        response = chat.send_message(user_message, stream=True)
 
         partial_text = ""
-        for chunk in stream:
+        async for chunk in response:
             partial_text += (chunk.text or "")
             yield partial_text
 
@@ -115,11 +107,10 @@ async def _respond_stream(model_name, user_message, history):
         yield f"Ошибка при запросе к API: {e}"
         return
 
-
 async def _respond_thinking(model_name, user_message, history):
     """
     Для thinking-моделей:
-    1) Выводим "Думаю..." 
+    1) Выводим "Думаю..."
     2) После завершения — финальный ответ в формате {output: ...} + размышления.
     """
     if model_name not in MODELS:
@@ -142,10 +133,14 @@ 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):
+                # Assuming a specific structure for "thinking" models' output
+                # This part might need adjustments based on actual model output
+                if isinstance(p, dict) and "output" not in p:  # Heuristic for thought
                     thinking_process_text += p.text or ""
+                elif isinstance(p, dict) and "output" in p:
+                    final_text += p["output"]  # Assuming output is directly in the dict
                 else:
-                    final_text += p.text or ""
+                    final_text += p.text or ""  # Fallback
 
         # Для thinking-моделей просили итоговый ответ в {output: ...}
         final_text_formatted = f"{{output: {final_text}}}"
@@ -168,7 +163,7 @@ async def user_send_message(
     thinking_text: str
 ):
     """
-    Колбэк, когда пользователь отправляет запрос. 
+    Колбэк, когда пользователь отправляет запрос.
     Добавляем в history новый (user_msg, None), затем генерируем ответ.
     """
     # Пустой ввод
@@ -179,9 +174,9 @@ async def user_send_message(
     # Добавляем (user_message, None)
     history.append((user_message, None))
 
-    # Если модель — thinking
+    # Если модель — thinking (adjust logic based on actual model names)
     if "thinking" in model_name.lower():
-        async for (assistant_text, thought_text) in _respond_thinking(model_name, user_message, history):
+        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
@@ -208,43 +203,24 @@ def clear_all():
 
 def on_model_change(selected_model, history, thinking_text):
     """
-    При переключении модели добавляем в историю developer-сообщение, 
+    При переключении модели добавляем в историю developer-сообщение,
     + добавляем дефолтный промпт этой модели (тоже developer).
     """
     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,
+        f"<developer>: Switched to model: {selected_model}",
         None
     ))
 
     # Добавляем дефолтный промпт (developer role)
     default_prompt = DEFAULT_DEVELOPER_PROMPTS.get(
-        selected_model, 
+        selected_model,
         "No default prompt for this model."
     )
     new_history.append((
-        "<developer>: " + default_prompt,
+        f"<developer>: {default_prompt}",
         None
     ))
 
@@ -254,112 +230,7 @@ def on_model_change(selected_model, history, thinking_text):
 # 9. Функция конвертации истории с учётом developer role
 ###############################################################################
 
-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
+# This part is already implemented in _history_to_genai
 
 ###############################################################################
 # 10. Построение интерфейса Gradio
@@ -371,7 +242,7 @@ with gr.Blocks() as demo:
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=AVAILABLE_MODELS,
-            value="gemini-2.0-flash-exp",
+            value=AVAILABLE_MODELS[0] if AVAILABLE_MODELS else None,
             label="Выберите модель"
         )
         clear_button = gr.Button("Очистить чат")
@@ -409,7 +280,7 @@ with gr.Blocks() as demo:
     # (B) При нажатии «Отправить»
     ################################################
     send_chain = send_btn.click(
-        fn=user_send_message_enh,
+        fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
         outputs=[history_state, thinking_store],
         queue=True
@@ -438,7 +309,7 @@ with gr.Blocks() as demo:
     # (C) При нажатии Enter в textbox
     ################################################
     submit_chain = user_input.submit(
-        fn=user_send_message_enh,
+        fn=user_send_message,
         inputs=[user_input, history_state, model_dropdown, thinking_store],
         outputs=[history_state, thinking_store],
         queue=True
@@ -490,4 +361,4 @@ with gr.Blocks() as demo:
 
 # Запуск
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()