app.py

import os
import gradio as gr
import google.generativeai as genai
import asyncio

###############################################################################
# 1. Настройка окружения и моделей
###############################################################################

GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "ВАШ_КЛЮЧ_ЗДЕСЬ")
if not GEMINI_API_KEY:
    print("Error: GEMINI_API_KEY is not set.")
    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",
]

# Инициализация моделей
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. Промпты по умолчанию (system message) для каждой модели
###############################################################################
MODEL_SYSTEM_PROMPTS = {
    # Пример: для обычной модели
    "gemini-2.0-flash-exp": (
        "You are a helpful assistant. "
        "You respond in Markdown. "
        "You do NOT wrap your final answer in {output: ...}."
    ),
    # Допустим, это тоже обычная модель
    "gemini-exp-1206": (
        "You are an experimental Gemini model. "
        "You respond in Markdown. "
        "You do NOT wrap your final answer in {output: ...}."
    ),
    # Думающая модель
    "gemini-2.0-flash-thinking-exp-1219": (
        "You are a thinking model. "
        "You respond with thoughts internally, but produce the **final** answer in JSON format, like `{output: (final text)}`."
    ),
}

###############################################################################
# 3. Определяем, какую роль использовать: 'assistant' или 'model'
###############################################################################
def _assistant_role(model_name: str) -> str:
    """
    Для некоторых моделей Gemini требуется role='model', а не 'assistant'.
    Если столкнетесь с ошибкой "Please use a valid role: user, model", 
    укажите здесь 'model'. 
    """
    # Условие — подстройте под свои модели:
    # Предположим, gemini-exp-1206 и gemini-2.0-flash-thinking-exp-1219
    # требуют role='model', а остальные — 'assistant'
    if model_name in ["gemini-exp-1206", "gemini-2.0-flash-thinking-exp-1219"]:
        return "model"
    return "assistant"

###############################################################################
# 4. Утилиты для конвертации истории
###############################################################################
def _history_gradio_to_genai(history, model_name):
    """
    Gradio: [(user_msg, bot_msg), (user_msg, bot_msg), ...]
    → Gemini: [{'role': 'user'|'assistant'|'model'|'system', 'parts': ...}, ...]

    Добавляем system-сообщение (промпт по умолчанию) в начало. 
    Далее все user-реплики = {'role': 'user'}, 
         все ответы ассистента = {'role': <assistant_role>}.
    """
    genai_history = []

    # Системное сообщение (промпт по умолчанию)
    system_prompt = MODEL_SYSTEM_PROMPTS.get(model_name, "")
    if system_prompt:
        genai_history.append({"role": "system", "parts": system_prompt})

    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": asst_role, "parts": bot_text})

    return genai_history

###############################################################################
# 5. Функции-генераторы для ответа
###############################################################################

async def _respond_stream(model_name, user_message, history):
    """
    Обычная модель, stream=True: выдаём ответ порциями.
    """
    if model_name not in MODELS:
        yield "Error: model not found."
        return

    model = MODELS[model_name]
    genai_history = _history_gradio_to_genai(history, model_name)

    try:
        chat = model.start_chat(history=genai_history)
        response_stream = chat.send_message(user_message, stream=True)
        partial_text = ""
        for chunk in response_stream:
            partial_text += (chunk.text or "")
            yield partial_text
        return
    except Exception as e:
        yield f"Ошибка при запросе к API: {e}"
        return


async def _respond_thinking(model_name, user_message, history):
    """
    Думающая модель:
    1) "Думаю..."
    2) По завершении:
       - собираем "мысли" (part.thought)
       - итоговый ответ оборачиваем в {output: ...} (согласно ТЗ)
    """
    if model_name not in MODELS:
        yield "Error: model not found.", ""
        return

    model = MODELS[model_name]
    genai_history = _history_gradio_to_genai(history, model_name)

    # Шаг 1: "Думаю..."
    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 ""

        # По ТЗ: оборачиваем финальный ответ в {output: ...}
        final_text_json = f"{{output: ({final_text})}}"

        yield final_text_json, thinking_process_text
        return
    except Exception as e:
        yield f"Ошибка при запросе к API: {e}", ""
        return

