app.py

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

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

# Задайте свой ключ через переменную окружения, например:
# export GEMINI_API_KEY="ваш-ключ"
GEMINI_API_KEY = "AIzaSyBoqoPX-9uzvXyxzse0gRwH8_P9xO6O3Bc" 
if not GEMINI_API_KEY:
    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",
]

# Пытаемся инициализировать все модели заранее и складываем в словарь
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. Утилиты для преобразования истории Gradio <-> Gemini
###############################################################################

def _history_gradio_to_genai(history):
    """
    Gradio хранит чат как [(user_msg, bot_msg), (user_msg, bot_msg), ...].
    Для Gemini нужен формат [{'role': 'user', '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


###############################################################################
# 3. Функции-генераторы для запроса ответа от моделей (обычный/thinking)
###############################################################################

async def _respond_stream(model_name, user_message, history):
    """
    Генерация ответа для Обычных моделей (без thinking) с помощью stream=True
    - Возвращаем куски текста через yield
    - В конце просто делаем return (без значения)
    """
    if model_name not in MODELS:
        yield "Ошибка: модель не найдена."
        return

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

    try:
        chat = model.start_chat(history=genai_history)
        response_stream = chat.send_message(user_message, stream=True)

        partial_text = ""
        for chunk in response_stream:
            chunk_text = chunk.text or ""
            partial_text += chunk_text
            # Возвращаем промежуточный вариант ответа
            yield partial_text

        return  # Завершить генератор без возвращения значения
    except Exception as e:
        yield f"Ошибка при запросе к API: {e}"
        return


async def _respond_thinking(model_name, user_message, history):
    """
    Генерация ответа для модели с "thinking" (например, gemini-2.0-flash-thinking-exp).
    1) Сначала выдаём "Думаю..."
    2) Затем, когда модель ответит (stream=False), выделяем thinking + финальный ответ.
    3) Возвращаем (полезный_ответ, размышления) в виде кортежа.
       В Gradio это обычно [(assistant_text, thinking_text), ...].
    """
    if model_name not in MODELS:
        # Выдаем ошибку через yield
        yield "Ошибка: модель не найдена.", ""
        return

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

    # Шаг 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 ""

        # Возвращаем готовый ответ и размышления
        yield final_text, thinking_process_text
        return
    except Exception as e:
        yield f"Ошибка при запросе к API: {e}", ""
        return


###############################################################################
# 4. Основная асинхронная функция Gradio, обрабатывающая новый пользовательский ввод
###############################################################################

async def user_send_message(
    user_message: str,
    history: list[tuple[str, str]],
    model_name: str,
    thinking_text: str
):
    """
    Параметры:
      user_message: вход от пользователя (текущая реплика)
      history: история [(user, assistant), ...]
      model_name: выбранная модель
      thinking_text: текущее содержимое поля «Размышления»

    Выход (через yield):
      (обновлённая история, обновлённое thinking_text)
    """
    # Если пользователь ничего не ввёл, просто возвращаем текущее состояние
    if not user_message.strip():
        yield history, thinking_text
        return

    # Добавляем новую запись в историю: ассистент пока None
    history.append((user_message, None))

    # Проверяем, thinking-модель ли
    if "thinking" in model_name.lower():
        # Обрабатываем через _respond_thinking
        async for (assistant_text, thought_text) in _respond_thinking(model_name, user_message, history):
            # Обновляем последнюю пару в истории
            history[-1] = (user_message, assistant_text)
            # Обновляем thinking_text
            yield history, thought_text
        return
    else:
        # Обычная модель (stream)
        partial_answer = ""
        async for chunk in _respond_stream(model_name, user_message, history):
            partial_answer = chunk
            history[-1] = (user_message, partial_answer)
            # В обычном режиме thinking_text = ""
            yield history, ""
        return


###############################################################################
# 5. Колбэки для очистки
###############################################################################

def clear_all():
    """
    Полная очистка чата и размышлений.
    """
    return [], ""  # (history, thinking_store)


###############################################################################
# 6. Определяем интерфейс Gradio
###############################################################################

with gr.Blocks() as demo:
    gr.Markdown("## Gemini Chatbot (с сохранением истории и thinking-режимом)")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=AVAILABLE_MODELS,
            value="gemini-1.5-flash",
            label="Выберите модель",
        )
        clear_button = gr.Button("Очистить чат")

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

    chatbot = gr.Chatbot(label="Диалог с Gemini")

    user_input = gr.Textbox(
        label="Ваш вопрос",
        placeholder="Введите вопрос и нажмите Enter...",
    )

    # Отдельный блок для показа «размышлений»
    thinking_output = gr.Textbox(
        label="Размышления (только для gemini-2.0-flash-thinking-exp)",
        interactive=False
    )

    send_btn = gr.Button("Отправить")

    # Связка: кнопка «Отправить» => user_send_message => обновление истории и размышлений
    send_btn.click(
        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],
        queue=True
    ).then(
        # После этого выводим thinking_store
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )

    # Аналогичное поведение при нажатии 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],
        queue=True
    ).then(
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )

    # Кнопка «Очистить» сбрасывает историю и thinking_store
    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(
        # И обнуляем thinking_output
        fn=lambda _: "",
        inputs=[],
        outputs=[thinking_output]
    )

# Запуск Gradio
if __name__ == "__main__":
    demo.launch()