###############################################################################
# 6. Основная функция для одного шага диалога
###############################################################################
async def user_send_message(user_message, history, model_name, thinking_text):
    """
    Пользователь ввёл user_message; на основе history + model_name генерируем ответ.
    Возвращаем (history, thinking_text) через yield на каждом шаге.
    """
    # Если пустая строка — ничего не делаем
    if not user_message.strip():
        yield history, thinking_text
        return

    # Добавляем ход в history, пока ответ ассистента None
    history.append((user_message, None))

    # Проверяем, является ли модель "думающей"
    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
            history[-1] = (user_message, partial_answer)
            yield history, ""
        return

###############################################################################
# 7. Логика переключения модели, чтобы «сообщать» о переключении
###############################################################################
def switch_model(old_model, new_model, history):
    """
    При смене модели добавляем в историю user-сообщение:
      "I’m switching from {old_model} to {new_model}."
    Сохраняем новое имя модели в state (prev_model_state).
    """
    if old_model and new_model and old_model != new_model:
        switch_message = f"I’m switching from {old_model} to {new_model}."
        history.append((switch_message, None))  # user_msg, None
    return new_model, history

###############################################################################
# 8. Функция очистки
###############################################################################
def clear_all():
    """Сбросить историю и размышления."""
    return [], ""

###############################################################################
# 9. Сборка Gradio-интерфейса
###############################################################################
with gr.Blocks() as demo:
    gr.Markdown("## Chat with Gemini (Thinking & Non-Thinking Models)")

    # Храним «предыдущую модель», чтобы отследить переключение
    prev_model_state = gr.State("")  
    # Храним историю [(user, assistant), ...]
    history_state = gr.State([])   
    # Храним размышления (только для думающей модели)
    thinking_store = gr.State("")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=AVAILABLE_MODELS,
            value="gemini-2.0-flash-exp",
            label="Choose a model"
        )
        clear_button = gr.Button("Очистить чат")

    chatbot = gr.Chatbot(
        label="Диалог с Gemini",
        markdown=True  # Включаем поддержку Markdown
    )
    user_input = gr.Textbox(
        label="Ваш вопрос",
        placeholder="Введите текст...",
    )
    thinking_output = gr.Textbox(
        label="Размышления (если модель думающая)",
        interactive=False
    )
    send_btn = gr.Button("Отправить")

    # --- Обработка переключения модели ---
    model_dropdown.change(
        fn=switch_model,
        inputs=[prev_model_state, model_dropdown, history_state],
        outputs=[prev_model_state, history_state],
        queue=False
    ).then(
        # После смены модели обновим чат (чтобы увидеть user-сообщение о переключении)
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=False
    )

    # --- При нажатии кнопки «Отправить» ---
    send_chain = send_btn.click(
        fn=user_send_message,
        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
    )
    # Очищаем поле ввода
    send_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input],
        queue=True
    )

    # --- При нажатии Enter в поле ввода ---
    submit_chain = user_input.submit(
        fn=user_send_message,
        inputs=[user_input, history_state, model_dropdown, thinking_store],
        outputs=[history_state, thinking_store],
        queue=True
    )
    submit_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=True
    )
    submit_chain.then(
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )
    # Очищаем поле
    submit_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input],
        queue=True
    )

    # --- Кнопка «Очистить» ---
    clear_chain = clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=[history_state, thinking_store],
        queue=False
    )
    clear_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot]
    )
    clear_chain.then(
        fn=lambda _: "",
        inputs=[],
        outputs=[thinking_output]
    )
    clear_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input]
    )

if __name__ == "__main__":
    demo.launch()